mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +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,
|
CancelVideoTaskError,
|
||||||
};
|
};
|
||||||
pub(crate) use query::{
|
pub(crate) use query::{
|
||||||
read_video_task_detail, read_video_task_page, read_video_task_stats,
|
read_video_task_detail, read_video_task_page, read_video_task_page_summary,
|
||||||
read_video_task_video_source, VideoTaskPageResponse, VideoTaskStatsResponse,
|
read_video_task_stats, read_video_task_video_source, VideoTaskPageResponse,
|
||||||
VideoTaskVideoSource,
|
VideoTaskStatsResponse, VideoTaskVideoSource,
|
||||||
};
|
};
|
||||||
pub(crate) use runtime::{
|
pub(crate) use runtime::{
|
||||||
execute_video_task_refresh_plan, finalize_video_task_if_terminal, spawn_video_task_poller,
|
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(
|
pub(crate) async fn read_video_task_detail(
|
||||||
state: &AppState,
|
state: &AppState,
|
||||||
task_id: &str,
|
task_id: &str,
|
||||||
|
|||||||
@@ -267,11 +267,14 @@ fn build_successful_poll_update(
|
|||||||
record.format_converted = task.format_converted;
|
record.format_converted = task.format_converted;
|
||||||
record.model = task.model.clone().or(record.model);
|
record.model = task.model.clone().or(record.model);
|
||||||
record.prompt = task.prompt.clone().or(record.prompt);
|
record.prompt = task.prompt.clone().or(record.prompt);
|
||||||
record.original_request_body = task.original_request_body.clone();
|
record.original_request_body = task
|
||||||
record.duration_seconds = task.duration_seconds;
|
.original_request_body
|
||||||
record.resolution = task.resolution.clone();
|
.clone()
|
||||||
record.aspect_ratio = task.aspect_ratio.clone();
|
.or(record.original_request_body);
|
||||||
record.size = task.size.clone();
|
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.created_at_unix_ms = task.created_at_unix_ms;
|
||||||
record.submitted_at_unix_secs = task.submitted_at_unix_secs;
|
record.submitted_at_unix_secs = task.submitted_at_unix_secs;
|
||||||
record.updated_at_unix_secs = now_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 {
|
fn stored_task_to_upsert(task: &StoredVideoTask) -> UpsertVideoTask {
|
||||||
|
let snapshot_record =
|
||||||
|
LocalVideoTaskSnapshot::from_stored_task(task).map(|snapshot| snapshot.to_upsert_record());
|
||||||
UpsertVideoTask {
|
UpsertVideoTask {
|
||||||
id: task.id.clone(),
|
id: task.id.clone(),
|
||||||
short_id: task.short_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(),
|
provider_api_format: task.provider_api_format.clone(),
|
||||||
format_converted: task.format_converted,
|
format_converted: task.format_converted,
|
||||||
model: task.model.clone(),
|
model: task.model.clone(),
|
||||||
prompt: task.prompt.clone(),
|
prompt: task.prompt.clone().or_else(|| {
|
||||||
original_request_body: task.original_request_body.clone(),
|
snapshot_record
|
||||||
duration_seconds: task.duration_seconds,
|
.as_ref()
|
||||||
resolution: task.resolution.clone(),
|
.and_then(|record| record.prompt.clone())
|
||||||
aspect_ratio: task.aspect_ratio.clone(),
|
}),
|
||||||
size: task.size.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,
|
status: task.status,
|
||||||
progress_percent: task.progress_percent,
|
progress_percent: task.progress_percent,
|
||||||
progress_message: task.progress_message.clone(),
|
progress_message: task.progress_message.clone(),
|
||||||
@@ -550,3 +579,154 @@ fn now_unix_secs() -> u64 {
|
|||||||
.unwrap_or_default()
|
.unwrap_or_default()
|
||||||
.as_secs()
|
.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_key_hash,
|
||||||
read_resolved_auth_api_key_snapshot_by_user_api_key_ids,
|
read_resolved_auth_api_key_snapshot_by_user_api_key_ids,
|
||||||
};
|
};
|
||||||
|
use futures_util::TryStreamExt;
|
||||||
use sqlx::Row;
|
use sqlx::Row;
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
@@ -1278,12 +1279,14 @@ impl GatewayDataState {
|
|||||||
let Some(pool) = self.postgres_pool() else {
|
let Some(pool) = self.postgres_pool() else {
|
||||||
return Ok(Vec::new());
|
return Ok(Vec::new());
|
||||||
};
|
};
|
||||||
let rows = sqlx::query(LIST_USER_SESSIONS_SQL)
|
let mut rows = sqlx::query(LIST_USER_SESSIONS_SQL)
|
||||||
.bind(user_id)
|
.bind(user_id)
|
||||||
.fetch_all(&pool)
|
.fetch(&pool);
|
||||||
.await
|
let mut sessions = Vec::new();
|
||||||
.map_postgres_err()?;
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||||
rows.iter().map(map_user_session_row).collect()
|
sessions.push(map_user_session_row(&row)?);
|
||||||
|
}
|
||||||
|
Ok(sessions)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn create_user_session(
|
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(
|
pub(crate) async fn list_provider_catalog_key_page(
|
||||||
&self,
|
&self,
|
||||||
query: &ProviderCatalogKeyListQuery,
|
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(
|
pub(crate) async fn count_video_tasks(
|
||||||
&self,
|
&self,
|
||||||
filter: &VideoTaskQueryFilter,
|
filter: &VideoTaskQueryFilter,
|
||||||
|
|||||||
@@ -315,7 +315,7 @@ pub(crate) async fn recover_all_admin_key_health(
|
|||||||
Vec::new()
|
Vec::new()
|
||||||
} else {
|
} else {
|
||||||
state
|
state
|
||||||
.list_provider_catalog_keys_by_provider_ids(&provider_ids)
|
.list_provider_catalog_key_summaries_by_provider_ids(&provider_ids)
|
||||||
.await
|
.await
|
||||||
.ok()
|
.ok()
|
||||||
.unwrap_or_default()
|
.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();
|
let mut health_scores_by_format = BTreeMap::<String, Vec<f64>>::new();
|
||||||
if !provider_ids.is_empty() {
|
if !provider_ids.is_empty() {
|
||||||
let keys = state
|
let keys = state
|
||||||
.list_provider_catalog_keys_by_provider_ids(&provider_ids)
|
.list_provider_catalog_key_summaries_by_provider_ids(&provider_ids)
|
||||||
.await
|
.await
|
||||||
.ok()
|
.ok()
|
||||||
.unwrap_or_default();
|
.unwrap_or_default();
|
||||||
@@ -235,7 +235,7 @@ pub(crate) async fn build_admin_health_summary_payload(
|
|||||||
Vec::new()
|
Vec::new()
|
||||||
} else {
|
} else {
|
||||||
state
|
state
|
||||||
.list_provider_catalog_keys_by_provider_ids(&provider_ids)
|
.list_provider_catalog_key_summaries_by_provider_ids(&provider_ids)
|
||||||
.await
|
.await
|
||||||
.ok()
|
.ok()
|
||||||
.unwrap_or_default()
|
.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")
|
let page_size = query_param_value(request_context.query_string(), "page_size")
|
||||||
.and_then(|value| value.parse::<usize>().ok())
|
.and_then(|value| value.parse::<usize>().ok())
|
||||||
.unwrap_or(20);
|
.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?;
|
let provider_names = build_admin_video_task_provider_names(state, &response.items).await?;
|
||||||
return Ok(Some(
|
return Ok(Some(
|
||||||
Json(json!({
|
Json(json!({
|
||||||
|
|||||||
@@ -58,7 +58,7 @@ pub(super) async fn build_admin_monitoring_provider_name_by_id_and_keys(
|
|||||||
Vec::new()
|
Vec::new()
|
||||||
} else {
|
} else {
|
||||||
state
|
state
|
||||||
.list_provider_catalog_keys_by_provider_ids(&provider_ids)
|
.list_provider_catalog_key_summaries_by_provider_ids(&provider_ids)
|
||||||
.await?
|
.await?
|
||||||
};
|
};
|
||||||
Ok((provider_name_by_id, keys))
|
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::provider::shared::support::AdminProviderPoolRuntimeState;
|
||||||
pub(crate) use crate::handlers::admin::shared::attach_admin_audit_response;
|
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(
|
pub(crate) async fn maybe_build_local_admin_pool_response(
|
||||||
state: &AdminAppState<'_>,
|
state: &AdminAppState<'_>,
|
||||||
|
|||||||
@@ -3,7 +3,7 @@ use super::{
|
|||||||
parse_admin_pool_page, parse_admin_pool_page_size, parse_admin_pool_search,
|
parse_admin_pool_page, parse_admin_pool_page_size, parse_admin_pool_search,
|
||||||
parse_admin_pool_status_filter, pool_payloads, pool_selection,
|
parse_admin_pool_status_filter, pool_payloads, pool_selection,
|
||||||
read_admin_provider_pool_cooldown_key_ids, read_admin_provider_pool_runtime_state,
|
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,
|
ADMIN_POOL_PROVIDER_CATALOG_READER_UNAVAILABLE_DETAIL,
|
||||||
};
|
};
|
||||||
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||||
@@ -118,6 +118,7 @@ pub(super) async fn build_admin_pool_list_keys_response(
|
|||||||
},
|
},
|
||||||
offset: page_offset,
|
offset: page_offset,
|
||||||
limit: page_size,
|
limit: page_size,
|
||||||
|
order: ProviderCatalogKeyListOrder::Name,
|
||||||
})
|
})
|
||||||
.await?;
|
.await?;
|
||||||
(key_page.items, key_page.total)
|
(key_page.items, key_page.total)
|
||||||
|
|||||||
@@ -31,7 +31,7 @@ pub(crate) async fn build_admin_provider_summary_payload(
|
|||||||
active_global_model_ids_result,
|
active_global_model_ids_result,
|
||||||
) = tokio::join!(
|
) = tokio::join!(
|
||||||
state.list_provider_catalog_endpoints_by_provider_ids(&provider_ids),
|
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.read_provider_quota_snapshot(provider_id),
|
||||||
state.list_provider_model_stats(&provider_ids),
|
state.list_provider_model_stats(&provider_ids),
|
||||||
state.list_active_global_model_ids_by_provider_ids(&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()
|
Vec::new()
|
||||||
} else {
|
} else {
|
||||||
state
|
state
|
||||||
.list_provider_catalog_keys_by_provider_ids(&provider_ids)
|
.list_provider_catalog_key_summaries_by_provider_ids(&provider_ids)
|
||||||
.await
|
.await
|
||||||
.ok()
|
.ok()
|
||||||
.unwrap_or_default()
|
.unwrap_or_default()
|
||||||
|
|||||||
@@ -48,11 +48,11 @@ pub(crate) async fn build_admin_providers_payload(
|
|||||||
.ok()
|
.ok()
|
||||||
.unwrap_or_default()
|
.unwrap_or_default()
|
||||||
};
|
};
|
||||||
let keys = if provider_ids.is_empty() {
|
let key_stats = if provider_ids.is_empty() {
|
||||||
Vec::new()
|
Vec::new()
|
||||||
} else {
|
} else {
|
||||||
state
|
state
|
||||||
.list_provider_catalog_keys_by_provider_ids(&provider_ids)
|
.list_provider_catalog_key_stats_by_provider_ids(&provider_ids)
|
||||||
.await
|
.await
|
||||||
.ok()
|
.ok()
|
||||||
.unwrap_or_default()
|
.unwrap_or_default()
|
||||||
@@ -68,9 +68,12 @@ pub(crate) async fn build_admin_providers_payload(
|
|||||||
},
|
},
|
||||||
);
|
);
|
||||||
let has_any_key_by_provider =
|
let has_any_key_by_provider =
|
||||||
keys.into_iter()
|
key_stats
|
||||||
.fold(BTreeSet::<String>::new(), |mut acc, key| {
|
.into_iter()
|
||||||
acc.insert(key.provider_id);
|
.fold(BTreeSet::<String>::new(), |mut acc, stats| {
|
||||||
|
if stats.total_keys > 0 {
|
||||||
|
acc.insert(stats.provider_id);
|
||||||
|
}
|
||||||
acc
|
acc
|
||||||
});
|
});
|
||||||
|
|
||||||
|
|||||||
@@ -1,4 +1,7 @@
|
|||||||
use crate::handlers::admin::request::AdminAppState;
|
use crate::handlers::admin::request::AdminAppState;
|
||||||
|
use aether_data_contracts::repository::provider_catalog::{
|
||||||
|
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery,
|
||||||
|
};
|
||||||
use std::time::{SystemTime, UNIX_EPOCH};
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||||||
|
|
||||||
pub(crate) async fn build_admin_provider_keys_payload(
|
pub(crate) async fn build_admin_provider_keys_payload(
|
||||||
@@ -15,30 +18,26 @@ pub(crate) async fn build_admin_provider_keys_payload(
|
|||||||
.await
|
.await
|
||||||
.ok()
|
.ok()
|
||||||
.and_then(|mut providers| providers.drain(..).next())?;
|
.and_then(|mut providers| providers.drain(..).next())?;
|
||||||
let mut keys = state
|
let key_page = state
|
||||||
.list_provider_catalog_keys_by_provider_ids(std::slice::from_ref(&provider.id))
|
.list_provider_catalog_key_page(&ProviderCatalogKeyListQuery {
|
||||||
|
provider_id: provider.id.clone(),
|
||||||
|
search: None,
|
||||||
|
is_active: None,
|
||||||
|
offset: skip,
|
||||||
|
limit,
|
||||||
|
order: ProviderCatalogKeyListOrder::CreatedAt,
|
||||||
|
})
|
||||||
.await
|
.await
|
||||||
.ok()
|
.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))
|
|
||||||
});
|
|
||||||
let now_unix_secs = SystemTime::now()
|
let now_unix_secs = SystemTime::now()
|
||||||
.duration_since(UNIX_EPOCH)
|
.duration_since(UNIX_EPOCH)
|
||||||
.ok()
|
.ok()
|
||||||
.map(|duration| duration.as_secs())
|
.map(|duration| duration.as_secs())
|
||||||
.unwrap_or(0);
|
.unwrap_or(0);
|
||||||
Some(serde_json::Value::Array(
|
Some(serde_json::Value::Array(
|
||||||
keys.into_iter()
|
key_page
|
||||||
.skip(skip)
|
.items
|
||||||
.take(limit)
|
.into_iter()
|
||||||
.map(|key| {
|
.map(|key| {
|
||||||
state.build_admin_provider_key_response(
|
state.build_admin_provider_key_response(
|
||||||
&key,
|
&key,
|
||||||
|
|||||||
@@ -55,6 +55,15 @@ impl<'a> AdminAppState<'a> {
|
|||||||
crate::async_task::read_video_task_page(self.app, filter, page, page_size).await
|
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(
|
pub(crate) async fn read_video_task_stats(
|
||||||
&self,
|
&self,
|
||||||
filter: &aether_data_contracts::repository::video_tasks::VideoTaskQueryFilter,
|
filter: &aether_data_contracts::repository::video_tasks::VideoTaskQueryFilter,
|
||||||
|
|||||||
@@ -69,6 +69,18 @@ impl<'a> AdminAppState<'a> {
|
|||||||
.await
|
.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(
|
pub(crate) async fn list_provider_catalog_keys_by_ids(
|
||||||
&self,
|
&self,
|
||||||
key_ids: &[String],
|
key_ids: &[String],
|
||||||
|
|||||||
@@ -148,7 +148,7 @@ pub(crate) async fn build_admin_module_runtime_state(
|
|||||||
false
|
false
|
||||||
} else {
|
} else {
|
||||||
state
|
state
|
||||||
.list_provider_catalog_keys_by_provider_ids(&provider_ids)
|
.list_provider_catalog_key_summaries_by_provider_ids(&provider_ids)
|
||||||
.await
|
.await
|
||||||
.ok()
|
.ok()
|
||||||
.unwrap_or_default()
|
.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();
|
let mut key_counts_by_format = BTreeMap::<String, usize>::new();
|
||||||
if options.include_key_count && !provider_ids.is_empty() {
|
if options.include_key_count && !provider_ids.is_empty() {
|
||||||
let keys = state
|
let keys = state
|
||||||
.list_provider_catalog_keys_by_provider_ids(&provider_ids)
|
.list_provider_catalog_key_summaries_by_provider_ids(&provider_ids)
|
||||||
.await
|
.await
|
||||||
.ok()
|
.ok()
|
||||||
.unwrap_or_default();
|
.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!(
|
let (endpoints_result, keys_result) = tokio::join!(
|
||||||
state.list_provider_catalog_endpoints_by_provider_ids(&provider_ids),
|
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
|
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 DB_MAINTENANCE_MINUTE: u32 = 0;
|
||||||
const MAINTENANCE_DEFAULT_TIMEZONE: &str = "Asia/Shanghai";
|
const MAINTENANCE_DEFAULT_TIMEZONE: &str = "Asia/Shanghai";
|
||||||
const DB_MAINTENANCE_TABLES: &[&str] = &["usage", "request_candidates", "audit_logs"];
|
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#"
|
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 (
|
INSERT INTO wallet_daily_usage_ledgers (
|
||||||
id,
|
id,
|
||||||
wallet_id,
|
wallet_id,
|
||||||
@@ -148,11 +148,23 @@ INSERT INTO wallet_daily_usage_ledgers (
|
|||||||
created_at,
|
created_at,
|
||||||
updated_at
|
updated_at
|
||||||
)
|
)
|
||||||
VALUES (
|
SELECT
|
||||||
$1, $2, $3, $4, $5,
|
md5(CONCAT('wallet-daily-usage:', aggregated.wallet_id, ':', CAST($3 AS TEXT), ':', $4)),
|
||||||
$6, $7, $8, $9, $10,
|
aggregated.wallet_id,
|
||||||
$11, $12, $13, $14, $15
|
$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)
|
ON CONFLICT (wallet_id, billing_date, billing_timezone)
|
||||||
DO UPDATE SET
|
DO UPDATE SET
|
||||||
total_cost_usd = EXCLUDED.total_cost_usd,
|
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#"
|
const SELECT_USAGE_BODY_COMPRESSION_BATCH_SQL: &str = r#"
|
||||||
SELECT
|
SELECT
|
||||||
id,
|
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
|
FROM usage
|
||||||
WHERE created_at < $1
|
WHERE created_at < $1
|
||||||
AND ($2::timestamptz IS NULL OR created_at >= $2)
|
AND ($2::timestamptz IS NULL OR created_at >= $2)
|
||||||
@@ -503,6 +505,22 @@ DO UPDATE SET
|
|||||||
END,
|
END,
|
||||||
updated_at = NOW()
|
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#"
|
const UPDATE_USAGE_BODY_COMPRESSION_SQL: &str = r#"
|
||||||
UPDATE usage
|
UPDATE usage
|
||||||
SET request_body = NULL,
|
SET request_body = NULL,
|
||||||
@@ -695,24 +713,24 @@ DO UPDATE SET
|
|||||||
aggregated_at = EXCLUDED.aggregated_at,
|
aggregated_at = EXCLUDED.aggregated_at,
|
||||||
updated_at = EXCLUDED.updated_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#"
|
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 (
|
INSERT INTO stats_daily_model (
|
||||||
id,
|
id,
|
||||||
date,
|
date,
|
||||||
@@ -727,7 +745,20 @@ INSERT INTO stats_daily_model (
|
|||||||
created_at,
|
created_at,
|
||||||
updated_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)
|
ON CONFLICT (date, model)
|
||||||
DO UPDATE SET
|
DO UPDATE SET
|
||||||
total_requests = EXCLUDED.total_requests,
|
total_requests = EXCLUDED.total_requests,
|
||||||
@@ -739,21 +770,21 @@ DO UPDATE SET
|
|||||||
avg_response_time_ms = EXCLUDED.avg_response_time_ms,
|
avg_response_time_ms = EXCLUDED.avg_response_time_ms,
|
||||||
updated_at = EXCLUDED.updated_at
|
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#"
|
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 (
|
INSERT INTO stats_daily_provider (
|
||||||
id,
|
id,
|
||||||
date,
|
date,
|
||||||
@@ -767,7 +798,19 @@ INSERT INTO stats_daily_provider (
|
|||||||
created_at,
|
created_at,
|
||||||
updated_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)
|
ON CONFLICT (date, provider_name)
|
||||||
DO UPDATE SET
|
DO UPDATE SET
|
||||||
total_requests = EXCLUDED.total_requests,
|
total_requests = EXCLUDED.total_requests,
|
||||||
@@ -778,24 +821,34 @@ DO UPDATE SET
|
|||||||
total_cost = EXCLUDED.total_cost,
|
total_cost = EXCLUDED.total_cost,
|
||||||
updated_at = EXCLUDED.updated_at
|
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#"
|
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 (
|
INSERT INTO stats_daily_api_key (
|
||||||
id,
|
id,
|
||||||
api_key_id,
|
api_key_id,
|
||||||
@@ -812,7 +865,22 @@ INSERT INTO stats_daily_api_key (
|
|||||||
created_at,
|
created_at,
|
||||||
updated_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)
|
ON CONFLICT (api_key_id, date)
|
||||||
DO UPDATE SET
|
DO UPDATE SET
|
||||||
api_key_name = COALESCE(EXCLUDED.api_key_name, stats_daily_api_key.api_key_name),
|
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
|
DELETE FROM stats_daily_error
|
||||||
WHERE date = $1
|
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#"
|
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 (
|
INSERT INTO stats_daily_error (
|
||||||
id,
|
id,
|
||||||
date,
|
date,
|
||||||
@@ -853,32 +921,56 @@ INSERT INTO stats_daily_error (
|
|||||||
created_at,
|
created_at,
|
||||||
updated_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
|
SELECT
|
||||||
user_id,
|
md5(
|
||||||
MAX(username) AS username,
|
CONCAT(
|
||||||
CAST(COUNT(id) AS BIGINT) AS total_requests,
|
'stats-daily-error:',
|
||||||
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($1 AS TEXT),
|
||||||
CAST(COALESCE(SUM(input_tokens), 0) AS BIGINT) AS input_tokens,
|
':',
|
||||||
CAST(COALESCE(SUM(output_tokens), 0) AS BIGINT) AS output_tokens,
|
aggregated.error_category,
|
||||||
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,
|
COALESCE(aggregated.provider_name, ''),
|
||||||
CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost
|
':',
|
||||||
FROM usage
|
COALESCE(aggregated.model, '')
|
||||||
WHERE created_at >= $1
|
)
|
||||||
AND created_at < $2
|
),
|
||||||
AND user_id IS NOT NULL
|
$1,
|
||||||
GROUP BY user_id
|
aggregated.error_category,
|
||||||
|
aggregated.provider_name,
|
||||||
|
aggregated.model,
|
||||||
|
aggregated.total_count,
|
||||||
|
$3,
|
||||||
|
$3
|
||||||
|
FROM aggregated
|
||||||
"#;
|
"#;
|
||||||
const UPSERT_STATS_USER_DAILY_SQL: &str = r#"
|
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 (
|
INSERT INTO stats_user_daily (
|
||||||
id,
|
id,
|
||||||
user_id,
|
user_id,
|
||||||
@@ -895,7 +987,24 @@ INSERT INTO stats_user_daily (
|
|||||||
created_at,
|
created_at,
|
||||||
updated_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)
|
ON CONFLICT (user_id, date)
|
||||||
DO UPDATE SET
|
DO UPDATE SET
|
||||||
username = COALESCE(EXCLUDED.username, stats_user_daily.username),
|
username = COALESCE(EXCLUDED.username, stats_user_daily.username),
|
||||||
@@ -1034,21 +1143,29 @@ DO UPDATE SET
|
|||||||
aggregated_at = EXCLUDED.aggregated_at,
|
aggregated_at = EXCLUDED.aggregated_at,
|
||||||
updated_at = EXCLUDED.updated_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#"
|
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 (
|
INSERT INTO stats_hourly_user (
|
||||||
id,
|
id,
|
||||||
hour_utc,
|
hour_utc,
|
||||||
@@ -1062,7 +1179,19 @@ INSERT INTO stats_hourly_user (
|
|||||||
created_at,
|
created_at,
|
||||||
updated_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)
|
ON CONFLICT (hour_utc, user_id)
|
||||||
DO UPDATE SET
|
DO UPDATE SET
|
||||||
total_requests = EXCLUDED.total_requests,
|
total_requests = EXCLUDED.total_requests,
|
||||||
@@ -1073,20 +1202,22 @@ DO UPDATE SET
|
|||||||
total_cost = EXCLUDED.total_cost,
|
total_cost = EXCLUDED.total_cost,
|
||||||
updated_at = EXCLUDED.updated_at
|
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#"
|
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 (
|
INSERT INTO stats_hourly_model (
|
||||||
id,
|
id,
|
||||||
hour_utc,
|
hour_utc,
|
||||||
@@ -1099,7 +1230,18 @@ INSERT INTO stats_hourly_model (
|
|||||||
created_at,
|
created_at,
|
||||||
updated_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)
|
ON CONFLICT (hour_utc, model)
|
||||||
DO UPDATE SET
|
DO UPDATE SET
|
||||||
total_requests = EXCLUDED.total_requests,
|
total_requests = EXCLUDED.total_requests,
|
||||||
@@ -1109,19 +1251,21 @@ DO UPDATE SET
|
|||||||
avg_response_time_ms = EXCLUDED.avg_response_time_ms,
|
avg_response_time_ms = EXCLUDED.avg_response_time_ms,
|
||||||
updated_at = EXCLUDED.updated_at
|
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#"
|
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 (
|
INSERT INTO stats_hourly_provider (
|
||||||
id,
|
id,
|
||||||
hour_utc,
|
hour_utc,
|
||||||
@@ -1133,7 +1277,17 @@ INSERT INTO stats_hourly_provider (
|
|||||||
created_at,
|
created_at,
|
||||||
updated_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)
|
ON CONFLICT (hour_utc, provider_name)
|
||||||
DO UPDATE SET
|
DO UPDATE SET
|
||||||
total_requests = EXCLUDED.total_requests,
|
total_requests = EXCLUDED.total_requests,
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
use std::collections::HashSet;
|
use std::collections::HashSet;
|
||||||
|
|
||||||
use chrono::Utc;
|
use chrono::Utc;
|
||||||
|
use futures_util::TryStreamExt;
|
||||||
use sqlx::Row;
|
use sqlx::Row;
|
||||||
|
|
||||||
use crate::data::GatewayDataState;
|
use crate::data::GatewayDataState;
|
||||||
@@ -59,13 +60,18 @@ pub(crate) async fn cleanup_stale_pending_requests_once(
|
|||||||
|
|
||||||
loop {
|
loop {
|
||||||
let mut tx = pool.begin().await.map_err(postgres_error)?;
|
let mut tx = pool.begin().await.map_err(postgres_error)?;
|
||||||
let stale_rows = sqlx::query(SELECT_STALE_PENDING_USAGE_BATCH_SQL)
|
let stale_rows = {
|
||||||
.bind(active_statuses.clone())
|
let mut stale_rows_stream = sqlx::query(SELECT_STALE_PENDING_USAGE_BATCH_SQL)
|
||||||
.bind(cutoff_time)
|
.bind(active_statuses.clone())
|
||||||
.bind(i64::try_from(batch_size).unwrap_or(i64::MAX))
|
.bind(cutoff_time)
|
||||||
.fetch_all(&mut *tx)
|
.bind(i64::try_from(batch_size).unwrap_or(i64::MAX))
|
||||||
.await
|
.fetch(&mut *tx);
|
||||||
.map_err(postgres_error)?;
|
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() {
|
if stale_rows.is_empty() {
|
||||||
tx.rollback().await.map_err(postgres_error)?;
|
tx.rollback().await.map_err(postgres_error)?;
|
||||||
break;
|
break;
|
||||||
@@ -93,14 +99,18 @@ pub(crate) async fn cleanup_stale_pending_requests_once(
|
|||||||
let completed_request_ids = if request_ids.is_empty() {
|
let completed_request_ids = if request_ids.is_empty() {
|
||||||
HashSet::new()
|
HashSet::new()
|
||||||
} else {
|
} else {
|
||||||
sqlx::query(SELECT_COMPLETED_PENDING_REQUEST_IDS_SQL)
|
{
|
||||||
.bind(request_ids)
|
let mut completed_rows = sqlx::query(SELECT_COMPLETED_PENDING_REQUEST_IDS_SQL)
|
||||||
.fetch_all(&mut *tx)
|
.bind(request_ids)
|
||||||
.await
|
.fetch(&mut *tx);
|
||||||
.map_err(postgres_error)?
|
let mut completed_request_ids = HashSet::new();
|
||||||
.into_iter()
|
while let Some(row) = completed_rows.try_next().await.map_err(postgres_error)? {
|
||||||
.filter_map(|row| row.try_get::<String, _>("request_id").ok())
|
if let Ok(request_id) = row.try_get::<String, _>("request_id") {
|
||||||
.collect::<HashSet<_>>()
|
completed_request_ids.insert(request_id);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
completed_request_ids
|
||||||
|
}
|
||||||
};
|
};
|
||||||
let plan = plan_pending_cleanup_batch(stale_rows, &completed_request_ids, timeout_minutes);
|
let plan = plan_pending_cleanup_batch(stale_rows, &completed_request_ids, timeout_minutes);
|
||||||
let now = Utc::now();
|
let now = Utc::now();
|
||||||
|
|||||||
@@ -1,5 +1,3 @@
|
|||||||
use std::collections::HashMap;
|
|
||||||
|
|
||||||
use chrono::{DateTime, Utc};
|
use chrono::{DateTime, Utc};
|
||||||
use sqlx::Row;
|
use sqlx::Row;
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
@@ -10,15 +8,13 @@ use aether_data_contracts::DataLayerError;
|
|||||||
use super::{
|
use super::{
|
||||||
postgres_error, stats_aggregation_target_day, system_config_bool, PercentileSummary,
|
postgres_error, stats_aggregation_target_day, system_config_bool, PercentileSummary,
|
||||||
StatsAggregationSummary, DELETE_STATS_DAILY_ERRORS_FOR_DATE_SQL, INSERT_STATS_DAILY_ERROR_SQL,
|
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,
|
INSERT_STATS_SUMMARY_SQL, SELECT_EXISTING_STATS_SUMMARY_ID_SQL,
|
||||||
SELECT_STATS_DAILY_AGGREGATE_SQL, SELECT_STATS_DAILY_API_KEY_AGGREGATES_SQL,
|
SELECT_STATS_DAILY_AGGREGATE_SQL, SELECT_STATS_DAILY_FALLBACK_COUNT_SQL,
|
||||||
SELECT_STATS_DAILY_ERROR_AGGREGATES_SQL, SELECT_STATS_DAILY_FALLBACK_COUNT_SQL,
|
SELECT_STATS_DAILY_FIRST_BYTE_PERCENTILES_SQL,
|
||||||
SELECT_STATS_DAILY_FIRST_BYTE_PERCENTILES_SQL, SELECT_STATS_DAILY_MODEL_AGGREGATES_SQL,
|
SELECT_STATS_DAILY_RESPONSE_TIME_PERCENTILES_SQL, SELECT_STATS_SUMMARY_ENTITY_COUNTS_SQL,
|
||||||
SELECT_STATS_DAILY_PROVIDER_AGGREGATES_SQL, SELECT_STATS_DAILY_RESPONSE_TIME_PERCENTILES_SQL,
|
SELECT_STATS_SUMMARY_TOTALS_SQL, UPDATE_STATS_SUMMARY_SQL, UPSERT_STATS_DAILY_API_KEY_SQL,
|
||||||
SELECT_STATS_SUMMARY_ENTITY_COUNTS_SQL, SELECT_STATS_SUMMARY_TOTALS_SQL,
|
UPSERT_STATS_DAILY_MODEL_SQL, UPSERT_STATS_DAILY_PROVIDER_SQL, UPSERT_STATS_DAILY_SQL,
|
||||||
SELECT_STATS_USER_DAILY_AGGREGATES_SQL, UPDATE_STATS_SUMMARY_SQL,
|
UPSERT_STATS_USER_DAILY_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(
|
pub(super) async fn perform_stats_aggregation_once(
|
||||||
@@ -226,31 +222,15 @@ async fn upsert_stats_daily_model_rows(
|
|||||||
day_end_utc: DateTime<Utc>,
|
day_end_utc: DateTime<Utc>,
|
||||||
now_utc: DateTime<Utc>,
|
now_utc: DateTime<Utc>,
|
||||||
) -> Result<usize, sqlx::Error> {
|
) -> 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_start_utc)
|
||||||
.bind(day_end_utc)
|
.bind(day_end_utc)
|
||||||
.fetch_all(&mut **tx)
|
.bind(now_utc)
|
||||||
.await?;
|
.execute(&mut **tx)
|
||||||
|
.await?
|
||||||
|
.rows_affected();
|
||||||
|
|
||||||
for row in &rows {
|
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||||
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())
|
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn upsert_stats_daily_provider_rows(
|
async fn upsert_stats_daily_provider_rows(
|
||||||
@@ -259,30 +239,15 @@ async fn upsert_stats_daily_provider_rows(
|
|||||||
day_end_utc: DateTime<Utc>,
|
day_end_utc: DateTime<Utc>,
|
||||||
now_utc: DateTime<Utc>,
|
now_utc: DateTime<Utc>,
|
||||||
) -> Result<usize, sqlx::Error> {
|
) -> 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_start_utc)
|
||||||
.bind(day_end_utc)
|
.bind(day_end_utc)
|
||||||
.fetch_all(&mut **tx)
|
.bind(now_utc)
|
||||||
.await?;
|
.execute(&mut **tx)
|
||||||
|
.await?
|
||||||
|
.rows_affected();
|
||||||
|
|
||||||
for row in &rows {
|
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||||
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())
|
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn upsert_stats_daily_api_key_rows(
|
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>,
|
day_end_utc: DateTime<Utc>,
|
||||||
now_utc: DateTime<Utc>,
|
now_utc: DateTime<Utc>,
|
||||||
) -> Result<usize, sqlx::Error> {
|
) -> 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_start_utc)
|
||||||
.bind(day_end_utc)
|
.bind(day_end_utc)
|
||||||
.fetch_all(&mut **tx)
|
.bind(now_utc)
|
||||||
.await?;
|
.execute(&mut **tx)
|
||||||
|
.await?
|
||||||
|
.rows_affected();
|
||||||
|
|
||||||
for row in &rows {
|
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||||
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())
|
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn refresh_stats_daily_error_rows(
|
async fn refresh_stats_daily_error_rows(
|
||||||
@@ -332,27 +277,15 @@ async fn refresh_stats_daily_error_rows(
|
|||||||
.bind(day_start_utc)
|
.bind(day_start_utc)
|
||||||
.execute(&mut **tx)
|
.execute(&mut **tx)
|
||||||
.await?;
|
.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_start_utc)
|
||||||
.bind(day_end_utc)
|
.bind(day_end_utc)
|
||||||
.fetch_all(&mut **tx)
|
.bind(now_utc)
|
||||||
.await?;
|
.execute(&mut **tx)
|
||||||
|
.await?
|
||||||
|
.rows_affected();
|
||||||
|
|
||||||
for row in &rows {
|
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||||
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())
|
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn upsert_stats_user_daily_rows(
|
async fn upsert_stats_user_daily_rows(
|
||||||
@@ -361,87 +294,15 @@ async fn upsert_stats_user_daily_rows(
|
|||||||
day_end_utc: DateTime<Utc>,
|
day_end_utc: DateTime<Utc>,
|
||||||
now_utc: DateTime<Utc>,
|
now_utc: DateTime<Utc>,
|
||||||
) -> Result<usize, sqlx::Error> {
|
) -> Result<usize, sqlx::Error> {
|
||||||
let active_user_ids = sqlx::query(SELECT_ACTIVE_USER_IDS_SQL)
|
let rows_affected = sqlx::query(UPSERT_STATS_USER_DAILY_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)
|
|
||||||
.bind(day_start_utc)
|
.bind(day_start_utc)
|
||||||
.bind(day_end_utc)
|
.bind(day_end_utc)
|
||||||
.fetch_all(&mut **tx)
|
.bind(now_utc)
|
||||||
.await?;
|
.execute(&mut **tx)
|
||||||
let mut aggregated_by_user = HashMap::with_capacity(aggregated_rows.len());
|
.await?
|
||||||
for row in aggregated_rows {
|
.rows_affected();
|
||||||
let user_id = row.try_get::<String, _>("user_id")?;
|
|
||||||
aggregated_by_user.insert(user_id, row);
|
|
||||||
}
|
|
||||||
|
|
||||||
for user_id in &active_user_ids {
|
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||||
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())
|
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn refresh_stats_summary_row(
|
async fn refresh_stats_summary_row(
|
||||||
|
|||||||
@@ -7,9 +7,8 @@ use aether_data_contracts::DataLayerError;
|
|||||||
|
|
||||||
use super::{
|
use super::{
|
||||||
stats_hourly_aggregation_target_hour, system_config_bool, SELECT_STATS_HOURLY_AGGREGATE_SQL,
|
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,
|
UPSERT_STATS_HOURLY_MODEL_SQL, UPSERT_STATS_HOURLY_PROVIDER_SQL, UPSERT_STATS_HOURLY_SQL,
|
||||||
SELECT_STATS_HOURLY_USER_AGGREGATES_SQL, UPSERT_STATS_HOURLY_MODEL_SQL,
|
UPSERT_STATS_HOURLY_USER_SQL,
|
||||||
UPSERT_STATS_HOURLY_PROVIDER_SQL, UPSERT_STATS_HOURLY_SQL, UPSERT_STATS_HOURLY_USER_SQL,
|
|
||||||
};
|
};
|
||||||
|
|
||||||
#[derive(Debug, Clone, PartialEq)]
|
#[derive(Debug, Clone, PartialEq)]
|
||||||
@@ -115,51 +114,16 @@ async fn upsert_stats_hourly_user_rows(
|
|||||||
hour_end: DateTime<Utc>,
|
hour_end: DateTime<Utc>,
|
||||||
now_utc: DateTime<Utc>,
|
now_utc: DateTime<Utc>,
|
||||||
) -> Result<usize, DataLayerError> {
|
) -> 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_utc)
|
||||||
.bind(hour_end)
|
.bind(hour_end)
|
||||||
.fetch_all(&mut **tx)
|
.bind(now_utc)
|
||||||
|
.execute(&mut **tx)
|
||||||
.await
|
.await
|
||||||
.map_err(postgres_error)?;
|
.map_err(postgres_error)?
|
||||||
|
.rows_affected();
|
||||||
|
|
||||||
for row in &rows {
|
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||||
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())
|
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn upsert_stats_hourly_model_rows(
|
async fn upsert_stats_hourly_model_rows(
|
||||||
@@ -168,54 +132,16 @@ async fn upsert_stats_hourly_model_rows(
|
|||||||
hour_end: DateTime<Utc>,
|
hour_end: DateTime<Utc>,
|
||||||
now_utc: DateTime<Utc>,
|
now_utc: DateTime<Utc>,
|
||||||
) -> Result<usize, DataLayerError> {
|
) -> 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_utc)
|
||||||
.bind(hour_end)
|
.bind(hour_end)
|
||||||
.fetch_all(&mut **tx)
|
.bind(now_utc)
|
||||||
|
.execute(&mut **tx)
|
||||||
.await
|
.await
|
||||||
.map_err(postgres_error)?;
|
.map_err(postgres_error)?
|
||||||
let mut inserted = 0usize;
|
.rows_affected();
|
||||||
|
|
||||||
for row in &rows {
|
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||||
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)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn upsert_stats_hourly_provider_rows(
|
async fn upsert_stats_hourly_provider_rows(
|
||||||
@@ -224,50 +150,16 @@ async fn upsert_stats_hourly_provider_rows(
|
|||||||
hour_end: DateTime<Utc>,
|
hour_end: DateTime<Utc>,
|
||||||
now_utc: DateTime<Utc>,
|
now_utc: DateTime<Utc>,
|
||||||
) -> Result<usize, DataLayerError> {
|
) -> 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_utc)
|
||||||
.bind(hour_end)
|
.bind(hour_end)
|
||||||
.fetch_all(&mut **tx)
|
.bind(now_utc)
|
||||||
|
.execute(&mut **tx)
|
||||||
.await
|
.await
|
||||||
.map_err(postgres_error)?;
|
.map_err(postgres_error)?
|
||||||
let mut inserted = 0usize;
|
.rows_affected();
|
||||||
|
|
||||||
for row in &rows {
|
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||||
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)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
fn postgres_error(error: sqlx::Error) -> DataLayerError {
|
fn postgres_error(error: sqlx::Error) -> DataLayerError {
|
||||||
|
|||||||
@@ -29,7 +29,7 @@ use super::{
|
|||||||
usage_cleanup_window, wallet_daily_usage_aggregation_target, AppState, DbMaintenanceRunSummary,
|
usage_cleanup_window, wallet_daily_usage_aggregation_target, AppState, DbMaintenanceRunSummary,
|
||||||
FailedPendingUsageRow, GatewayDataState, ProxyUpgradeRolloutProbeConfig, StalePendingUsageRow,
|
FailedPendingUsageRow, GatewayDataState, ProxyUpgradeRolloutProbeConfig, StalePendingUsageRow,
|
||||||
UsageCleanupSettings, DELETE_STALE_WALLET_DAILY_USAGE_LEDGERS_SQL,
|
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,
|
UPDATE_FAILED_VOID_STALE_USAGE_SQL, USAGE_CLEANUP_HOUR, USAGE_CLEANUP_MINUTE,
|
||||||
WALLET_DAILY_USAGE_AGGREGATION_HOUR, WALLET_DAILY_USAGE_AGGREGATION_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]
|
#[test]
|
||||||
fn wallet_daily_usage_queries_use_settlement_snapshots_for_wallet_identity() {
|
fn wallet_daily_usage_queries_use_settlement_snapshots_for_wallet_identity() {
|
||||||
assert!(
|
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"));
|
.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.contains("JOIN usage_settlement_snapshots"));
|
||||||
assert!(DELETE_STALE_WALLET_DAILY_USAGE_LEDGERS_SQL
|
assert!(DELETE_STALE_WALLET_DAILY_USAGE_LEDGERS_SQL
|
||||||
|
|||||||
@@ -6,6 +6,7 @@ use aether_data_contracts::repository::usage::{
|
|||||||
use aether_data_contracts::DataLayerError;
|
use aether_data_contracts::DataLayerError;
|
||||||
use chrono::{DateTime, Utc};
|
use chrono::{DateTime, Utc};
|
||||||
use flate2::{write::GzEncoder, Compression};
|
use flate2::{write::GzEncoder, Compression};
|
||||||
|
use futures_util::TryStreamExt;
|
||||||
use serde_json::{Map, Value};
|
use serde_json::{Map, Value};
|
||||||
use sqlx::Row;
|
use sqlx::Row;
|
||||||
use tracing::warn;
|
use tracing::warn;
|
||||||
@@ -21,7 +22,8 @@ use super::{
|
|||||||
DISABLE_EXPIRED_API_KEY_SQL, EXPIRED_API_KEY_PRE_CLEAN_BATCH_SIZE,
|
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,
|
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_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,
|
SELECT_USAGE_STALE_BODY_BATCH_SQL, UPDATE_USAGE_BODY_COMPRESSION_SQL,
|
||||||
UPDATE_USAGE_REQUEST_METADATA_SQL, UPSERT_USAGE_BODY_BLOB_SQL,
|
UPDATE_USAGE_REQUEST_METADATA_SQL, UPSERT_USAGE_BODY_BLOB_SQL,
|
||||||
UPSERT_USAGE_HTTP_AUDIT_BODY_REFS_SQL,
|
UPSERT_USAGE_HTTP_AUDIT_BODY_REFS_SQL,
|
||||||
@@ -211,23 +213,20 @@ async fn cleanup_usage_header_fields(
|
|||||||
|
|
||||||
let mut total_cleaned = 0usize;
|
let mut total_cleaned = 0usize;
|
||||||
loop {
|
loop {
|
||||||
let rows = sqlx::query(SELECT_USAGE_HEADER_BATCH_SQL)
|
let mut stream = sqlx::query(SELECT_USAGE_HEADER_BATCH_SQL)
|
||||||
.bind(cutoff_time)
|
.bind(cutoff_time)
|
||||||
.bind(newer_than)
|
.bind(newer_than)
|
||||||
.bind(i64::try_from(batch_size).unwrap_or(i64::MAX))
|
.bind(i64::try_from(batch_size).unwrap_or(i64::MAX))
|
||||||
.fetch_all(pool)
|
.fetch(pool);
|
||||||
.await
|
let mut rows = Vec::new();
|
||||||
.map_err(postgres_error)?
|
while let Some(row) = stream.try_next().await.map_err(postgres_error)? {
|
||||||
.into_iter()
|
rows.push(UsageBodyCleanupRow {
|
||||||
.map(|row| {
|
id: row.try_get::<String, _>("id").map_err(postgres_error)?,
|
||||||
Ok(UsageBodyCleanupRow {
|
request_id: row
|
||||||
id: row.try_get::<String, _>("id").map_err(postgres_error)?,
|
.try_get::<String, _>("request_id")
|
||||||
request_id: row
|
.map_err(postgres_error)?,
|
||||||
.try_get::<String, _>("request_id")
|
});
|
||||||
.map_err(postgres_error)?,
|
}
|
||||||
})
|
|
||||||
})
|
|
||||||
.collect::<Result<Vec<_>, DataLayerError>>()?;
|
|
||||||
if rows.is_empty() {
|
if rows.is_empty() {
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
@@ -279,23 +278,20 @@ async fn cleanup_usage_stale_body_fields(
|
|||||||
|
|
||||||
let mut total_cleaned = 0usize;
|
let mut total_cleaned = 0usize;
|
||||||
loop {
|
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(cutoff_time)
|
||||||
.bind(newer_than)
|
.bind(newer_than)
|
||||||
.bind(i64::try_from(batch_size).unwrap_or(i64::MAX))
|
.bind(i64::try_from(batch_size).unwrap_or(i64::MAX))
|
||||||
.fetch_all(pool)
|
.fetch(pool);
|
||||||
.await
|
let mut rows = Vec::new();
|
||||||
.map_err(postgres_error)?
|
while let Some(row) = stream.try_next().await.map_err(postgres_error)? {
|
||||||
.into_iter()
|
rows.push(UsageBodyCleanupRow {
|
||||||
.map(|row| {
|
id: row.try_get::<String, _>("id").map_err(postgres_error)?,
|
||||||
Ok(UsageBodyCleanupRow {
|
request_id: row
|
||||||
id: row.try_get::<String, _>("id").map_err(postgres_error)?,
|
.try_get::<String, _>("request_id")
|
||||||
request_id: row
|
.map_err(postgres_error)?,
|
||||||
.try_get::<String, _>("request_id")
|
});
|
||||||
.map_err(postgres_error)?,
|
}
|
||||||
})
|
|
||||||
})
|
|
||||||
.collect::<Result<Vec<_>, DataLayerError>>()?;
|
|
||||||
if rows.is_empty() {
|
if rows.is_empty() {
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
@@ -354,53 +350,59 @@ async fn compress_usage_body_fields(
|
|||||||
let mut no_progress_count = 0usize;
|
let mut no_progress_count = 0usize;
|
||||||
let batch_size = batch_size.clamp(1, 25);
|
let batch_size = batch_size.clamp(1, 25);
|
||||||
loop {
|
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(cutoff_time)
|
||||||
.bind(newer_than)
|
.bind(newer_than)
|
||||||
.bind(i64::try_from(batch_size).unwrap_or(i64::MAX))
|
.bind(i64::try_from(batch_size).unwrap_or(i64::MAX))
|
||||||
.fetch_all(pool)
|
.fetch(pool);
|
||||||
.await
|
let mut ids = Vec::new();
|
||||||
.map_err(postgres_error)?
|
while let Some(row) = stream.try_next().await.map_err(postgres_error)? {
|
||||||
.into_iter()
|
ids.push(row.try_get::<String, _>("id").map_err(postgres_error)?);
|
||||||
.map(|row| {
|
}
|
||||||
Ok(UsageBodyCompressionRow {
|
if ids.is_empty() {
|
||||||
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() {
|
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
|
|
||||||
let mut batch_success = 0usize;
|
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)?;
|
let detached = build_usage_body_externalization(&row)?;
|
||||||
if detached.refs.any_present() {
|
if detached.refs.any_present() {
|
||||||
let mut tx = pool.begin().await.map_err(postgres_error)?;
|
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,
|
pool: &aether_data::postgres::PostgresPool,
|
||||||
auto_delete_expired_keys: bool,
|
auto_delete_expired_keys: bool,
|
||||||
) -> Result<usize, DataLayerError> {
|
) -> Result<usize, DataLayerError> {
|
||||||
let expired_keys = sqlx::query(SELECT_EXPIRED_ACTIVE_API_KEYS_SQL)
|
let mut expired_keys = sqlx::query(SELECT_EXPIRED_ACTIVE_API_KEYS_SQL).fetch(pool);
|
||||||
.fetch_all(pool)
|
|
||||||
.await
|
|
||||||
.map_err(postgres_error)?;
|
|
||||||
let mut cleaned = 0usize;
|
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 api_key_id = row.try_get::<String, _>("id").map_err(postgres_error)?;
|
||||||
let key = ExpiredApiKeyRow {
|
let key = ExpiredApiKeyRow {
|
||||||
id: api_key_id.as_str(),
|
id: api_key_id.as_str(),
|
||||||
|
|||||||
@@ -1,14 +1,11 @@
|
|||||||
use chrono::{DateTime, Utc};
|
use chrono::{DateTime, Utc};
|
||||||
use sqlx::Row;
|
|
||||||
use uuid::Uuid;
|
|
||||||
|
|
||||||
use crate::data::GatewayDataState;
|
use crate::data::GatewayDataState;
|
||||||
use aether_data_contracts::DataLayerError;
|
use aether_data_contracts::DataLayerError;
|
||||||
|
|
||||||
use super::{
|
use super::{
|
||||||
maintenance_timezone, wallet_daily_usage_aggregation_target,
|
maintenance_timezone, wallet_daily_usage_aggregation_target,
|
||||||
DELETE_STALE_WALLET_DAILY_USAGE_LEDGERS_SQL, SELECT_WALLET_DAILY_USAGE_AGGREGATION_ROWS_SQL,
|
DELETE_STALE_WALLET_DAILY_USAGE_LEDGERS_SQL, UPSERT_WALLET_DAILY_USAGE_LEDGER_SQL,
|
||||||
UPSERT_WALLET_DAILY_USAGE_LEDGER_SQL,
|
|
||||||
};
|
};
|
||||||
|
|
||||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
#[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 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_start_utc)
|
||||||
.bind(target.window_end_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
|
.await
|
||||||
.map_err(postgres_error)?;
|
.map_err(postgres_error)?
|
||||||
for row in &rows {
|
.rows_affected();
|
||||||
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)?;
|
|
||||||
}
|
|
||||||
|
|
||||||
let deleted_stale_ledgers = sqlx::query(DELETE_STALE_WALLET_DAILY_USAGE_LEDGERS_SQL)
|
let deleted_stale_ledgers = sqlx::query(DELETE_STALE_WALLET_DAILY_USAGE_LEDGERS_SQL)
|
||||||
.bind(target.billing_date)
|
.bind(target.billing_date)
|
||||||
@@ -112,7 +65,7 @@ pub(super) async fn perform_wallet_daily_usage_aggregation_once(
|
|||||||
Ok(WalletDailyUsageAggregationSummary {
|
Ok(WalletDailyUsageAggregationSummary {
|
||||||
billing_date: target.billing_date,
|
billing_date: target.billing_date,
|
||||||
billing_timezone: target.billing_timezone,
|
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),
|
deleted_stale_ledgers: usize::try_from(deleted_stale_ledgers).unwrap_or(usize::MAX),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -4,6 +4,7 @@ use crate::{
|
|||||||
AdminBillingRuleWriteInput, GatewayError, LocalMutationOutcome,
|
AdminBillingRuleWriteInput, GatewayError, LocalMutationOutcome,
|
||||||
};
|
};
|
||||||
use aether_data::postgres::PostgresPool;
|
use aether_data::postgres::PostgresPool;
|
||||||
|
use futures_util::TryStreamExt;
|
||||||
use sqlx::Row;
|
use sqlx::Row;
|
||||||
|
|
||||||
fn internal(err: impl ToString) -> GatewayError {
|
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 offset = u64::from(page.saturating_sub(1) * page_size);
|
||||||
let rows = sqlx::query(
|
let mut rows = sqlx::query(
|
||||||
r#"
|
r#"
|
||||||
SELECT
|
SELECT
|
||||||
id,
|
id,
|
||||||
@@ -162,16 +163,13 @@ LIMIT $4
|
|||||||
.bind(is_enabled)
|
.bind(is_enabled)
|
||||||
.bind(i64::try_from(offset).map_err(|err| GatewayError::Internal(err.to_string()))?)
|
.bind(i64::try_from(offset).map_err(|err| GatewayError::Internal(err.to_string()))?)
|
||||||
.bind(i64::from(page_size))
|
.bind(i64::from(page_size))
|
||||||
.fetch_all(pool)
|
.fetch(pool);
|
||||||
.await
|
let mut items = Vec::new();
|
||||||
.map_err(internal)?;
|
while let Some(row) = rows.try_next().await.map_err(internal)? {
|
||||||
|
items.push(admin_billing_rule_from_row(&row)?);
|
||||||
|
}
|
||||||
|
|
||||||
Ok((
|
Ok((items, total))
|
||||||
rows.iter()
|
|
||||||
.map(admin_billing_rule_from_row)
|
|
||||||
.collect::<Result<Vec<_>, _>>()?,
|
|
||||||
total,
|
|
||||||
))
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn find_admin_billing_rule(
|
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 offset = u64::from(page.saturating_sub(1) * page_size);
|
||||||
let rows = sqlx::query(
|
let mut rows = sqlx::query(
|
||||||
r#"
|
r#"
|
||||||
SELECT
|
SELECT
|
||||||
id,
|
id,
|
||||||
@@ -406,16 +404,13 @@ LIMIT $6
|
|||||||
.bind(is_enabled)
|
.bind(is_enabled)
|
||||||
.bind(i64::try_from(offset).map_err(|err| GatewayError::Internal(err.to_string()))?)
|
.bind(i64::try_from(offset).map_err(|err| GatewayError::Internal(err.to_string()))?)
|
||||||
.bind(i64::from(page_size))
|
.bind(i64::from(page_size))
|
||||||
.fetch_all(pool)
|
.fetch(pool);
|
||||||
.await
|
let mut items = Vec::new();
|
||||||
.map_err(internal)?;
|
while let Some(row) = rows.try_next().await.map_err(internal)? {
|
||||||
|
items.push(admin_billing_collector_from_row(&row)?);
|
||||||
|
}
|
||||||
|
|
||||||
Ok((
|
Ok((items, total))
|
||||||
rows.iter()
|
|
||||||
.map(admin_billing_collector_from_row)
|
|
||||||
.collect::<Result<Vec<_>, _>>()?,
|
|
||||||
total,
|
|
||||||
))
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn find_admin_billing_collector(
|
pub(crate) async fn find_admin_billing_collector(
|
||||||
|
|||||||
@@ -2,6 +2,7 @@ use std::collections::BTreeMap;
|
|||||||
|
|
||||||
use aether_data::postgres::PostgresPool;
|
use aether_data::postgres::PostgresPool;
|
||||||
use chrono::{DateTime, Utc};
|
use chrono::{DateTime, Utc};
|
||||||
|
use futures_util::TryStreamExt;
|
||||||
use serde_json::{json, Value};
|
use serde_json::{json, Value};
|
||||||
use sqlx::Row;
|
use sqlx::Row;
|
||||||
|
|
||||||
@@ -36,7 +37,7 @@ WHERE a.created_at >= $1
|
|||||||
.await
|
.await
|
||||||
.map_err(|err| GatewayError::Internal(format!("admin audit logs count failed: {err}")))?;
|
.map_err(|err| GatewayError::Internal(format!("admin audit logs count failed: {err}")))?;
|
||||||
|
|
||||||
let rows = sqlx::query(
|
let mut rows = sqlx::query(
|
||||||
r#"
|
r#"
|
||||||
SELECT
|
SELECT
|
||||||
a.id,
|
a.id,
|
||||||
@@ -64,21 +65,24 @@ LIMIT $4 OFFSET $5
|
|||||||
.bind(event_type)
|
.bind(event_type)
|
||||||
.bind(i64::try_from(limit).unwrap_or(i64::MAX))
|
.bind(i64::try_from(limit).unwrap_or(i64::MAX))
|
||||||
.bind(i64::try_from(offset).unwrap_or(i64::MAX))
|
.bind(i64::try_from(offset).unwrap_or(i64::MAX))
|
||||||
.fetch_all(pool)
|
.fetch(pool);
|
||||||
.await
|
let mut items = Vec::new();
|
||||||
.map_err(|err| GatewayError::Internal(format!("admin audit logs read failed: {err}")))?;
|
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((
|
Ok((items, usize::try_from(total.max(0)).unwrap_or(usize::MAX)))
|
||||||
rows.into_iter().map(admin_audit_log_row_to_json).collect(),
|
|
||||||
usize::try_from(total.max(0)).unwrap_or(usize::MAX),
|
|
||||||
))
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn list_admin_suspicious_activities(
|
pub(crate) async fn list_admin_suspicious_activities(
|
||||||
pool: &PostgresPool,
|
pool: &PostgresPool,
|
||||||
cutoff_time: DateTime<Utc>,
|
cutoff_time: DateTime<Utc>,
|
||||||
) -> Result<Vec<Value>, GatewayError> {
|
) -> Result<Vec<Value>, GatewayError> {
|
||||||
let rows = sqlx::query(
|
let mut rows = sqlx::query(
|
||||||
r#"
|
r#"
|
||||||
SELECT
|
SELECT
|
||||||
id,
|
id,
|
||||||
@@ -102,13 +106,15 @@ LIMIT 100
|
|||||||
"login_failed",
|
"login_failed",
|
||||||
"request_rate_limited",
|
"request_rate_limited",
|
||||||
])
|
])
|
||||||
.fetch_all(pool)
|
.fetch(pool);
|
||||||
.await
|
let mut items = Vec::new();
|
||||||
.map_err(|err| {
|
while let Some(row) = rows.try_next().await.map_err(|err| {
|
||||||
GatewayError::Internal(format!("admin suspicious activities read failed: {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(
|
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,
|
user_id: &str,
|
||||||
cutoff_time: DateTime<Utc>,
|
cutoff_time: DateTime<Utc>,
|
||||||
) -> Result<BTreeMap<String, u64>, GatewayError> {
|
) -> Result<BTreeMap<String, u64>, GatewayError> {
|
||||||
let rows = sqlx::query(
|
let mut rows = sqlx::query(
|
||||||
r#"
|
r#"
|
||||||
SELECT event_type, COUNT(*)::bigint AS count
|
SELECT event_type, COUNT(*)::bigint AS count
|
||||||
FROM audit_logs
|
FROM audit_logs
|
||||||
@@ -127,22 +133,25 @@ GROUP BY event_type
|
|||||||
)
|
)
|
||||||
.bind(user_id)
|
.bind(user_id)
|
||||||
.bind(cutoff_time)
|
.bind(cutoff_time)
|
||||||
.fetch_all(pool)
|
.fetch(pool);
|
||||||
.await
|
let mut counts = BTreeMap::new();
|
||||||
.map_err(|err| GatewayError::Internal(format!("admin user behavior read failed: {err}")))?;
|
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
|
Ok(counts)
|
||||||
.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())
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn list_user_audit_logs(
|
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#"
|
r#"
|
||||||
SELECT id, event_type, description, ip_address, status_code, created_at
|
SELECT id, event_type, description, ip_address, status_code, created_at
|
||||||
FROM audit_logs
|
FROM audit_logs
|
||||||
@@ -192,21 +201,17 @@ LIMIT $4 OFFSET $5
|
|||||||
.bind(event_type)
|
.bind(event_type)
|
||||||
.bind(i64::try_from(limit).unwrap_or(i64::MAX))
|
.bind(i64::try_from(limit).unwrap_or(i64::MAX))
|
||||||
.bind(i64::try_from(offset).unwrap_or(i64::MAX))
|
.bind(i64::try_from(offset).unwrap_or(i64::MAX))
|
||||||
.fetch_all(pool)
|
.fetch(pool);
|
||||||
.await
|
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,
|
items.push(user_audit_log_row_to_json(row));
|
||||||
Err(err) => {
|
}
|
||||||
return Err(GatewayError::Internal(format!(
|
|
||||||
"user audit logs read failed: {err}"
|
|
||||||
)))
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
Ok((
|
Ok((items, total))
|
||||||
rows.into_iter().map(user_audit_log_row_to_json).collect(),
|
|
||||||
total,
|
|
||||||
))
|
|
||||||
}
|
}
|
||||||
|
|
||||||
fn admin_audit_log_row_to_json(row: sqlx::postgres::PgRow) -> Value {
|
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()))
|
.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(
|
pub(crate) async fn list_provider_catalog_keys_by_ids(
|
||||||
&self,
|
&self,
|
||||||
key_ids: &[String],
|
key_ids: &[String],
|
||||||
|
|||||||
@@ -120,6 +120,18 @@ impl AppState {
|
|||||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
.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(
|
pub(crate) async fn count_video_tasks(
|
||||||
&self,
|
&self,
|
||||||
filter: &VideoTaskQueryFilter,
|
filter: &VideoTaskQueryFilter,
|
||||||
|
|||||||
@@ -1,7 +1,7 @@
|
|||||||
mod types;
|
mod types;
|
||||||
|
|
||||||
pub use types::{
|
pub use types::{
|
||||||
ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
|
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogReadRepository,
|
||||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyPage,
|
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||||
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
|
StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
|
||||||
};
|
};
|
||||||
|
|||||||
@@ -450,6 +450,13 @@ impl StoredProviderCatalogKey {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq, Default)]
|
||||||
|
pub enum ProviderCatalogKeyListOrder {
|
||||||
|
#[default]
|
||||||
|
Name,
|
||||||
|
CreatedAt,
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, PartialEq, Eq, Default)]
|
#[derive(Debug, Clone, PartialEq, Eq, Default)]
|
||||||
pub struct ProviderCatalogKeyListQuery {
|
pub struct ProviderCatalogKeyListQuery {
|
||||||
pub provider_id: String,
|
pub provider_id: String,
|
||||||
@@ -457,6 +464,7 @@ pub struct ProviderCatalogKeyListQuery {
|
|||||||
pub is_active: Option<bool>,
|
pub is_active: Option<bool>,
|
||||||
pub offset: usize,
|
pub offset: usize,
|
||||||
pub limit: usize,
|
pub limit: usize,
|
||||||
|
pub order: ProviderCatalogKeyListOrder,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||||
@@ -529,6 +537,11 @@ pub trait ProviderCatalogReadRepository: Send + Sync {
|
|||||||
provider_ids: &[String],
|
provider_ids: &[String],
|
||||||
) -> Result<Vec<StoredProviderCatalogKey>, crate::DataLayerError>;
|
) -> Result<Vec<StoredProviderCatalogKey>, crate::DataLayerError>;
|
||||||
|
|
||||||
|
async fn list_key_summaries_by_provider_ids(
|
||||||
|
&self,
|
||||||
|
provider_ids: &[String],
|
||||||
|
) -> Result<Vec<StoredProviderCatalogKey>, crate::DataLayerError>;
|
||||||
|
|
||||||
async fn list_keys_page(
|
async fn list_keys_page(
|
||||||
&self,
|
&self,
|
||||||
query: &ProviderCatalogKeyListQuery,
|
query: &ProviderCatalogKeyListQuery,
|
||||||
|
|||||||
@@ -394,6 +394,13 @@ pub trait VideoTaskReadRepository: Send + Sync {
|
|||||||
limit: usize,
|
limit: usize,
|
||||||
) -> Result<Vec<StoredVideoTask>, crate::DataLayerError>;
|
) -> Result<Vec<StoredVideoTask>, crate::DataLayerError>;
|
||||||
|
|
||||||
|
async fn list_page_summary(
|
||||||
|
&self,
|
||||||
|
filter: &VideoTaskQueryFilter,
|
||||||
|
offset: usize,
|
||||||
|
limit: usize,
|
||||||
|
) -> Result<Vec<StoredVideoTask>, crate::DataLayerError>;
|
||||||
|
|
||||||
async fn count(&self, filter: &VideoTaskQueryFilter) -> Result<u64, crate::DataLayerError>;
|
async fn count(&self, filter: &VideoTaskQueryFilter) -> Result<u64, crate::DataLayerError>;
|
||||||
|
|
||||||
async fn count_by_status(
|
async fn count_by_status(
|
||||||
|
|||||||
@@ -63,6 +63,7 @@ use crate::repository::wallet::{
|
|||||||
SqlxWalletRepository, WalletReadRepository, WalletWriteRepository,
|
SqlxWalletRepository, WalletReadRepository, WalletWriteRepository,
|
||||||
};
|
};
|
||||||
use crate::DataLayerError;
|
use crate::DataLayerError;
|
||||||
|
use futures_util::TryStreamExt;
|
||||||
use sqlx::Row;
|
use sqlx::Row;
|
||||||
|
|
||||||
const FIND_SYSTEM_CONFIG_VALUE_SQL: &str = r#"
|
const FIND_SYSTEM_CONFIG_VALUE_SQL: &str = r#"
|
||||||
@@ -332,23 +333,20 @@ impl PostgresBackend {
|
|||||||
pub async fn list_system_config_entries(
|
pub async fn list_system_config_entries(
|
||||||
&self,
|
&self,
|
||||||
) -> Result<Vec<StoredSystemConfigEntry>, DataLayerError> {
|
) -> Result<Vec<StoredSystemConfigEntry>, DataLayerError> {
|
||||||
let rows = sqlx::query(LIST_SYSTEM_CONFIG_ENTRIES_SQL)
|
let mut rows = sqlx::query(LIST_SYSTEM_CONFIG_ENTRIES_SQL).fetch(&self.pool);
|
||||||
.fetch_all(&self.pool)
|
let mut entries = Vec::new();
|
||||||
.await
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||||
.map_postgres_err()?;
|
entries.push(StoredSystemConfigEntry {
|
||||||
rows.into_iter()
|
key: row.try_get("key").map_postgres_err()?,
|
||||||
.map(|row| {
|
value: row.try_get("value").map_postgres_err()?,
|
||||||
Ok(StoredSystemConfigEntry {
|
description: row.try_get("description").map_postgres_err()?,
|
||||||
key: row.try_get("key").map_postgres_err()?,
|
updated_at_unix_secs: row
|
||||||
value: row.try_get("value").map_postgres_err()?,
|
.try_get::<Option<i64>, _>("updated_at_unix_secs")
|
||||||
description: row.try_get("description").map_postgres_err()?,
|
.map_postgres_err()?
|
||||||
updated_at_unix_secs: row
|
.map(|value| value.max(0) as u64),
|
||||||
.try_get::<Option<i64>, _>("updated_at_unix_secs")
|
});
|
||||||
.map_postgres_err()?
|
}
|
||||||
.map(|value| value.max(0) as u64),
|
Ok(entries)
|
||||||
})
|
|
||||||
})
|
|
||||||
.collect::<Result<Vec<_>, DataLayerError>>()
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn upsert_system_config_entry(
|
pub async fn upsert_system_config_entry(
|
||||||
|
|||||||
@@ -1,7 +1,7 @@
|
|||||||
use crate::error::SqlxResultExt;
|
use crate::error::SqlxResultExt;
|
||||||
use crate::postgres::{DatabaseRecordId, PostgresTransactionOptions, PostgresTransactionRunner};
|
use crate::postgres::{DatabaseRecordId, PostgresTransactionOptions, PostgresTransactionRunner};
|
||||||
use crate::DataLayerError;
|
use crate::DataLayerError;
|
||||||
use futures_util::FutureExt;
|
use futures_util::{FutureExt, TryStreamExt};
|
||||||
use sqlx::query_scalar;
|
use sqlx::query_scalar;
|
||||||
|
|
||||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
@@ -105,13 +105,15 @@ impl PostgresLeaseRunner {
|
|||||||
self.transaction_runner
|
self.transaction_runner
|
||||||
.run(tx_options, |tx| {
|
.run(tx_options, |tx| {
|
||||||
async move {
|
async move {
|
||||||
let rows = query_scalar::<_, String>(&sql)
|
let mut rows = query_scalar::<_, String>(&sql)
|
||||||
.bind(owner)
|
.bind(owner)
|
||||||
.bind(lease_ms)
|
.bind(lease_ms)
|
||||||
.fetch_all(&mut **tx)
|
.fetch(&mut **tx);
|
||||||
.await
|
let mut ids = Vec::new();
|
||||||
.map_postgres_err()?;
|
while let Some(id) = rows.try_next().await.map_postgres_err()? {
|
||||||
Ok(rows.into_iter().map(DatabaseRecordId).collect())
|
ids.push(DatabaseRecordId(id));
|
||||||
|
}
|
||||||
|
Ok(ids)
|
||||||
}
|
}
|
||||||
.boxed()
|
.boxed()
|
||||||
})
|
})
|
||||||
@@ -140,13 +142,15 @@ impl PostgresLeaseRunner {
|
|||||||
self.transaction_runner
|
self.transaction_runner
|
||||||
.run(tx_options, |tx| {
|
.run(tx_options, |tx| {
|
||||||
async move {
|
async move {
|
||||||
let rows = query_scalar::<_, String>(&sql)
|
let mut rows = query_scalar::<_, String>(&sql)
|
||||||
.bind(ids)
|
.bind(ids)
|
||||||
.bind(owner)
|
.bind(owner)
|
||||||
.fetch_all(&mut **tx)
|
.fetch(&mut **tx);
|
||||||
.await
|
let mut released = Vec::new();
|
||||||
.map_postgres_err()?;
|
while let Some(id) = rows.try_next().await.map_postgres_err()? {
|
||||||
Ok(rows.into_iter().map(DatabaseRecordId).collect())
|
released.push(DatabaseRecordId(id));
|
||||||
|
}
|
||||||
|
Ok(released)
|
||||||
}
|
}
|
||||||
.boxed()
|
.boxed()
|
||||||
})
|
})
|
||||||
@@ -184,14 +188,16 @@ impl PostgresLeaseRunner {
|
|||||||
self.transaction_runner
|
self.transaction_runner
|
||||||
.run(tx_options, |tx| {
|
.run(tx_options, |tx| {
|
||||||
async move {
|
async move {
|
||||||
let rows = query_scalar::<_, String>(&sql)
|
let mut rows = query_scalar::<_, String>(&sql)
|
||||||
.bind(ids)
|
.bind(ids)
|
||||||
.bind(owner)
|
.bind(owner)
|
||||||
.bind(lease_ms)
|
.bind(lease_ms)
|
||||||
.fetch_all(&mut **tx)
|
.fetch(&mut **tx);
|
||||||
.await
|
let mut renewed = Vec::new();
|
||||||
.map_postgres_err()?;
|
while let Some(id) = rows.try_next().await.map_postgres_err()? {
|
||||||
Ok(rows.into_iter().map(DatabaseRecordId).collect())
|
renewed.push(DatabaseRecordId(id));
|
||||||
|
}
|
||||||
|
Ok(renewed)
|
||||||
}
|
}
|
||||||
.boxed()
|
.boxed()
|
||||||
})
|
})
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use chrono::{TimeZone, Utc};
|
use chrono::{TimeZone, Utc};
|
||||||
|
use futures_util::TryStreamExt;
|
||||||
use sqlx::{postgres::PgRow, PgPool, Row};
|
use sqlx::{postgres::PgRow, PgPool, Row};
|
||||||
|
|
||||||
use super::types::{
|
use super::types::{
|
||||||
@@ -220,18 +221,16 @@ impl AnnouncementReadRepository for SqlxAnnouncementReadRepository {
|
|||||||
.map_postgres_err()?
|
.map_postgres_err()?
|
||||||
.max(0) as u64;
|
.max(0) as u64;
|
||||||
|
|
||||||
let rows = sqlx::query(LIST_ANNOUNCEMENTS_SQL)
|
let mut rows = sqlx::query(LIST_ANNOUNCEMENTS_SQL)
|
||||||
.bind(query.active_only)
|
.bind(query.active_only)
|
||||||
.bind(now_unix_secs as f64)
|
.bind(now_unix_secs as f64)
|
||||||
.bind(query.offset as i64)
|
.bind(query.offset as i64)
|
||||||
.bind(query.limit as i64)
|
.bind(query.limit as i64)
|
||||||
.fetch_all(&self.pool)
|
.fetch(&self.pool);
|
||||||
.await
|
let mut items = Vec::new();
|
||||||
.map_postgres_err()?;
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||||
let items = rows
|
items.push(map_announcement_row(&row)?);
|
||||||
.iter()
|
}
|
||||||
.map(map_announcement_row)
|
|
||||||
.collect::<Result<Vec<_>, _>>()?;
|
|
||||||
|
|
||||||
Ok(StoredAnnouncementPage { items, total })
|
Ok(StoredAnnouncementPage { items, total })
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use sqlx::{PgPool, Row};
|
use futures_util::{stream::TryStream, TryStreamExt};
|
||||||
|
use sqlx::{postgres::PgRow, PgPool, Row};
|
||||||
|
|
||||||
use super::types::{
|
use super::types::{
|
||||||
AuthApiKeyExportSummary, AuthApiKeyLookupKey, AuthApiKeyReadRepository,
|
AuthApiKeyExportSummary, AuthApiKeyLookupKey, AuthApiKeyReadRepository,
|
||||||
@@ -654,6 +655,20 @@ impl SqlxAuthApiKeySnapshotReadRepository {
|
|||||||
&self.pool
|
&self.pool
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn collect_query_rows<T, S>(
|
||||||
|
mut rows: S,
|
||||||
|
map_row: fn(&PgRow) -> Result<T, DataLayerError>,
|
||||||
|
) -> Result<Vec<T>, DataLayerError>
|
||||||
|
where
|
||||||
|
S: TryStream<Ok = PgRow, Error = sqlx::Error> + Unpin,
|
||||||
|
{
|
||||||
|
let mut items = Vec::new();
|
||||||
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||||
|
items.push(map_row(&row)?);
|
||||||
|
}
|
||||||
|
Ok(items)
|
||||||
|
}
|
||||||
|
|
||||||
pub async fn find_api_key_snapshot(
|
pub async fn find_api_key_snapshot(
|
||||||
&self,
|
&self,
|
||||||
key: AuthApiKeyLookupKey<'_>,
|
key: AuthApiKeyLookupKey<'_>,
|
||||||
@@ -691,12 +706,13 @@ impl SqlxAuthApiKeySnapshotReadRepository {
|
|||||||
return Ok(Vec::new());
|
return Ok(Vec::new());
|
||||||
}
|
}
|
||||||
|
|
||||||
let rows = sqlx::query(LIST_BY_API_KEY_IDS_SQL)
|
Self::collect_query_rows(
|
||||||
.bind(api_key_ids)
|
sqlx::query(LIST_BY_API_KEY_IDS_SQL)
|
||||||
.fetch_all(&self.pool)
|
.bind(api_key_ids)
|
||||||
.await
|
.fetch(&self.pool),
|
||||||
.map_postgres_err()?;
|
map_auth_api_key_snapshot_row,
|
||||||
rows.iter().map(map_auth_api_key_snapshot_row).collect()
|
)
|
||||||
|
.await
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn list_export_api_keys_by_user_ids(
|
pub async fn list_export_api_keys_by_user_ids(
|
||||||
@@ -707,12 +723,13 @@ impl SqlxAuthApiKeySnapshotReadRepository {
|
|||||||
return Ok(Vec::new());
|
return Ok(Vec::new());
|
||||||
}
|
}
|
||||||
|
|
||||||
let rows = sqlx::query(LIST_EXPORT_BY_USER_IDS_SQL)
|
Self::collect_query_rows(
|
||||||
.bind(user_ids)
|
sqlx::query(LIST_EXPORT_BY_USER_IDS_SQL)
|
||||||
.fetch_all(&self.pool)
|
.bind(user_ids)
|
||||||
.await
|
.fetch(&self.pool),
|
||||||
.map_postgres_err()?;
|
map_auth_api_key_export_row,
|
||||||
rows.iter().map(map_auth_api_key_export_row).collect()
|
)
|
||||||
|
.await
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn list_export_api_keys_by_ids(
|
pub async fn list_export_api_keys_by_ids(
|
||||||
@@ -723,12 +740,13 @@ impl SqlxAuthApiKeySnapshotReadRepository {
|
|||||||
return Ok(Vec::new());
|
return Ok(Vec::new());
|
||||||
}
|
}
|
||||||
|
|
||||||
let rows = sqlx::query(LIST_EXPORT_BY_API_KEY_IDS_SQL)
|
Self::collect_query_rows(
|
||||||
.bind(api_key_ids)
|
sqlx::query(LIST_EXPORT_BY_API_KEY_IDS_SQL)
|
||||||
.fetch_all(&self.pool)
|
.bind(api_key_ids)
|
||||||
.await
|
.fetch(&self.pool),
|
||||||
.map_postgres_err()?;
|
map_auth_api_key_export_row,
|
||||||
rows.iter().map(map_auth_api_key_export_row).collect()
|
)
|
||||||
|
.await
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn summarize_export_api_keys_by_user_ids(
|
pub async fn summarize_export_api_keys_by_user_ids(
|
||||||
@@ -770,11 +788,11 @@ impl SqlxAuthApiKeySnapshotReadRepository {
|
|||||||
pub async fn list_export_standalone_api_keys(
|
pub async fn list_export_standalone_api_keys(
|
||||||
&self,
|
&self,
|
||||||
) -> Result<Vec<StoredAuthApiKeyExportRecord>, DataLayerError> {
|
) -> Result<Vec<StoredAuthApiKeyExportRecord>, DataLayerError> {
|
||||||
let rows = sqlx::query(LIST_EXPORT_STANDALONE_SQL)
|
Self::collect_query_rows(
|
||||||
.fetch_all(&self.pool)
|
sqlx::query(LIST_EXPORT_STANDALONE_SQL).fetch(&self.pool),
|
||||||
.await
|
map_auth_api_key_export_row,
|
||||||
.map_postgres_err()?;
|
)
|
||||||
rows.iter().map(map_auth_api_key_export_row).collect()
|
.await
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn list_export_standalone_api_keys_page(
|
pub async fn list_export_standalone_api_keys_page(
|
||||||
@@ -785,14 +803,15 @@ impl SqlxAuthApiKeySnapshotReadRepository {
|
|||||||
.map_err(|_| DataLayerError::InvalidInput("limit is too large".to_string()))?;
|
.map_err(|_| DataLayerError::InvalidInput("limit is too large".to_string()))?;
|
||||||
let skip = i64::try_from(query.skip)
|
let skip = i64::try_from(query.skip)
|
||||||
.map_err(|_| DataLayerError::InvalidInput("skip is too large".to_string()))?;
|
.map_err(|_| DataLayerError::InvalidInput("skip is too large".to_string()))?;
|
||||||
let rows = sqlx::query(LIST_EXPORT_STANDALONE_PAGE_SQL)
|
Self::collect_query_rows(
|
||||||
.bind(query.is_active)
|
sqlx::query(LIST_EXPORT_STANDALONE_PAGE_SQL)
|
||||||
.bind(skip)
|
.bind(query.is_active)
|
||||||
.bind(limit)
|
.bind(skip)
|
||||||
.fetch_all(&self.pool)
|
.bind(limit)
|
||||||
.await
|
.fetch(&self.pool),
|
||||||
.map_postgres_err()?;
|
map_auth_api_key_export_row,
|
||||||
rows.iter().map(map_auth_api_key_export_row).collect()
|
)
|
||||||
|
.await
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn count_export_standalone_api_keys(
|
pub async fn count_export_standalone_api_keys(
|
||||||
|
|||||||
@@ -1,4 +1,5 @@
|
|||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
|
use futures_util::{stream::TryStream, TryStreamExt};
|
||||||
use sqlx::{postgres::PgRow, PgPool, Row};
|
use sqlx::{postgres::PgRow, PgPool, Row};
|
||||||
|
|
||||||
use super::types::{
|
use super::types::{
|
||||||
@@ -145,16 +146,30 @@ impl SqlxAuthModuleRepository {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn collect_query_rows<T, S>(
|
||||||
|
mut rows: S,
|
||||||
|
map_row: fn(&PgRow) -> Result<T, DataLayerError>,
|
||||||
|
) -> Result<Vec<T>, DataLayerError>
|
||||||
|
where
|
||||||
|
S: TryStream<Ok = PgRow, Error = sqlx::Error> + Unpin,
|
||||||
|
{
|
||||||
|
let mut items = Vec::new();
|
||||||
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||||
|
items.push(map_row(&row)?);
|
||||||
|
}
|
||||||
|
Ok(items)
|
||||||
|
}
|
||||||
|
|
||||||
#[async_trait]
|
#[async_trait]
|
||||||
impl AuthModuleReadRepository for SqlxAuthModuleReadRepository {
|
impl AuthModuleReadRepository for SqlxAuthModuleReadRepository {
|
||||||
async fn list_enabled_oauth_providers(
|
async fn list_enabled_oauth_providers(
|
||||||
&self,
|
&self,
|
||||||
) -> Result<Vec<StoredOAuthProviderModuleConfig>, DataLayerError> {
|
) -> Result<Vec<StoredOAuthProviderModuleConfig>, DataLayerError> {
|
||||||
let rows = sqlx::query(LIST_ENABLED_OAUTH_PROVIDERS_SQL)
|
collect_query_rows(
|
||||||
.fetch_all(&self.pool)
|
sqlx::query(LIST_ENABLED_OAUTH_PROVIDERS_SQL).fetch(&self.pool),
|
||||||
.await
|
map_oauth_row,
|
||||||
.map_postgres_err()?;
|
)
|
||||||
rows.iter().map(map_oauth_row).collect()
|
.await
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn get_ldap_config(&self) -> Result<Option<StoredLdapModuleConfig>, DataLayerError> {
|
async fn get_ldap_config(&self) -> Result<Option<StoredLdapModuleConfig>, DataLayerError> {
|
||||||
@@ -171,11 +186,11 @@ impl AuthModuleReadRepository for SqlxAuthModuleRepository {
|
|||||||
async fn list_enabled_oauth_providers(
|
async fn list_enabled_oauth_providers(
|
||||||
&self,
|
&self,
|
||||||
) -> Result<Vec<StoredOAuthProviderModuleConfig>, DataLayerError> {
|
) -> Result<Vec<StoredOAuthProviderModuleConfig>, DataLayerError> {
|
||||||
let rows = sqlx::query(LIST_ENABLED_OAUTH_PROVIDERS_SQL)
|
collect_query_rows(
|
||||||
.fetch_all(&self.pool)
|
sqlx::query(LIST_ENABLED_OAUTH_PROVIDERS_SQL).fetch(&self.pool),
|
||||||
.await
|
map_oauth_row,
|
||||||
.map_postgres_err()?;
|
)
|
||||||
rows.iter().map(map_oauth_row).collect()
|
.await
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn get_ldap_config(&self) -> Result<Option<StoredLdapModuleConfig>, DataLayerError> {
|
async fn get_ldap_config(&self) -> Result<Option<StoredLdapModuleConfig>, DataLayerError> {
|
||||||
|
|||||||
@@ -1,4 +1,5 @@
|
|||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
|
use futures_util::{stream::TryStream, TryStreamExt};
|
||||||
use sqlx::{PgPool, Row};
|
use sqlx::{PgPool, Row};
|
||||||
|
|
||||||
use super::{
|
use super::{
|
||||||
@@ -165,16 +166,31 @@ impl SqlxMinimalCandidateSelectionReadRepository {
|
|||||||
&self.pool
|
&self.pool
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn collect_query_rows<T, S>(
|
||||||
|
mut rows: S,
|
||||||
|
map_row: fn(&sqlx::postgres::PgRow) -> Result<T, DataLayerError>,
|
||||||
|
) -> Result<Vec<T>, DataLayerError>
|
||||||
|
where
|
||||||
|
S: TryStream<Ok = sqlx::postgres::PgRow, Error = sqlx::Error> + Unpin,
|
||||||
|
{
|
||||||
|
let mut items = Vec::new();
|
||||||
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||||
|
items.push(map_row(&row)?);
|
||||||
|
}
|
||||||
|
Ok(items)
|
||||||
|
}
|
||||||
|
|
||||||
pub async fn list_for_exact_api_format(
|
pub async fn list_for_exact_api_format(
|
||||||
&self,
|
&self,
|
||||||
api_format: &str,
|
api_format: &str,
|
||||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||||
let rows = sqlx::query(LIST_FOR_EXACT_API_FORMAT_SQL)
|
Self::collect_query_rows(
|
||||||
.bind(api_format)
|
sqlx::query(LIST_FOR_EXACT_API_FORMAT_SQL)
|
||||||
.fetch_all(&self.pool)
|
.bind(api_format)
|
||||||
.await
|
.fetch(&self.pool),
|
||||||
.map_postgres_err()?;
|
map_candidate_selection_row,
|
||||||
rows.iter().map(map_candidate_selection_row).collect()
|
)
|
||||||
|
.await
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn list_for_exact_api_format_and_global_model(
|
pub async fn list_for_exact_api_format_and_global_model(
|
||||||
@@ -182,13 +198,14 @@ impl SqlxMinimalCandidateSelectionReadRepository {
|
|||||||
api_format: &str,
|
api_format: &str,
|
||||||
global_model_name: &str,
|
global_model_name: &str,
|
||||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||||
let rows = sqlx::query(LIST_FOR_EXACT_API_FORMAT_AND_GLOBAL_MODEL_SQL)
|
Self::collect_query_rows(
|
||||||
.bind(api_format)
|
sqlx::query(LIST_FOR_EXACT_API_FORMAT_AND_GLOBAL_MODEL_SQL)
|
||||||
.bind(global_model_name)
|
.bind(api_format)
|
||||||
.fetch_all(&self.pool)
|
.bind(global_model_name)
|
||||||
.await
|
.fetch(&self.pool),
|
||||||
.map_postgres_err()?;
|
map_candidate_selection_row,
|
||||||
rows.iter().map(map_candidate_selection_row).collect()
|
)
|
||||||
|
.await
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use futures_util::future::BoxFuture;
|
use futures_util::{future::BoxFuture, stream::TryStream, TryStreamExt};
|
||||||
use sqlx::{PgPool, Row};
|
use sqlx::{postgres::PgRow, PgPool, Row};
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
use super::{
|
use super::{
|
||||||
@@ -306,12 +306,13 @@ impl SqlxRequestCandidateReadRepository {
|
|||||||
&self,
|
&self,
|
||||||
request_id: &str,
|
request_id: &str,
|
||||||
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
|
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
|
||||||
let rows = sqlx::query(LIST_BY_REQUEST_ID_SQL)
|
collect_query_rows(
|
||||||
.bind(request_id)
|
sqlx::query(LIST_BY_REQUEST_ID_SQL)
|
||||||
.fetch_all(&self.pool)
|
.bind(request_id)
|
||||||
.await
|
.fetch(&self.pool),
|
||||||
.map_postgres_err()?;
|
map_request_candidate_row,
|
||||||
rows.iter().map(map_request_candidate_row).collect()
|
)
|
||||||
|
.await
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn list_recent(
|
pub async fn list_recent(
|
||||||
@@ -322,16 +323,17 @@ impl SqlxRequestCandidateReadRepository {
|
|||||||
return Ok(Vec::new());
|
return Ok(Vec::new());
|
||||||
}
|
}
|
||||||
|
|
||||||
let rows = sqlx::query(LIST_RECENT_SQL)
|
collect_query_rows(
|
||||||
.bind(i64::try_from(limit).map_err(|_| {
|
sqlx::query(LIST_RECENT_SQL)
|
||||||
DataLayerError::UnexpectedValue(format!(
|
.bind(i64::try_from(limit).map_err(|_| {
|
||||||
"invalid recent request candidate limit: {limit}"
|
DataLayerError::UnexpectedValue(format!(
|
||||||
))
|
"invalid recent request candidate limit: {limit}"
|
||||||
})?)
|
))
|
||||||
.fetch_all(&self.pool)
|
})?)
|
||||||
.await
|
.fetch(&self.pool),
|
||||||
.map_postgres_err()?;
|
map_request_candidate_row,
|
||||||
rows.iter().map(map_request_candidate_row).collect()
|
)
|
||||||
|
.await
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn list_by_provider_id(
|
pub async fn list_by_provider_id(
|
||||||
@@ -349,13 +351,14 @@ impl SqlxRequestCandidateReadRepository {
|
|||||||
))
|
))
|
||||||
})?;
|
})?;
|
||||||
|
|
||||||
let rows = sqlx::query(LIST_BY_PROVIDER_ID_SQL)
|
collect_query_rows(
|
||||||
.bind(provider_id)
|
sqlx::query(LIST_BY_PROVIDER_ID_SQL)
|
||||||
.bind(limit_value)
|
.bind(provider_id)
|
||||||
.fetch_all(&self.pool)
|
.bind(limit_value)
|
||||||
.await
|
.fetch(&self.pool),
|
||||||
.map_postgres_err()?;
|
map_request_candidate_row,
|
||||||
rows.iter().map(map_request_candidate_row).collect()
|
)
|
||||||
|
.await
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn list_finalized_by_endpoint_ids_since(
|
pub async fn list_finalized_by_endpoint_ids_since(
|
||||||
@@ -368,18 +371,19 @@ impl SqlxRequestCandidateReadRepository {
|
|||||||
return Ok(Vec::new());
|
return Ok(Vec::new());
|
||||||
}
|
}
|
||||||
|
|
||||||
let rows = sqlx::query(LIST_FINALIZED_BY_ENDPOINT_IDS_SINCE_SQL)
|
collect_query_rows(
|
||||||
.bind(endpoint_ids)
|
sqlx::query(LIST_FINALIZED_BY_ENDPOINT_IDS_SINCE_SQL)
|
||||||
.bind(since_unix_secs as f64)
|
.bind(endpoint_ids)
|
||||||
.bind(i64::try_from(limit).map_err(|_| {
|
.bind(since_unix_secs as f64)
|
||||||
DataLayerError::UnexpectedValue(format!(
|
.bind(i64::try_from(limit).map_err(|_| {
|
||||||
"invalid finalized request candidate limit: {limit}"
|
DataLayerError::UnexpectedValue(format!(
|
||||||
))
|
"invalid finalized request candidate limit: {limit}"
|
||||||
})?)
|
))
|
||||||
.fetch_all(&self.pool)
|
})?)
|
||||||
.await
|
.fetch(&self.pool),
|
||||||
.map_postgres_err()?;
|
map_request_candidate_row,
|
||||||
rows.iter().map(map_request_candidate_row).collect()
|
)
|
||||||
|
.await
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn count_finalized_statuses_by_endpoint_ids_since(
|
pub async fn count_finalized_statuses_by_endpoint_ids_since(
|
||||||
@@ -391,29 +395,29 @@ impl SqlxRequestCandidateReadRepository {
|
|||||||
return Ok(Vec::new());
|
return Ok(Vec::new());
|
||||||
}
|
}
|
||||||
|
|
||||||
let rows = sqlx::query(COUNT_FINALIZED_STATUSES_BY_ENDPOINT_IDS_SINCE_SQL)
|
let mut rows = sqlx::query(COUNT_FINALIZED_STATUSES_BY_ENDPOINT_IDS_SINCE_SQL)
|
||||||
.bind(endpoint_ids)
|
.bind(endpoint_ids)
|
||||||
.bind(since_unix_secs as f64)
|
.bind(since_unix_secs as f64)
|
||||||
.fetch_all(&self.pool)
|
.fetch(&self.pool);
|
||||||
.await
|
let mut counts = Vec::new();
|
||||||
.map_postgres_err()?;
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||||
|
let entry = {
|
||||||
rows.iter()
|
|
||||||
.map(|row| {
|
|
||||||
let status = RequestCandidateStatus::from_database(
|
let status = RequestCandidateStatus::from_database(
|
||||||
row_get::<String>(row, "status")?.as_str(),
|
row_get::<String>(&row, "status")?.as_str(),
|
||||||
)?;
|
)?;
|
||||||
Ok(PublicHealthStatusCount {
|
PublicHealthStatusCount {
|
||||||
endpoint_id: row_get(row, "endpoint_id")?,
|
endpoint_id: row_get(&row, "endpoint_id")?,
|
||||||
status,
|
status,
|
||||||
count: u64::try_from(row_get::<i64>(row, "count")?).map_err(|_| {
|
count: u64::try_from(row_get::<i64>(&row, "count")?).map_err(|_| {
|
||||||
DataLayerError::UnexpectedValue(
|
DataLayerError::UnexpectedValue(
|
||||||
"public health status count out of range".to_string(),
|
"public health status count out of range".to_string(),
|
||||||
)
|
)
|
||||||
})?,
|
})?,
|
||||||
})
|
}
|
||||||
})
|
};
|
||||||
.collect()
|
counts.push(entry);
|
||||||
|
}
|
||||||
|
Ok(counts)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn aggregate_finalized_timeline_by_endpoint_ids_since(
|
pub async fn aggregate_finalized_timeline_by_endpoint_ids_since(
|
||||||
@@ -434,18 +438,16 @@ impl SqlxRequestCandidateReadRepository {
|
|||||||
(span_seconds as f64) / (segments as f64)
|
(span_seconds as f64) / (segments as f64)
|
||||||
};
|
};
|
||||||
|
|
||||||
let rows = sqlx::query(AGGREGATE_FINALIZED_TIMELINE_BY_ENDPOINT_IDS_SINCE_SQL)
|
let mut rows = sqlx::query(AGGREGATE_FINALIZED_TIMELINE_BY_ENDPOINT_IDS_SINCE_SQL)
|
||||||
.bind(endpoint_ids)
|
.bind(endpoint_ids)
|
||||||
.bind(since_unix_secs as f64)
|
.bind(since_unix_secs as f64)
|
||||||
.bind(until_unix_secs as f64)
|
.bind(until_unix_secs as f64)
|
||||||
.bind(segment_seconds)
|
.bind(segment_seconds)
|
||||||
.fetch_all(&self.pool)
|
.fetch(&self.pool);
|
||||||
.await
|
let mut buckets = Vec::new();
|
||||||
.map_postgres_err()?;
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||||
|
let bucket = {
|
||||||
rows.iter()
|
let raw_segment_idx = row_get::<i64>(&row, "segment_idx")?;
|
||||||
.map(|row| {
|
|
||||||
let raw_segment_idx = row_get::<i64>(row, "segment_idx")?;
|
|
||||||
let segment_idx = if raw_segment_idx < 0 {
|
let segment_idx = if raw_segment_idx < 0 {
|
||||||
0
|
0
|
||||||
} else {
|
} else {
|
||||||
@@ -457,31 +459,31 @@ impl SqlxRequestCandidateReadRepository {
|
|||||||
}
|
}
|
||||||
.min(segments.saturating_sub(1));
|
.min(segments.saturating_sub(1));
|
||||||
|
|
||||||
Ok(PublicHealthTimelineBucket {
|
PublicHealthTimelineBucket {
|
||||||
endpoint_id: row_get(row, "endpoint_id")?,
|
endpoint_id: row_get(&row, "endpoint_id")?,
|
||||||
segment_idx,
|
segment_idx,
|
||||||
total_count: u64::try_from(row_get::<i64>(row, "total_count")?).map_err(
|
total_count: u64::try_from(row_get::<i64>(&row, "total_count")?).map_err(
|
||||||
|_| {
|
|_| {
|
||||||
DataLayerError::UnexpectedValue(
|
DataLayerError::UnexpectedValue(
|
||||||
"public health total_count out of range".to_string(),
|
"public health total_count out of range".to_string(),
|
||||||
)
|
)
|
||||||
},
|
},
|
||||||
)?,
|
)?,
|
||||||
success_count: u64::try_from(row_get::<i64>(row, "success_count")?).map_err(
|
success_count: u64::try_from(row_get::<i64>(&row, "success_count")?).map_err(
|
||||||
|_| {
|
|_| {
|
||||||
DataLayerError::UnexpectedValue(
|
DataLayerError::UnexpectedValue(
|
||||||
"public health success_count out of range".to_string(),
|
"public health success_count out of range".to_string(),
|
||||||
)
|
)
|
||||||
},
|
},
|
||||||
)?,
|
)?,
|
||||||
failed_count: u64::try_from(row_get::<i64>(row, "failed_count")?).map_err(
|
failed_count: u64::try_from(row_get::<i64>(&row, "failed_count")?).map_err(
|
||||||
|_| {
|
|_| {
|
||||||
DataLayerError::UnexpectedValue(
|
DataLayerError::UnexpectedValue(
|
||||||
"public health failed_count out of range".to_string(),
|
"public health failed_count out of range".to_string(),
|
||||||
)
|
)
|
||||||
},
|
},
|
||||||
)?,
|
)?,
|
||||||
min_created_at_unix_ms: row_get::<Option<i64>>(row, "min_created_at_unix_ms")?
|
min_created_at_unix_ms: row_get::<Option<i64>>(&row, "min_created_at_unix_ms")?
|
||||||
.map(|value| {
|
.map(|value| {
|
||||||
u64::try_from(value).map_err(|_| {
|
u64::try_from(value).map_err(|_| {
|
||||||
DataLayerError::UnexpectedValue(format!(
|
DataLayerError::UnexpectedValue(format!(
|
||||||
@@ -490,7 +492,7 @@ impl SqlxRequestCandidateReadRepository {
|
|||||||
})
|
})
|
||||||
})
|
})
|
||||||
.transpose()?,
|
.transpose()?,
|
||||||
max_created_at_unix_ms: row_get::<Option<i64>>(row, "max_created_at_unix_ms")?
|
max_created_at_unix_ms: row_get::<Option<i64>>(&row, "max_created_at_unix_ms")?
|
||||||
.map(|value| {
|
.map(|value| {
|
||||||
u64::try_from(value).map_err(|_| {
|
u64::try_from(value).map_err(|_| {
|
||||||
DataLayerError::UnexpectedValue(format!(
|
DataLayerError::UnexpectedValue(format!(
|
||||||
@@ -499,9 +501,11 @@ impl SqlxRequestCandidateReadRepository {
|
|||||||
})
|
})
|
||||||
})
|
})
|
||||||
.transpose()?,
|
.transpose()?,
|
||||||
})
|
}
|
||||||
})
|
};
|
||||||
.collect()
|
buckets.push(bucket);
|
||||||
|
}
|
||||||
|
Ok(buckets)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn upsert(
|
pub async fn upsert(
|
||||||
@@ -651,9 +655,21 @@ impl RequestCandidateWriteRepository for SqlxRequestCandidateReadRepository {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn map_request_candidate_row(
|
async fn collect_query_rows<T, S>(
|
||||||
row: &sqlx::postgres::PgRow,
|
mut rows: S,
|
||||||
) -> Result<StoredRequestCandidate, DataLayerError> {
|
map_row: fn(&PgRow) -> Result<T, DataLayerError>,
|
||||||
|
) -> Result<Vec<T>, DataLayerError>
|
||||||
|
where
|
||||||
|
S: TryStream<Ok = PgRow, Error = sqlx::Error> + Unpin,
|
||||||
|
{
|
||||||
|
let mut items = Vec::new();
|
||||||
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||||
|
items.push(map_row(&row)?);
|
||||||
|
}
|
||||||
|
Ok(items)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn map_request_candidate_row(row: &PgRow) -> Result<StoredRequestCandidate, DataLayerError> {
|
||||||
let status = RequestCandidateStatus::from_database(row_get::<String>(row, "status")?.as_str())?;
|
let status = RequestCandidateStatus::from_database(row_get::<String>(row, "status")?.as_str())?;
|
||||||
StoredRequestCandidate::new(
|
StoredRequestCandidate::new(
|
||||||
row_get(row, "id")?,
|
row_get(row, "id")?,
|
||||||
@@ -683,7 +699,7 @@ fn map_request_candidate_row(
|
|||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
fn row_get<T>(row: &sqlx::postgres::PgRow, column: &str) -> Result<T, DataLayerError>
|
fn row_get<T>(row: &PgRow, column: &str) -> Result<T, DataLayerError>
|
||||||
where
|
where
|
||||||
for<'r> T: sqlx::Decode<'r, sqlx::Postgres> + sqlx::Type<sqlx::Postgres>,
|
for<'r> T: sqlx::Decode<'r, sqlx::Postgres> + sqlx::Type<sqlx::Postgres>,
|
||||||
{
|
{
|
||||||
|
|||||||
@@ -1,4 +1,5 @@
|
|||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
|
use futures_util::TryStreamExt;
|
||||||
use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row};
|
use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row};
|
||||||
|
|
||||||
use super::types::{
|
use super::types::{
|
||||||
@@ -91,16 +92,15 @@ WHERE file_name = $1
|
|||||||
.fetch_one(&self.pool)
|
.fetch_one(&self.pool)
|
||||||
.await
|
.await
|
||||||
.map_postgres_err()?;
|
.map_postgres_err()?;
|
||||||
let rows = build_list_rows_query(query)
|
let mut builder = build_list_rows_query(query);
|
||||||
.build()
|
let built_query = builder.build();
|
||||||
.fetch_all(&self.pool)
|
let mut rows = built_query.fetch(&self.pool);
|
||||||
.await
|
let mut items = Vec::new();
|
||||||
.map_postgres_err()?;
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||||
|
items.push(Self::map_row(&row)?);
|
||||||
|
}
|
||||||
Ok(StoredGeminiFileMappingListPage {
|
Ok(StoredGeminiFileMappingListPage {
|
||||||
items: rows
|
items,
|
||||||
.iter()
|
|
||||||
.map(Self::map_row)
|
|
||||||
.collect::<Result<Vec<_>, _>>()?,
|
|
||||||
total: usize::try_from(total).unwrap_or_default(),
|
total: usize::try_from(total).unwrap_or_default(),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
@@ -133,7 +133,7 @@ FROM gemini_file_mappings
|
|||||||
.map_postgres_err()?,
|
.map_postgres_err()?,
|
||||||
)
|
)
|
||||||
.unwrap_or_default();
|
.unwrap_or_default();
|
||||||
let by_mime_type_rows = sqlx::query(
|
let mut by_mime_type_rows = sqlx::query(
|
||||||
r#"
|
r#"
|
||||||
SELECT
|
SELECT
|
||||||
COALESCE(NULLIF(TRIM(mime_type), ''), 'unknown') AS mime_type,
|
COALESCE(NULLIF(TRIM(mime_type), ''), 'unknown') AS mime_type,
|
||||||
@@ -145,23 +145,20 @@ ORDER BY mime_type ASC
|
|||||||
"#,
|
"#,
|
||||||
)
|
)
|
||||||
.bind(now_unix_secs as f64)
|
.bind(now_unix_secs as f64)
|
||||||
.fetch_all(&self.pool)
|
.fetch(&self.pool);
|
||||||
.await
|
let mut by_mime_type = Vec::new();
|
||||||
.map_postgres_err()?;
|
while let Some(row) = by_mime_type_rows.try_next().await.map_postgres_err()? {
|
||||||
|
by_mime_type.push(GeminiFileMappingMimeTypeCount {
|
||||||
|
mime_type: row.try_get("mime_type").map_postgres_err()?,
|
||||||
|
count: usize::try_from(row.try_get::<i64, _>("count").map_postgres_err()?)
|
||||||
|
.unwrap_or_default(),
|
||||||
|
});
|
||||||
|
}
|
||||||
Ok(GeminiFileMappingStats {
|
Ok(GeminiFileMappingStats {
|
||||||
total_mappings,
|
total_mappings,
|
||||||
active_mappings,
|
active_mappings,
|
||||||
expired_mappings: total_mappings.saturating_sub(active_mappings),
|
expired_mappings: total_mappings.saturating_sub(active_mappings),
|
||||||
by_mime_type: by_mime_type_rows
|
by_mime_type,
|
||||||
.into_iter()
|
|
||||||
.map(|row| {
|
|
||||||
Ok(GeminiFileMappingMimeTypeCount {
|
|
||||||
mime_type: row.try_get("mime_type").map_postgres_err()?,
|
|
||||||
count: usize::try_from(row.try_get::<i64, _>("count").map_postgres_err()?)
|
|
||||||
.unwrap_or_default(),
|
|
||||||
})
|
|
||||||
})
|
|
||||||
.collect::<Result<Vec<_>, DataLayerError>>()?,
|
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,4 +1,5 @@
|
|||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
|
use futures_util::{stream::TryStream, TryStreamExt};
|
||||||
use serde_json::Value;
|
use serde_json::Value;
|
||||||
use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row};
|
use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row};
|
||||||
|
|
||||||
@@ -180,12 +181,8 @@ impl SqlxGlobalModelReadRepository {
|
|||||||
.push_bind(query.offset as i64)
|
.push_bind(query.offset as i64)
|
||||||
.push(" LIMIT ")
|
.push(" LIMIT ")
|
||||||
.push_bind(query.limit as i64);
|
.push_bind(query.limit as i64);
|
||||||
let rows = list_builder
|
let query = list_builder.build();
|
||||||
.build()
|
let items = collect_query_rows(query.fetch(&self.pool), map_row).await?;
|
||||||
.fetch_all(&self.pool)
|
|
||||||
.await
|
|
||||||
.map_postgres_err()?;
|
|
||||||
let items = rows.iter().map(map_row).collect::<Result<Vec<_>, _>>()?;
|
|
||||||
|
|
||||||
Ok(StoredPublicGlobalModelPage { items, total })
|
Ok(StoredPublicGlobalModelPage { items, total })
|
||||||
}
|
}
|
||||||
@@ -198,17 +195,13 @@ impl SqlxGlobalModelReadRepository {
|
|||||||
return Ok(Vec::new());
|
return Ok(Vec::new());
|
||||||
}
|
}
|
||||||
|
|
||||||
let rows = build_provider_id_list_query(
|
let mut builder = build_provider_id_list_query(
|
||||||
LIST_PROVIDER_MODEL_STATS_PREFIX,
|
LIST_PROVIDER_MODEL_STATS_PREFIX,
|
||||||
provider_ids,
|
provider_ids,
|
||||||
")\nGROUP BY provider_id\nORDER BY provider_id ASC",
|
")\nGROUP BY provider_id\nORDER BY provider_id ASC",
|
||||||
)
|
);
|
||||||
.build()
|
let query = builder.build();
|
||||||
.fetch_all(&self.pool)
|
collect_query_rows(query.fetch(&self.pool), map_provider_model_stats_row).await
|
||||||
.await
|
|
||||||
.map_postgres_err()?;
|
|
||||||
|
|
||||||
rows.iter().map(map_provider_model_stats_row).collect()
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn list_active_global_model_ids_by_provider_ids(
|
pub async fn list_active_global_model_ids_by_provider_ids(
|
||||||
@@ -219,19 +212,17 @@ impl SqlxGlobalModelReadRepository {
|
|||||||
return Ok(Vec::new());
|
return Ok(Vec::new());
|
||||||
}
|
}
|
||||||
|
|
||||||
let rows = build_provider_id_list_query(
|
let mut builder = build_provider_id_list_query(
|
||||||
LIST_ACTIVE_GLOBAL_MODEL_IDS_BY_PROVIDER_IDS_PREFIX,
|
LIST_ACTIVE_GLOBAL_MODEL_IDS_BY_PROVIDER_IDS_PREFIX,
|
||||||
provider_ids,
|
provider_ids,
|
||||||
")\nAND is_active = TRUE\nAND global_model_id IS NOT NULL\nORDER BY provider_id ASC, global_model_id ASC",
|
")\nAND is_active = TRUE\nAND global_model_id IS NOT NULL\nORDER BY provider_id ASC, global_model_id ASC",
|
||||||
|
);
|
||||||
|
let query = builder.build();
|
||||||
|
collect_query_rows(
|
||||||
|
query.fetch(&self.pool),
|
||||||
|
map_provider_active_global_model_row,
|
||||||
)
|
)
|
||||||
.build()
|
|
||||||
.fetch_all(&self.pool)
|
|
||||||
.await
|
.await
|
||||||
.map_postgres_err()?;
|
|
||||||
|
|
||||||
rows.iter()
|
|
||||||
.map(map_provider_active_global_model_row)
|
|
||||||
.collect()
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn list_admin_provider_models(
|
pub async fn list_admin_provider_models(
|
||||||
@@ -250,12 +241,8 @@ impl SqlxGlobalModelReadRepository {
|
|||||||
.push_bind(query.offset as i64)
|
.push_bind(query.offset as i64)
|
||||||
.push(" LIMIT ")
|
.push(" LIMIT ")
|
||||||
.push_bind(query.limit as i64);
|
.push_bind(query.limit as i64);
|
||||||
let rows = builder
|
let query = builder.build();
|
||||||
.build()
|
collect_query_rows(query.fetch(&self.pool), map_admin_provider_model_row).await
|
||||||
.fetch_all(&self.pool)
|
|
||||||
.await
|
|
||||||
.map_postgres_err()?;
|
|
||||||
rows.iter().map(map_admin_provider_model_row).collect()
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn list_admin_global_models(
|
pub async fn list_admin_global_models(
|
||||||
@@ -281,15 +268,8 @@ impl SqlxGlobalModelReadRepository {
|
|||||||
.push_bind(query.offset as i64)
|
.push_bind(query.offset as i64)
|
||||||
.push(" LIMIT ")
|
.push(" LIMIT ")
|
||||||
.push_bind(query.limit as i64);
|
.push_bind(query.limit as i64);
|
||||||
let rows = list_builder
|
let query = list_builder.build();
|
||||||
.build()
|
let items = collect_query_rows(query.fetch(&self.pool), map_admin_global_model_row).await?;
|
||||||
.fetch_all(&self.pool)
|
|
||||||
.await
|
|
||||||
.map_postgres_err()?;
|
|
||||||
let items = rows
|
|
||||||
.iter()
|
|
||||||
.map(map_admin_global_model_row)
|
|
||||||
.collect::<Result<Vec<_>, _>>()?;
|
|
||||||
Ok(StoredAdminGlobalModelPage { items, total })
|
Ok(StoredAdminGlobalModelPage { items, total })
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -343,8 +323,9 @@ LIMIT 1
|
|||||||
&self,
|
&self,
|
||||||
provider_id: &str,
|
provider_id: &str,
|
||||||
) -> Result<Vec<StoredAdminProviderModel>, DataLayerError> {
|
) -> Result<Vec<StoredAdminProviderModel>, DataLayerError> {
|
||||||
let rows = sqlx::query(
|
collect_query_rows(
|
||||||
r#"
|
sqlx::query(
|
||||||
|
r#"
|
||||||
SELECT
|
SELECT
|
||||||
m.id,
|
m.id,
|
||||||
m.provider_id,
|
m.provider_id,
|
||||||
@@ -375,13 +356,12 @@ WHERE m.provider_id = $1
|
|||||||
AND gm.is_active = TRUE
|
AND gm.is_active = TRUE
|
||||||
ORDER BY gm.name ASC, m.created_at DESC, m.id ASC
|
ORDER BY gm.name ASC, m.created_at DESC, m.id ASC
|
||||||
"#,
|
"#,
|
||||||
|
)
|
||||||
|
.bind(provider_id)
|
||||||
|
.fetch(&self.pool),
|
||||||
|
map_admin_provider_model_row,
|
||||||
)
|
)
|
||||||
.bind(provider_id)
|
|
||||||
.fetch_all(&self.pool)
|
|
||||||
.await
|
.await
|
||||||
.map_postgres_err()?;
|
|
||||||
|
|
||||||
rows.iter().map(map_admin_provider_model_row).collect()
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn get_admin_global_model_by_id(
|
pub async fn get_admin_global_model_by_id(
|
||||||
@@ -482,8 +462,9 @@ LIMIT 1
|
|||||||
&self,
|
&self,
|
||||||
global_model_id: &str,
|
global_model_id: &str,
|
||||||
) -> Result<Vec<StoredAdminProviderModel>, DataLayerError> {
|
) -> Result<Vec<StoredAdminProviderModel>, DataLayerError> {
|
||||||
let rows = sqlx::query(
|
collect_query_rows(
|
||||||
r#"
|
sqlx::query(
|
||||||
|
r#"
|
||||||
SELECT
|
SELECT
|
||||||
m.id,
|
m.id,
|
||||||
m.provider_id,
|
m.provider_id,
|
||||||
@@ -512,13 +493,12 @@ LEFT JOIN global_models gm ON gm.id = m.global_model_id
|
|||||||
WHERE m.global_model_id = $1
|
WHERE m.global_model_id = $1
|
||||||
ORDER BY m.created_at DESC, m.id ASC
|
ORDER BY m.created_at DESC, m.id ASC
|
||||||
"#,
|
"#,
|
||||||
|
)
|
||||||
|
.bind(global_model_id)
|
||||||
|
.fetch(&self.pool),
|
||||||
|
map_admin_provider_model_row,
|
||||||
)
|
)
|
||||||
.bind(global_model_id)
|
|
||||||
.fetch_all(&self.pool)
|
|
||||||
.await
|
.await
|
||||||
.map_postgres_err()?;
|
|
||||||
|
|
||||||
rows.iter().map(map_admin_provider_model_row).collect()
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn create_admin_provider_model(
|
pub async fn create_admin_provider_model(
|
||||||
@@ -801,12 +781,8 @@ LIMIT 1
|
|||||||
.push_bind(query.offset as i64)
|
.push_bind(query.offset as i64)
|
||||||
.push(" LIMIT ")
|
.push(" LIMIT ")
|
||||||
.push_bind(query.limit as i64);
|
.push_bind(query.limit as i64);
|
||||||
let rows = builder
|
let query = builder.build();
|
||||||
.build()
|
collect_query_rows(query.fetch(&self.pool), map_public_catalog_model_row).await
|
||||||
.fetch_all(&self.pool)
|
|
||||||
.await
|
|
||||||
.map_postgres_err()?;
|
|
||||||
rows.iter().map(map_public_catalog_model_row).collect()
|
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn search_public_catalog_models(
|
async fn search_public_catalog_models(
|
||||||
@@ -822,12 +798,8 @@ LIMIT 1
|
|||||||
builder
|
builder
|
||||||
.push(" ORDER BY p.provider_priority ASC, p.name ASC, COALESCE(gm.name, m.provider_model_name) ASC, m.id ASC LIMIT ")
|
.push(" ORDER BY p.provider_priority ASC, p.name ASC, COALESCE(gm.name, m.provider_model_name) ASC, m.id ASC LIMIT ")
|
||||||
.push_bind(query.limit as i64);
|
.push_bind(query.limit as i64);
|
||||||
let rows = builder
|
let query = builder.build();
|
||||||
.build()
|
collect_query_rows(query.fetch(&self.pool), map_public_catalog_model_row).await
|
||||||
.fetch_all(&self.pool)
|
|
||||||
.await
|
|
||||||
.map_postgres_err()?;
|
|
||||||
rows.iter().map(map_public_catalog_model_row).collect()
|
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn list_admin_global_models(
|
async fn list_admin_global_models(
|
||||||
@@ -1012,6 +984,20 @@ fn map_row(row: &PgRow) -> Result<StoredPublicGlobalModel, DataLayerError> {
|
|||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn collect_query_rows<T, S>(
|
||||||
|
mut rows: S,
|
||||||
|
map_row: fn(&PgRow) -> Result<T, DataLayerError>,
|
||||||
|
) -> Result<Vec<T>, DataLayerError>
|
||||||
|
where
|
||||||
|
S: TryStream<Ok = PgRow, Error = sqlx::Error> + Unpin,
|
||||||
|
{
|
||||||
|
let mut items = Vec::new();
|
||||||
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||||
|
items.push(map_row(&row)?);
|
||||||
|
}
|
||||||
|
Ok(items)
|
||||||
|
}
|
||||||
|
|
||||||
fn apply_public_catalog_model_filters(
|
fn apply_public_catalog_model_filters(
|
||||||
builder: &mut QueryBuilder<'_, Postgres>,
|
builder: &mut QueryBuilder<'_, Postgres>,
|
||||||
provider_id: Option<&str>,
|
provider_id: Option<&str>,
|
||||||
|
|||||||
@@ -1,4 +1,5 @@
|
|||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
|
use futures_util::TryStreamExt;
|
||||||
use sqlx::{postgres::PgRow, PgPool, Row};
|
use sqlx::{postgres::PgRow, PgPool, Row};
|
||||||
|
|
||||||
use super::types::{
|
use super::types::{
|
||||||
@@ -269,20 +270,19 @@ impl ManagementTokenReadRepository for SqlxManagementTokenRepository {
|
|||||||
.map_postgres_err()?;
|
.map_postgres_err()?;
|
||||||
let total = count_row.try_get::<i64, _>("total").map_postgres_err()?;
|
let total = count_row.try_get::<i64, _>("total").map_postgres_err()?;
|
||||||
|
|
||||||
let rows = sqlx::query(LIST_MANAGEMENT_TOKENS_SQL)
|
let mut rows = sqlx::query(LIST_MANAGEMENT_TOKENS_SQL)
|
||||||
.bind(query.user_id.as_deref())
|
.bind(query.user_id.as_deref())
|
||||||
.bind(query.is_active)
|
.bind(query.is_active)
|
||||||
.bind(i64::try_from(query.offset).unwrap_or(i64::MAX))
|
.bind(i64::try_from(query.offset).unwrap_or(i64::MAX))
|
||||||
.bind(i64::try_from(query.limit).unwrap_or(i64::MAX))
|
.bind(i64::try_from(query.limit).unwrap_or(i64::MAX))
|
||||||
.fetch_all(&self.pool)
|
.fetch(&self.pool);
|
||||||
.await
|
let mut items = Vec::new();
|
||||||
.map_postgres_err()?;
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||||
|
items.push(map_token_with_user_row(&row)?);
|
||||||
|
}
|
||||||
|
|
||||||
Ok(StoredManagementTokenListPage {
|
Ok(StoredManagementTokenListPage {
|
||||||
items: rows
|
items,
|
||||||
.iter()
|
|
||||||
.map(map_token_with_user_row)
|
|
||||||
.collect::<Result<Vec<_>, _>>()?,
|
|
||||||
total: usize::try_from(total.max(0)).unwrap_or(usize::MAX),
|
total: usize::try_from(total.max(0)).unwrap_or(usize::MAX),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,4 +1,5 @@
|
|||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
|
use futures_util::TryStreamExt;
|
||||||
use sqlx::{postgres::PgRow, PgPool, Row};
|
use sqlx::{postgres::PgRow, PgPool, Row};
|
||||||
|
|
||||||
use super::types::{
|
use super::types::{
|
||||||
@@ -181,11 +182,12 @@ impl OAuthProviderReadRepository for SqlxOAuthProviderRepository {
|
|||||||
async fn list_oauth_provider_configs(
|
async fn list_oauth_provider_configs(
|
||||||
&self,
|
&self,
|
||||||
) -> Result<Vec<StoredOAuthProviderConfig>, DataLayerError> {
|
) -> Result<Vec<StoredOAuthProviderConfig>, DataLayerError> {
|
||||||
let rows = sqlx::query(LIST_OAUTH_PROVIDER_CONFIGS_SQL)
|
let mut rows = sqlx::query(LIST_OAUTH_PROVIDER_CONFIGS_SQL).fetch(&self.pool);
|
||||||
.fetch_all(&self.pool)
|
let mut items = Vec::new();
|
||||||
.await
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||||
.map_postgres_err()?;
|
items.push(map_oauth_provider_row(&row)?);
|
||||||
rows.iter().map(map_oauth_provider_row).collect()
|
}
|
||||||
|
Ok(items)
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn get_oauth_provider_config(
|
async fn get_oauth_provider_config(
|
||||||
|
|||||||
@@ -4,9 +4,9 @@ use std::sync::RwLock;
|
|||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
|
|
||||||
use super::{
|
use super::{
|
||||||
ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
|
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogReadRepository,
|
||||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyPage,
|
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||||
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
|
StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
|
||||||
};
|
};
|
||||||
use crate::DataLayerError;
|
use crate::DataLayerError;
|
||||||
|
|
||||||
@@ -146,6 +146,13 @@ impl ProviderCatalogReadRepository for InMemoryProviderCatalogReadRepository {
|
|||||||
Ok(keys)
|
Ok(keys)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn list_key_summaries_by_provider_ids(
|
||||||
|
&self,
|
||||||
|
provider_ids: &[String],
|
||||||
|
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||||
|
Self::list_keys_by_provider_ids(self, provider_ids).await
|
||||||
|
}
|
||||||
|
|
||||||
async fn list_keys_page(
|
async fn list_keys_page(
|
||||||
&self,
|
&self,
|
||||||
query: &ProviderCatalogKeyListQuery,
|
query: &ProviderCatalogKeyListQuery,
|
||||||
@@ -170,12 +177,28 @@ impl ProviderCatalogReadRepository for InMemoryProviderCatalogReadRepository {
|
|||||||
})
|
})
|
||||||
.cloned()
|
.cloned()
|
||||||
.collect::<Vec<_>>();
|
.collect::<Vec<_>>();
|
||||||
keys.sort_by(|left, right| {
|
match query.order {
|
||||||
left.internal_priority
|
ProviderCatalogKeyListOrder::Name => {
|
||||||
.cmp(&right.internal_priority)
|
keys.sort_by(|left, right| {
|
||||||
.then(left.name.cmp(&right.name))
|
left.internal_priority
|
||||||
.then(left.id.cmp(&right.id))
|
.cmp(&right.internal_priority)
|
||||||
});
|
.then(left.name.cmp(&right.name))
|
||||||
|
.then(left.id.cmp(&right.id))
|
||||||
|
});
|
||||||
|
}
|
||||||
|
ProviderCatalogKeyListOrder::CreatedAt => {
|
||||||
|
keys.sort_by(|left, right| {
|
||||||
|
left.internal_priority
|
||||||
|
.cmp(&right.internal_priority)
|
||||||
|
.then(
|
||||||
|
left.created_at_unix_ms
|
||||||
|
.unwrap_or_default()
|
||||||
|
.cmp(&right.created_at_unix_ms.unwrap_or_default()),
|
||||||
|
)
|
||||||
|
.then(left.id.cmp(&right.id))
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
let total = keys.len();
|
let total = keys.len();
|
||||||
let items = keys
|
let items = keys
|
||||||
.into_iter()
|
.into_iter()
|
||||||
@@ -415,8 +438,9 @@ impl ProviderCatalogWriteRepository for InMemoryProviderCatalogReadRepository {
|
|||||||
mod tests {
|
mod tests {
|
||||||
use super::InMemoryProviderCatalogReadRepository;
|
use super::InMemoryProviderCatalogReadRepository;
|
||||||
use crate::repository::provider_catalog::{
|
use crate::repository::provider_catalog::{
|
||||||
ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
|
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogReadRepository,
|
||||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||||
|
StoredProviderCatalogProvider,
|
||||||
};
|
};
|
||||||
|
|
||||||
fn sample_provider(id: &str) -> StoredProviderCatalogProvider {
|
fn sample_provider(id: &str) -> StoredProviderCatalogProvider {
|
||||||
@@ -583,6 +607,7 @@ mod tests {
|
|||||||
is_active: Some(true),
|
is_active: Some(true),
|
||||||
offset: 0,
|
offset: 0,
|
||||||
limit: 10,
|
limit: 10,
|
||||||
|
order: ProviderCatalogKeyListOrder::Name,
|
||||||
})
|
})
|
||||||
.await
|
.await
|
||||||
.expect("keys should page");
|
.expect("keys should page");
|
||||||
@@ -598,6 +623,44 @@ mod tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn paginates_provider_keys_by_created_at_when_requested() {
|
||||||
|
let mut early = sample_key("key-1", "provider-1");
|
||||||
|
early.name = "zeta".to_string();
|
||||||
|
early.internal_priority = 10;
|
||||||
|
early.created_at_unix_ms = Some(10);
|
||||||
|
let mut late = sample_key("key-2", "provider-1");
|
||||||
|
late.name = "alpha".to_string();
|
||||||
|
late.internal_priority = 10;
|
||||||
|
late.created_at_unix_ms = Some(20);
|
||||||
|
let repository = InMemoryProviderCatalogReadRepository::seed(
|
||||||
|
vec![sample_provider("provider-1")],
|
||||||
|
vec![],
|
||||||
|
vec![late, early],
|
||||||
|
);
|
||||||
|
|
||||||
|
let page = repository
|
||||||
|
.list_keys_page(&ProviderCatalogKeyListQuery {
|
||||||
|
provider_id: "provider-1".to_string(),
|
||||||
|
search: None,
|
||||||
|
is_active: None,
|
||||||
|
offset: 0,
|
||||||
|
limit: 10,
|
||||||
|
order: ProviderCatalogKeyListOrder::CreatedAt,
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.expect("keys should page");
|
||||||
|
|
||||||
|
assert_eq!(page.total, 2);
|
||||||
|
assert_eq!(
|
||||||
|
page.items
|
||||||
|
.iter()
|
||||||
|
.map(|item| item.id.as_str())
|
||||||
|
.collect::<Vec<_>>(),
|
||||||
|
vec!["key-1", "key-2"]
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn summarizes_provider_key_stats() {
|
async fn summarizes_provider_key_stats() {
|
||||||
let mut inactive = sample_key("key-2", "provider-1");
|
let mut inactive = sample_key("key-2", "provider-1");
|
||||||
|
|||||||
@@ -3,9 +3,9 @@ mod sql;
|
|||||||
|
|
||||||
#[allow(unused_imports)]
|
#[allow(unused_imports)]
|
||||||
pub(crate) use aether_data_contracts::repository::provider_catalog::{
|
pub(crate) use aether_data_contracts::repository::provider_catalog::{
|
||||||
ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
|
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogReadRepository,
|
||||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyPage,
|
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||||
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
|
StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
|
||||||
};
|
};
|
||||||
pub use memory::InMemoryProviderCatalogReadRepository;
|
pub use memory::InMemoryProviderCatalogReadRepository;
|
||||||
pub use sql::SqlxProviderCatalogReadRepository;
|
pub use sql::SqlxProviderCatalogReadRepository;
|
||||||
|
|||||||
@@ -1,10 +1,11 @@
|
|||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
|
use futures_util::TryStreamExt;
|
||||||
use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row};
|
use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row};
|
||||||
|
|
||||||
use super::{
|
use super::{
|
||||||
ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
|
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogReadRepository,
|
||||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyPage,
|
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||||
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
|
StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
|
||||||
};
|
};
|
||||||
use crate::{
|
use crate::{
|
||||||
error::{postgres_error, SqlxResultExt},
|
error::{postgres_error, SqlxResultExt},
|
||||||
@@ -240,6 +241,61 @@ FROM provider_api_keys
|
|||||||
WHERE provider_id IN (
|
WHERE provider_id IN (
|
||||||
"#;
|
"#;
|
||||||
|
|
||||||
|
const LIST_KEY_SUMMARIES_BY_PROVIDER_IDS_PREFIX: &str = r#"
|
||||||
|
SELECT
|
||||||
|
id,
|
||||||
|
provider_id,
|
||||||
|
COALESCE(NULLIF(name, ''), id) AS name,
|
||||||
|
COALESCE(NULLIF(auth_type, ''), 'summary') AS auth_type,
|
||||||
|
capabilities,
|
||||||
|
is_active,
|
||||||
|
api_formats,
|
||||||
|
'summary' AS api_key,
|
||||||
|
NULL::text AS auth_config,
|
||||||
|
NULL::text AS note,
|
||||||
|
internal_priority,
|
||||||
|
rate_multipliers,
|
||||||
|
global_priority_by_format,
|
||||||
|
NULL::jsonb AS allowed_models,
|
||||||
|
NULL::bigint AS expires_at_unix_secs,
|
||||||
|
cache_ttl_minutes,
|
||||||
|
max_probe_interval_minutes,
|
||||||
|
NULL::jsonb AS proxy,
|
||||||
|
NULL::jsonb AS fingerprint,
|
||||||
|
NULL::integer AS rpm_limit,
|
||||||
|
NULL::integer AS learned_rpm_limit,
|
||||||
|
NULL::integer AS concurrent_429_count,
|
||||||
|
NULL::integer AS rpm_429_count,
|
||||||
|
NULL::bigint AS last_429_at_unix_secs,
|
||||||
|
NULL::text AS last_429_type,
|
||||||
|
NULL::jsonb AS adjustment_history,
|
||||||
|
NULL::jsonb AS utilization_samples,
|
||||||
|
NULL::bigint AS last_probe_increase_at_unix_secs,
|
||||||
|
request_count,
|
||||||
|
0::bigint AS total_tokens,
|
||||||
|
0::double precision AS total_cost_usd,
|
||||||
|
success_count,
|
||||||
|
NULL::integer AS error_count,
|
||||||
|
total_response_time_ms,
|
||||||
|
EXTRACT(EPOCH FROM last_used_at)::bigint AS last_used_at_unix_secs,
|
||||||
|
auto_fetch_models,
|
||||||
|
NULL::bigint AS last_models_fetch_at_unix_secs,
|
||||||
|
NULL::text AS last_models_fetch_error,
|
||||||
|
NULL::jsonb AS locked_models,
|
||||||
|
NULL::jsonb AS model_include_patterns,
|
||||||
|
NULL::jsonb AS model_exclude_patterns,
|
||||||
|
NULL::jsonb AS upstream_metadata,
|
||||||
|
EXTRACT(EPOCH FROM oauth_invalid_at)::bigint AS oauth_invalid_at_unix_secs,
|
||||||
|
oauth_invalid_reason,
|
||||||
|
NULL::jsonb AS status_snapshot,
|
||||||
|
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_ms,
|
||||||
|
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs,
|
||||||
|
health_by_format,
|
||||||
|
circuit_breaker_by_format
|
||||||
|
FROM provider_api_keys
|
||||||
|
WHERE provider_id IN (
|
||||||
|
"#;
|
||||||
|
|
||||||
const LIST_KEY_STATS_BY_PROVIDER_IDS_PREFIX: &str = r#"
|
const LIST_KEY_STATS_BY_PROVIDER_IDS_PREFIX: &str = r#"
|
||||||
SELECT
|
SELECT
|
||||||
provider_id,
|
provider_id,
|
||||||
@@ -271,24 +327,26 @@ impl SqlxProviderCatalogReadRepository {
|
|||||||
return Ok(Vec::new());
|
return Ok(Vec::new());
|
||||||
}
|
}
|
||||||
|
|
||||||
let rows = build_list_query(
|
collect_query_rows(
|
||||||
LIST_PROVIDERS_BY_IDS_PREFIX,
|
build_list_query(
|
||||||
provider_ids,
|
LIST_PROVIDERS_BY_IDS_PREFIX,
|
||||||
" ORDER BY name ASC",
|
provider_ids,
|
||||||
|
" ORDER BY name ASC",
|
||||||
|
)
|
||||||
|
.build()
|
||||||
|
.fetch(&self.pool),
|
||||||
|
map_provider_row,
|
||||||
)
|
)
|
||||||
.build()
|
|
||||||
.fetch_all(&self.pool)
|
|
||||||
.await
|
.await
|
||||||
.map_postgres_err()?;
|
|
||||||
rows.iter().map(map_provider_row).collect()
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn list_providers(
|
pub async fn list_providers(
|
||||||
&self,
|
&self,
|
||||||
active_only: bool,
|
active_only: bool,
|
||||||
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
|
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
|
||||||
let rows = sqlx::query(
|
collect_query_rows(
|
||||||
r#"
|
sqlx::query(
|
||||||
|
r#"
|
||||||
SELECT
|
SELECT
|
||||||
id,
|
id,
|
||||||
name,
|
name,
|
||||||
@@ -317,12 +375,12 @@ FROM providers
|
|||||||
WHERE ($1::boolean = false OR is_active = true)
|
WHERE ($1::boolean = false OR is_active = true)
|
||||||
ORDER BY provider_priority ASC, name ASC
|
ORDER BY provider_priority ASC, name ASC
|
||||||
"#,
|
"#,
|
||||||
|
)
|
||||||
|
.bind(active_only)
|
||||||
|
.fetch(&self.pool),
|
||||||
|
map_provider_row,
|
||||||
)
|
)
|
||||||
.bind(active_only)
|
|
||||||
.fetch_all(&self.pool)
|
|
||||||
.await
|
.await
|
||||||
.map_postgres_err()?;
|
|
||||||
rows.iter().map(map_provider_row).collect()
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn list_endpoints_by_ids(
|
pub async fn list_endpoints_by_ids(
|
||||||
@@ -333,28 +391,35 @@ ORDER BY provider_priority ASC, name ASC
|
|||||||
return Ok(Vec::new());
|
return Ok(Vec::new());
|
||||||
}
|
}
|
||||||
|
|
||||||
let rows = match build_list_query(
|
let rows = match collect_query_rows(
|
||||||
LIST_ENDPOINTS_BY_IDS_PREFIX,
|
build_list_query(
|
||||||
endpoint_ids,
|
LIST_ENDPOINTS_BY_IDS_PREFIX,
|
||||||
" ORDER BY api_format ASC, id ASC",
|
|
||||||
)
|
|
||||||
.build()
|
|
||||||
.fetch_all(&self.pool)
|
|
||||||
.await
|
|
||||||
{
|
|
||||||
Ok(rows) => rows,
|
|
||||||
Err(error) if is_missing_endpoint_health_score_column(&error) => build_list_query(
|
|
||||||
LIST_ENDPOINTS_BY_IDS_PREFIX_LEGACY,
|
|
||||||
endpoint_ids,
|
endpoint_ids,
|
||||||
" ORDER BY api_format ASC, id ASC",
|
" ORDER BY api_format ASC, id ASC",
|
||||||
)
|
)
|
||||||
.build()
|
.build()
|
||||||
.fetch_all(&self.pool)
|
.fetch(&self.pool),
|
||||||
.await
|
map_endpoint_row,
|
||||||
.map_postgres_err()?,
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(rows) => rows,
|
||||||
|
Err(error) if is_missing_endpoint_health_score_column_sql(&error) => {
|
||||||
|
collect_query_rows(
|
||||||
|
build_list_query(
|
||||||
|
LIST_ENDPOINTS_BY_IDS_PREFIX_LEGACY,
|
||||||
|
endpoint_ids,
|
||||||
|
" ORDER BY api_format ASC, id ASC",
|
||||||
|
)
|
||||||
|
.build()
|
||||||
|
.fetch(&self.pool),
|
||||||
|
map_endpoint_row,
|
||||||
|
)
|
||||||
|
.await?
|
||||||
|
}
|
||||||
Err(error) => return Err(postgres_error(error)),
|
Err(error) => return Err(postgres_error(error)),
|
||||||
};
|
};
|
||||||
rows.iter().map(map_endpoint_row).collect()
|
Ok(rows)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn list_endpoints_by_provider_ids(
|
pub async fn list_endpoints_by_provider_ids(
|
||||||
@@ -365,28 +430,35 @@ ORDER BY provider_priority ASC, name ASC
|
|||||||
return Ok(Vec::new());
|
return Ok(Vec::new());
|
||||||
}
|
}
|
||||||
|
|
||||||
let rows = match build_list_query(
|
let rows = match collect_query_rows(
|
||||||
LIST_ENDPOINTS_BY_PROVIDER_IDS_PREFIX,
|
build_list_query(
|
||||||
provider_ids,
|
LIST_ENDPOINTS_BY_PROVIDER_IDS_PREFIX,
|
||||||
" ORDER BY provider_id ASC, api_format ASC, id ASC",
|
|
||||||
)
|
|
||||||
.build()
|
|
||||||
.fetch_all(&self.pool)
|
|
||||||
.await
|
|
||||||
{
|
|
||||||
Ok(rows) => rows,
|
|
||||||
Err(error) if is_missing_endpoint_health_score_column(&error) => build_list_query(
|
|
||||||
LIST_ENDPOINTS_BY_PROVIDER_IDS_PREFIX_LEGACY,
|
|
||||||
provider_ids,
|
provider_ids,
|
||||||
" ORDER BY provider_id ASC, api_format ASC, id ASC",
|
" ORDER BY provider_id ASC, api_format ASC, id ASC",
|
||||||
)
|
)
|
||||||
.build()
|
.build()
|
||||||
.fetch_all(&self.pool)
|
.fetch(&self.pool),
|
||||||
.await
|
map_endpoint_row,
|
||||||
.map_postgres_err()?,
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(rows) => rows,
|
||||||
|
Err(error) if is_missing_endpoint_health_score_column_sql(&error) => {
|
||||||
|
collect_query_rows(
|
||||||
|
build_list_query(
|
||||||
|
LIST_ENDPOINTS_BY_PROVIDER_IDS_PREFIX_LEGACY,
|
||||||
|
provider_ids,
|
||||||
|
" ORDER BY provider_id ASC, api_format ASC, id ASC",
|
||||||
|
)
|
||||||
|
.build()
|
||||||
|
.fetch(&self.pool),
|
||||||
|
map_endpoint_row,
|
||||||
|
)
|
||||||
|
.await?
|
||||||
|
}
|
||||||
Err(error) => return Err(postgres_error(error)),
|
Err(error) => return Err(postgres_error(error)),
|
||||||
};
|
};
|
||||||
rows.iter().map(map_endpoint_row).collect()
|
Ok(rows)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn list_keys_by_ids(
|
pub async fn list_keys_by_ids(
|
||||||
@@ -397,16 +469,17 @@ ORDER BY provider_priority ASC, name ASC
|
|||||||
return Ok(Vec::new());
|
return Ok(Vec::new());
|
||||||
}
|
}
|
||||||
|
|
||||||
let rows = build_list_query(
|
collect_query_rows(
|
||||||
LIST_KEYS_BY_IDS_PREFIX,
|
build_list_query(
|
||||||
key_ids,
|
LIST_KEYS_BY_IDS_PREFIX,
|
||||||
" ORDER BY name ASC, id ASC",
|
key_ids,
|
||||||
|
" ORDER BY name ASC, id ASC",
|
||||||
|
)
|
||||||
|
.build()
|
||||||
|
.fetch(&self.pool),
|
||||||
|
map_key_row,
|
||||||
)
|
)
|
||||||
.build()
|
|
||||||
.fetch_all(&self.pool)
|
|
||||||
.await
|
.await
|
||||||
.map_postgres_err()?;
|
|
||||||
rows.iter().map(map_key_row).collect()
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn list_keys_by_provider_ids(
|
pub async fn list_keys_by_provider_ids(
|
||||||
@@ -417,16 +490,38 @@ ORDER BY provider_priority ASC, name ASC
|
|||||||
return Ok(Vec::new());
|
return Ok(Vec::new());
|
||||||
}
|
}
|
||||||
|
|
||||||
let rows = build_list_query(
|
collect_query_rows(
|
||||||
LIST_KEYS_BY_PROVIDER_IDS_PREFIX,
|
build_list_query(
|
||||||
provider_ids,
|
LIST_KEYS_BY_PROVIDER_IDS_PREFIX,
|
||||||
" ORDER BY provider_id ASC, name ASC, id ASC",
|
provider_ids,
|
||||||
|
" ORDER BY provider_id ASC, name ASC, id ASC",
|
||||||
|
)
|
||||||
|
.build()
|
||||||
|
.fetch(&self.pool),
|
||||||
|
map_key_row,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn list_key_summaries_by_provider_ids(
|
||||||
|
&self,
|
||||||
|
provider_ids: &[String],
|
||||||
|
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||||
|
if provider_ids.is_empty() {
|
||||||
|
return Ok(Vec::new());
|
||||||
|
}
|
||||||
|
|
||||||
|
collect_query_rows(
|
||||||
|
build_list_query(
|
||||||
|
LIST_KEY_SUMMARIES_BY_PROVIDER_IDS_PREFIX,
|
||||||
|
provider_ids,
|
||||||
|
" ORDER BY provider_id ASC, id ASC",
|
||||||
|
)
|
||||||
|
.build()
|
||||||
|
.fetch(&self.pool),
|
||||||
|
map_key_row,
|
||||||
)
|
)
|
||||||
.build()
|
|
||||||
.fetch_all(&self.pool)
|
|
||||||
.await
|
.await
|
||||||
.map_postgres_err()?;
|
|
||||||
rows.iter().map(map_key_row).collect()
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn list_keys_page(
|
pub async fn list_keys_page(
|
||||||
@@ -457,6 +552,12 @@ ORDER BY provider_priority ASC, name ASC
|
|||||||
.map(str::trim)
|
.map(str::trim)
|
||||||
.filter(|value| !value.is_empty())
|
.filter(|value| !value.is_empty())
|
||||||
.map(|value| format!("%{}%", value.to_ascii_lowercase()));
|
.map(|value| format!("%{}%", value.to_ascii_lowercase()));
|
||||||
|
let order_by = match query.order {
|
||||||
|
ProviderCatalogKeyListOrder::Name => "internal_priority ASC, name ASC, id ASC",
|
||||||
|
ProviderCatalogKeyListOrder::CreatedAt => {
|
||||||
|
"internal_priority ASC, COALESCE(created_at, TO_TIMESTAMP(0)) ASC, id ASC"
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
let count_row = sqlx::query(
|
let count_row = sqlx::query(
|
||||||
r#"
|
r#"
|
||||||
@@ -475,7 +576,7 @@ WHERE provider_id = $1
|
|||||||
.map_postgres_err()?;
|
.map_postgres_err()?;
|
||||||
let total = row_get::<i64>(&count_row, "total")?.max(0) as usize;
|
let total = row_get::<i64>(&count_row, "total")?.max(0) as usize;
|
||||||
|
|
||||||
let rows = sqlx::query(
|
let sql = format!(
|
||||||
r#"
|
r#"
|
||||||
SELECT
|
SELECT
|
||||||
id,
|
id,
|
||||||
@@ -531,23 +632,22 @@ FROM provider_api_keys
|
|||||||
WHERE provider_id = $1
|
WHERE provider_id = $1
|
||||||
AND ($2::TEXT IS NULL OR LOWER(name) LIKE $2 OR LOWER(id) LIKE $2)
|
AND ($2::TEXT IS NULL OR LOWER(name) LIKE $2 OR LOWER(id) LIKE $2)
|
||||||
AND ($3::BOOLEAN IS NULL OR is_active = $3)
|
AND ($3::BOOLEAN IS NULL OR is_active = $3)
|
||||||
ORDER BY internal_priority ASC, name ASC, id ASC
|
ORDER BY {order_by}
|
||||||
OFFSET $4
|
OFFSET $4
|
||||||
LIMIT $5
|
LIMIT $5
|
||||||
"#,
|
"#,
|
||||||
|
);
|
||||||
|
let items = collect_query_rows(
|
||||||
|
sqlx::query(&sql)
|
||||||
|
.bind(&query.provider_id)
|
||||||
|
.bind(search_pattern.as_deref())
|
||||||
|
.bind(query.is_active)
|
||||||
|
.bind(offset)
|
||||||
|
.bind(limit)
|
||||||
|
.fetch(&self.pool),
|
||||||
|
map_key_row,
|
||||||
)
|
)
|
||||||
.bind(&query.provider_id)
|
.await?;
|
||||||
.bind(search_pattern.as_deref())
|
|
||||||
.bind(query.is_active)
|
|
||||||
.bind(offset)
|
|
||||||
.bind(limit)
|
|
||||||
.fetch_all(&self.pool)
|
|
||||||
.await
|
|
||||||
.map_postgres_err()?;
|
|
||||||
let items = rows
|
|
||||||
.iter()
|
|
||||||
.map(map_key_row)
|
|
||||||
.collect::<Result<Vec<_>, _>>()?;
|
|
||||||
|
|
||||||
Ok(StoredProviderCatalogKeyPage { items, total })
|
Ok(StoredProviderCatalogKeyPage { items, total })
|
||||||
}
|
}
|
||||||
@@ -560,16 +660,17 @@ LIMIT $5
|
|||||||
return Ok(Vec::new());
|
return Ok(Vec::new());
|
||||||
}
|
}
|
||||||
|
|
||||||
let rows = build_list_query(
|
collect_query_rows(
|
||||||
LIST_KEY_STATS_BY_PROVIDER_IDS_PREFIX,
|
build_list_query(
|
||||||
provider_ids,
|
LIST_KEY_STATS_BY_PROVIDER_IDS_PREFIX,
|
||||||
"\nGROUP BY provider_id\nORDER BY provider_id ASC",
|
provider_ids,
|
||||||
|
"\nGROUP BY provider_id\nORDER BY provider_id ASC",
|
||||||
|
)
|
||||||
|
.build()
|
||||||
|
.fetch(&self.pool),
|
||||||
|
map_key_stats_row,
|
||||||
)
|
)
|
||||||
.build()
|
|
||||||
.fetch_all(&self.pool)
|
|
||||||
.await
|
.await
|
||||||
.map_postgres_err()?;
|
|
||||||
rows.iter().map(map_key_stats_row).collect()
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn update_key_oauth_credentials(
|
pub async fn update_key_oauth_credentials(
|
||||||
@@ -1794,6 +1895,13 @@ impl ProviderCatalogReadRepository for SqlxProviderCatalogReadRepository {
|
|||||||
Self::list_keys_by_provider_ids(self, provider_ids).await
|
Self::list_keys_by_provider_ids(self, provider_ids).await
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn list_key_summaries_by_provider_ids(
|
||||||
|
&self,
|
||||||
|
provider_ids: &[String],
|
||||||
|
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||||
|
Self::list_key_summaries_by_provider_ids(self, provider_ids).await
|
||||||
|
}
|
||||||
|
|
||||||
async fn list_keys_page(
|
async fn list_keys_page(
|
||||||
&self,
|
&self,
|
||||||
query: &ProviderCatalogKeyListQuery,
|
query: &ProviderCatalogKeyListQuery,
|
||||||
@@ -2051,6 +2159,29 @@ fn is_missing_endpoint_health_score_column(error: &sqlx::Error) -> bool {
|
|||||||
.unwrap_or(false)
|
.unwrap_or(false)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn is_missing_endpoint_health_score_column_sql(error: &DataLayerError) -> bool {
|
||||||
|
match error {
|
||||||
|
DataLayerError::Postgres(message) => {
|
||||||
|
message.contains("endpoint_health_score") && message.contains("does not exist")
|
||||||
|
}
|
||||||
|
_ => false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn collect_query_rows<T, S>(
|
||||||
|
mut rows: S,
|
||||||
|
mapper: fn(&PgRow) -> Result<T, DataLayerError>,
|
||||||
|
) -> Result<Vec<T>, DataLayerError>
|
||||||
|
where
|
||||||
|
S: futures_util::TryStream<Ok = PgRow, Error = sqlx::Error> + Unpin,
|
||||||
|
{
|
||||||
|
let mut items = Vec::new();
|
||||||
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||||
|
items.push(mapper(&row)?);
|
||||||
|
}
|
||||||
|
Ok(items)
|
||||||
|
}
|
||||||
|
|
||||||
fn map_key_stats_row(row: &PgRow) -> Result<StoredProviderCatalogKeyStats, DataLayerError> {
|
fn map_key_stats_row(row: &PgRow) -> Result<StoredProviderCatalogKeyStats, DataLayerError> {
|
||||||
StoredProviderCatalogKeyStats::new(
|
StoredProviderCatalogKeyStats::new(
|
||||||
row_get(row, "provider_id")?,
|
row_get(row, "provider_id")?,
|
||||||
|
|||||||
@@ -1,4 +1,5 @@
|
|||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
|
use futures_util::TryStreamExt;
|
||||||
use sha2::{Digest, Sha256};
|
use sha2::{Digest, Sha256};
|
||||||
use sqlx::{postgres::PgRow, PgPool, Row};
|
use sqlx::{postgres::PgRow, PgPool, Row};
|
||||||
|
|
||||||
@@ -414,11 +415,12 @@ impl SqlxProxyNodeRepository {
|
|||||||
#[async_trait]
|
#[async_trait]
|
||||||
impl ProxyNodeReadRepository for SqlxProxyNodeRepository {
|
impl ProxyNodeReadRepository for SqlxProxyNodeRepository {
|
||||||
async fn list_proxy_nodes(&self) -> Result<Vec<StoredProxyNode>, DataLayerError> {
|
async fn list_proxy_nodes(&self) -> Result<Vec<StoredProxyNode>, DataLayerError> {
|
||||||
let rows = sqlx::query(LIST_PROXY_NODES_SQL)
|
let mut rows = sqlx::query(LIST_PROXY_NODES_SQL).fetch(&self.pool);
|
||||||
.fetch_all(&self.pool)
|
let mut items = Vec::new();
|
||||||
.await
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||||
.map_postgres_err()?;
|
items.push(Self::row_to_stored(&row)?);
|
||||||
rows.iter().map(Self::row_to_stored).collect()
|
}
|
||||||
|
Ok(items)
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn find_proxy_node(
|
async fn find_proxy_node(
|
||||||
@@ -438,13 +440,15 @@ impl ProxyNodeReadRepository for SqlxProxyNodeRepository {
|
|||||||
node_id: &str,
|
node_id: &str,
|
||||||
limit: usize,
|
limit: usize,
|
||||||
) -> Result<Vec<StoredProxyNodeEvent>, DataLayerError> {
|
) -> Result<Vec<StoredProxyNodeEvent>, DataLayerError> {
|
||||||
let rows = sqlx::query(LIST_PROXY_NODE_EVENTS_SQL)
|
let mut rows = sqlx::query(LIST_PROXY_NODE_EVENTS_SQL)
|
||||||
.bind(node_id)
|
.bind(node_id)
|
||||||
.bind(i64::try_from(limit).unwrap_or(i64::MAX))
|
.bind(i64::try_from(limit).unwrap_or(i64::MAX))
|
||||||
.fetch_all(&self.pool)
|
.fetch(&self.pool);
|
||||||
.await
|
let mut items = Vec::new();
|
||||||
.map_postgres_err()?;
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||||
rows.iter().map(Self::row_to_event).collect()
|
items.push(Self::row_to_event(&row)?);
|
||||||
|
}
|
||||||
|
Ok(items)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use futures_util::future::BoxFuture;
|
use futures_util::future::BoxFuture;
|
||||||
|
use futures_util::TryStreamExt;
|
||||||
use sqlx::{PgPool, Row};
|
use sqlx::{PgPool, Row};
|
||||||
|
|
||||||
use super::types::{
|
use super::types::{
|
||||||
@@ -161,17 +162,18 @@ impl SqlxShadowResultRepository {
|
|||||||
return Ok(Vec::new());
|
return Ok(Vec::new());
|
||||||
}
|
}
|
||||||
|
|
||||||
let rows = sqlx::query(LIST_RECENT_SQL)
|
let mut rows = sqlx::query(LIST_RECENT_SQL)
|
||||||
.bind(i64::try_from(limit).map_err(|_| {
|
.bind(i64::try_from(limit).map_err(|_| {
|
||||||
DataLayerError::UnexpectedValue(format!(
|
DataLayerError::UnexpectedValue(format!(
|
||||||
"invalid recent shadow result limit: {limit}"
|
"invalid recent shadow result limit: {limit}"
|
||||||
))
|
))
|
||||||
})?)
|
})?)
|
||||||
.fetch_all(&self.pool)
|
.fetch(&self.pool);
|
||||||
.await
|
let mut items = Vec::new();
|
||||||
.map_postgres_err()?;
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||||
|
items.push(map_shadow_result_row(&row)?);
|
||||||
rows.iter().map(map_shadow_result_row).collect()
|
}
|
||||||
|
Ok(items)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn upsert(
|
pub async fn upsert(
|
||||||
|
|||||||
@@ -4,6 +4,7 @@ use aether_data_contracts::repository::usage::{
|
|||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use flate2::{read::GzDecoder, write::GzEncoder, Compression};
|
use flate2::{read::GzDecoder, write::GzEncoder, Compression};
|
||||||
use futures_util::future::BoxFuture;
|
use futures_util::future::BoxFuture;
|
||||||
|
use futures_util::TryStreamExt;
|
||||||
use serde_json::Map;
|
use serde_json::Map;
|
||||||
use serde_json::Value;
|
use serde_json::Value;
|
||||||
use sqlx::{PgPool, Postgres, QueryBuilder, Row};
|
use sqlx::{PgPool, Postgres, QueryBuilder, Row};
|
||||||
@@ -1181,12 +1182,13 @@ impl SqlxUsageReadRepository {
|
|||||||
}
|
}
|
||||||
|
|
||||||
builder.push(" ORDER BY \"usage\".created_at ASC, \"usage\".request_id ASC");
|
builder.push(" ORDER BY \"usage\".created_at ASC, \"usage\".request_id ASC");
|
||||||
let rows = builder
|
let query = builder.build();
|
||||||
.build()
|
let mut rows = query.fetch(&self.pool);
|
||||||
.fetch_all(&self.pool)
|
let mut items = Vec::new();
|
||||||
.await
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||||
.map_postgres_err()?;
|
items.push(map_usage_row(&row, false)?);
|
||||||
rows.iter().map(|row| map_usage_row(row, false)).collect()
|
}
|
||||||
|
Ok(items)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn list_recent_usage_audits(
|
pub async fn list_recent_usage_audits(
|
||||||
@@ -1205,12 +1207,13 @@ impl SqlxUsageReadRepository {
|
|||||||
.push_bind(i64::try_from(limit).map_err(|_| {
|
.push_bind(i64::try_from(limit).map_err(|_| {
|
||||||
DataLayerError::InvalidInput(format!("invalid recent usage limit: {limit}"))
|
DataLayerError::InvalidInput(format!("invalid recent usage limit: {limit}"))
|
||||||
})?);
|
})?);
|
||||||
let rows = builder
|
let query = builder.build();
|
||||||
.build()
|
let mut rows = query.fetch(&self.pool);
|
||||||
.fetch_all(&self.pool)
|
let mut items = Vec::new();
|
||||||
.await
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||||
.map_postgres_err()?;
|
items.push(map_usage_row(&row, false)?);
|
||||||
rows.iter().map(|row| map_usage_row(row, false)).collect()
|
}
|
||||||
|
Ok(items)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn summarize_total_tokens_by_api_key_ids(
|
pub async fn summarize_total_tokens_by_api_key_ids(
|
||||||
@@ -1221,14 +1224,12 @@ impl SqlxUsageReadRepository {
|
|||||||
return Ok(std::collections::BTreeMap::new());
|
return Ok(std::collections::BTreeMap::new());
|
||||||
}
|
}
|
||||||
|
|
||||||
let rows = sqlx::query(SUMMARIZE_TOTAL_TOKENS_BY_API_KEY_IDS_SQL)
|
let mut rows = sqlx::query(SUMMARIZE_TOTAL_TOKENS_BY_API_KEY_IDS_SQL)
|
||||||
.bind(api_key_ids)
|
.bind(api_key_ids)
|
||||||
.fetch_all(&self.pool)
|
.fetch(&self.pool);
|
||||||
.await
|
|
||||||
.map_postgres_err()?;
|
|
||||||
|
|
||||||
let mut totals = std::collections::BTreeMap::new();
|
let mut totals = std::collections::BTreeMap::new();
|
||||||
for row in rows {
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||||
let api_key_id: String = row.try_get("api_key_id").map_postgres_err()?;
|
let api_key_id: String = row.try_get("api_key_id").map_postgres_err()?;
|
||||||
let total_tokens = row
|
let total_tokens = row
|
||||||
.try_get::<i64, _>("total_tokens")
|
.try_get::<i64, _>("total_tokens")
|
||||||
@@ -1248,14 +1249,12 @@ impl SqlxUsageReadRepository {
|
|||||||
return Ok(std::collections::BTreeMap::new());
|
return Ok(std::collections::BTreeMap::new());
|
||||||
}
|
}
|
||||||
|
|
||||||
let rows = sqlx::query(SUMMARIZE_USAGE_BY_PROVIDER_API_KEY_IDS_SQL)
|
let mut rows = sqlx::query(SUMMARIZE_USAGE_BY_PROVIDER_API_KEY_IDS_SQL)
|
||||||
.bind(provider_api_key_ids)
|
.bind(provider_api_key_ids)
|
||||||
.fetch_all(&self.pool)
|
.fetch(&self.pool);
|
||||||
.await
|
|
||||||
.map_postgres_err()?;
|
|
||||||
|
|
||||||
let mut summaries = std::collections::BTreeMap::new();
|
let mut summaries = std::collections::BTreeMap::new();
|
||||||
for row in rows {
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||||
let provider_api_key_id: String =
|
let provider_api_key_id: String =
|
||||||
row.try_get("provider_api_key_id").map_postgres_err()?;
|
row.try_get("provider_api_key_id").map_postgres_err()?;
|
||||||
let request_count = row
|
let request_count = row
|
||||||
@@ -1988,7 +1987,7 @@ fn prepare_usage_body_storage(value: Option<&Value>) -> Result<UsageBodyStorage,
|
|||||||
let bytes = serde_json::to_vec(value).map_err(|err| {
|
let bytes = serde_json::to_vec(value).map_err(|err| {
|
||||||
DataLayerError::UnexpectedValue(format!("failed to serialize usage json: {err}"))
|
DataLayerError::UnexpectedValue(format!("failed to serialize usage json: {err}"))
|
||||||
})?;
|
})?;
|
||||||
if bytes.len() <= MAX_INLINE_USAGE_BODY_BYTES {
|
if bytes.len() == MAX_INLINE_USAGE_BODY_BYTES {
|
||||||
return Ok(UsageBodyStorage {
|
return Ok(UsageBodyStorage {
|
||||||
inline_json: Some(String::from_utf8(bytes).map_err(|err| {
|
inline_json: Some(String::from_utf8(bytes).map_err(|err| {
|
||||||
DataLayerError::UnexpectedValue(format!(
|
DataLayerError::UnexpectedValue(format!(
|
||||||
|
|||||||
@@ -1,4 +1,5 @@
|
|||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
|
use futures_util::TryStreamExt;
|
||||||
use sqlx::{PgPool, Postgres, QueryBuilder, Row};
|
use sqlx::{PgPool, Postgres, QueryBuilder, Row};
|
||||||
|
|
||||||
use super::types::{
|
use super::types::{
|
||||||
@@ -189,30 +190,31 @@ impl SqlxUserReadRepository {
|
|||||||
if user_ids.is_empty() {
|
if user_ids.is_empty() {
|
||||||
return Ok(Vec::new());
|
return Ok(Vec::new());
|
||||||
}
|
}
|
||||||
let rows = sqlx::query(LIST_USERS_BY_IDS_SQL)
|
collect_query_rows(
|
||||||
.bind(user_ids)
|
sqlx::query(LIST_USERS_BY_IDS_SQL)
|
||||||
.fetch_all(&self.pool)
|
.bind(user_ids)
|
||||||
.await
|
.fetch(&self.pool),
|
||||||
.map_postgres_err()?;
|
map_user_row,
|
||||||
rows.iter().map(map_user_row).collect()
|
)
|
||||||
|
.await
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn list_non_admin_export_users(
|
pub async fn list_non_admin_export_users(
|
||||||
&self,
|
&self,
|
||||||
) -> Result<Vec<StoredUserExportRow>, DataLayerError> {
|
) -> Result<Vec<StoredUserExportRow>, DataLayerError> {
|
||||||
let rows = sqlx::query(LIST_NON_ADMIN_EXPORT_USERS_SQL)
|
collect_query_rows(
|
||||||
.fetch_all(&self.pool)
|
sqlx::query(LIST_NON_ADMIN_EXPORT_USERS_SQL).fetch(&self.pool),
|
||||||
.await
|
map_user_export_row,
|
||||||
.map_postgres_err()?;
|
)
|
||||||
rows.iter().map(map_user_export_row).collect()
|
.await
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn list_export_users(&self) -> Result<Vec<StoredUserExportRow>, DataLayerError> {
|
pub async fn list_export_users(&self) -> Result<Vec<StoredUserExportRow>, DataLayerError> {
|
||||||
let rows = sqlx::query(LIST_EXPORT_USERS_SQL)
|
collect_query_rows(
|
||||||
.fetch_all(&self.pool)
|
sqlx::query(LIST_EXPORT_USERS_SQL).fetch(&self.pool),
|
||||||
.await
|
map_user_export_row,
|
||||||
.map_postgres_err()?;
|
)
|
||||||
rows.iter().map(map_user_export_row).collect()
|
.await
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn list_export_users_page(
|
pub async fn list_export_users_page(
|
||||||
@@ -240,12 +242,8 @@ impl SqlxUserReadRepository {
|
|||||||
DataLayerError::InvalidInput(format!("invalid user export limit: {}", query.limit))
|
DataLayerError::InvalidInput(format!("invalid user export limit: {}", query.limit))
|
||||||
})?);
|
})?);
|
||||||
|
|
||||||
let rows = builder
|
let query = builder.build();
|
||||||
.build()
|
collect_query_rows(query.fetch(&self.pool), map_user_export_row).await
|
||||||
.fetch_all(&self.pool)
|
|
||||||
.await
|
|
||||||
.map_postgres_err()?;
|
|
||||||
rows.iter().map(map_user_export_row).collect()
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn summarize_export_users(&self) -> Result<UserExportSummary, DataLayerError> {
|
pub async fn summarize_export_users(&self) -> Result<UserExportSummary, DataLayerError> {
|
||||||
@@ -279,12 +277,13 @@ impl SqlxUserReadRepository {
|
|||||||
return Ok(Vec::new());
|
return Ok(Vec::new());
|
||||||
}
|
}
|
||||||
|
|
||||||
let rows = sqlx::query(LIST_USER_AUTH_BY_IDS_SQL)
|
collect_query_rows(
|
||||||
.bind(user_ids)
|
sqlx::query(LIST_USER_AUTH_BY_IDS_SQL)
|
||||||
.fetch_all(&self.pool)
|
.bind(user_ids)
|
||||||
.await
|
.fetch(&self.pool),
|
||||||
.map_postgres_err()?;
|
map_user_auth_row,
|
||||||
rows.iter().map(map_user_auth_row).collect()
|
)
|
||||||
|
.await
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn find_user_auth_by_id(
|
pub async fn find_user_auth_by_id(
|
||||||
@@ -361,6 +360,20 @@ fn map_user_auth_row(row: &sqlx::postgres::PgRow) -> Result<StoredUserAuthRecord
|
|||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn collect_query_rows<T, S>(
|
||||||
|
mut rows: S,
|
||||||
|
mapper: fn(&sqlx::postgres::PgRow) -> Result<T, DataLayerError>,
|
||||||
|
) -> Result<Vec<T>, DataLayerError>
|
||||||
|
where
|
||||||
|
S: futures_util::TryStream<Ok = sqlx::postgres::PgRow, Error = sqlx::Error> + Unpin,
|
||||||
|
{
|
||||||
|
let mut items = Vec::new();
|
||||||
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||||
|
items.push(mapper(&row)?);
|
||||||
|
}
|
||||||
|
Ok(items)
|
||||||
|
}
|
||||||
|
|
||||||
#[async_trait]
|
#[async_trait]
|
||||||
impl UserReadRepository for SqlxUserReadRepository {
|
impl UserReadRepository for SqlxUserReadRepository {
|
||||||
async fn list_users_by_ids(
|
async fn list_users_by_ids(
|
||||||
|
|||||||
@@ -187,6 +187,15 @@ impl VideoTaskReadRepository for InMemoryVideoTaskRepository {
|
|||||||
Ok(tasks.into_iter().skip(offset).take(limit).collect())
|
Ok(tasks.into_iter().skip(offset).take(limit).collect())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn list_page_summary(
|
||||||
|
&self,
|
||||||
|
filter: &VideoTaskQueryFilter,
|
||||||
|
offset: usize,
|
||||||
|
limit: usize,
|
||||||
|
) -> Result<Vec<StoredVideoTask>, DataLayerError> {
|
||||||
|
Self::list_page(self, filter, offset, limit).await
|
||||||
|
}
|
||||||
|
|
||||||
async fn count(&self, filter: &VideoTaskQueryFilter) -> Result<u64, DataLayerError> {
|
async fn count(&self, filter: &VideoTaskQueryFilter) -> Result<u64, DataLayerError> {
|
||||||
Ok(self
|
Ok(self
|
||||||
.index
|
.index
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
use sqlx::{PgPool, Postgres, QueryBuilder, Row};
|
|
||||||
|
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
|
use futures_util::{stream::TryStream, TryStreamExt};
|
||||||
|
use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row};
|
||||||
|
|
||||||
use crate::error::SqlxResultExt;
|
use crate::error::SqlxResultExt;
|
||||||
use crate::repository::video_tasks::{
|
use crate::repository::video_tasks::{
|
||||||
@@ -10,7 +10,7 @@ use crate::repository::video_tasks::{
|
|||||||
};
|
};
|
||||||
use crate::DataLayerError;
|
use crate::DataLayerError;
|
||||||
|
|
||||||
const SELECT_VIDEO_TASK_COLUMNS: &str = r#"
|
const SELECT_VIDEO_TASK_COLUMNS_PREFIX: &str = r#"
|
||||||
id,
|
id,
|
||||||
short_id,
|
short_id,
|
||||||
request_id,
|
request_id,
|
||||||
@@ -26,12 +26,9 @@ const SELECT_VIDEO_TASK_COLUMNS: &str = r#"
|
|||||||
provider_api_format,
|
provider_api_format,
|
||||||
format_converted,
|
format_converted,
|
||||||
model,
|
model,
|
||||||
prompt,
|
"#;
|
||||||
original_request_body,
|
|
||||||
duration_seconds,
|
const SELECT_VIDEO_TASK_COLUMNS_SUFFIX: &str = r#"
|
||||||
resolution,
|
|
||||||
aspect_ratio,
|
|
||||||
size,
|
|
||||||
status,
|
status,
|
||||||
progress_percent,
|
progress_percent,
|
||||||
progress_message,
|
progress_message,
|
||||||
@@ -50,8 +47,53 @@ const SELECT_VIDEO_TASK_COLUMNS: &str = r#"
|
|||||||
request_metadata
|
request_metadata
|
||||||
"#;
|
"#;
|
||||||
|
|
||||||
|
fn select_video_task_columns(
|
||||||
|
prompt_sql: &str,
|
||||||
|
original_request_body_sql: &str,
|
||||||
|
duration_seconds_sql: &str,
|
||||||
|
resolution_sql: &str,
|
||||||
|
aspect_ratio_sql: &str,
|
||||||
|
size_sql: &str,
|
||||||
|
) -> String {
|
||||||
|
format!(
|
||||||
|
"{SELECT_VIDEO_TASK_COLUMNS_PREFIX}
|
||||||
|
{prompt_sql} AS prompt,
|
||||||
|
{original_request_body_sql} AS original_request_body,
|
||||||
|
{duration_seconds_sql} AS duration_seconds,
|
||||||
|
{resolution_sql} AS resolution,
|
||||||
|
{aspect_ratio_sql} AS aspect_ratio,
|
||||||
|
{size_sql} AS size,
|
||||||
|
{SELECT_VIDEO_TASK_COLUMNS_SUFFIX}"
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn select_video_task_full_columns() -> String {
|
||||||
|
select_video_task_columns(
|
||||||
|
"prompt",
|
||||||
|
"original_request_body",
|
||||||
|
"duration_seconds",
|
||||||
|
"resolution",
|
||||||
|
"aspect_ratio",
|
||||||
|
"size",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn select_video_task_claim_columns() -> String {
|
||||||
|
select_video_task_columns(
|
||||||
|
"NULL::TEXT",
|
||||||
|
"NULL::jsonb",
|
||||||
|
"NULL::INTEGER",
|
||||||
|
"NULL::TEXT",
|
||||||
|
"NULL::TEXT",
|
||||||
|
"NULL::TEXT",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
fn select_video_task_sql(where_clause: &str) -> String {
|
fn select_video_task_sql(where_clause: &str) -> String {
|
||||||
format!("SELECT\n{SELECT_VIDEO_TASK_COLUMNS}\nFROM video_tasks\n{where_clause}\n")
|
format!(
|
||||||
|
"SELECT\n{}\nFROM video_tasks\n{where_clause}\n",
|
||||||
|
select_video_task_full_columns()
|
||||||
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
fn find_by_id_sql() -> String {
|
fn find_by_id_sql() -> String {
|
||||||
@@ -76,7 +118,54 @@ fn list_due_sql() -> String {
|
|||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn select_video_task_page_summary_columns() -> &'static str {
|
||||||
|
r#"
|
||||||
|
id,
|
||||||
|
NULL::TEXT AS short_id,
|
||||||
|
request_id,
|
||||||
|
user_id,
|
||||||
|
NULL::TEXT AS api_key_id,
|
||||||
|
username,
|
||||||
|
NULL::TEXT AS api_key_name,
|
||||||
|
external_task_id,
|
||||||
|
provider_id,
|
||||||
|
NULL::TEXT AS endpoint_id,
|
||||||
|
NULL::TEXT AS key_id,
|
||||||
|
NULL::TEXT AS client_api_format,
|
||||||
|
NULL::TEXT AS provider_api_format,
|
||||||
|
FALSE AS format_converted,
|
||||||
|
model,
|
||||||
|
CASE
|
||||||
|
WHEN prompt IS NULL THEN NULL
|
||||||
|
WHEN char_length(prompt) <= 100 THEN prompt
|
||||||
|
ELSE LEFT(prompt, 100) || '...'
|
||||||
|
END AS prompt,
|
||||||
|
NULL::jsonb AS original_request_body,
|
||||||
|
duration_seconds,
|
||||||
|
resolution,
|
||||||
|
aspect_ratio,
|
||||||
|
NULL::TEXT AS size,
|
||||||
|
status,
|
||||||
|
progress_percent,
|
||||||
|
progress_message,
|
||||||
|
0::INTEGER AS retry_count,
|
||||||
|
1::INTEGER AS poll_interval_seconds,
|
||||||
|
NULL::BIGINT AS next_poll_at_unix_secs,
|
||||||
|
poll_count,
|
||||||
|
max_poll_count,
|
||||||
|
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_ms,
|
||||||
|
CAST(EXTRACT(EPOCH FROM submitted_at) AS BIGINT) AS submitted_at_unix_secs,
|
||||||
|
CAST(EXTRACT(EPOCH FROM completed_at) AS BIGINT) AS completed_at_unix_secs,
|
||||||
|
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs,
|
||||||
|
error_code,
|
||||||
|
error_message,
|
||||||
|
video_url,
|
||||||
|
NULL::jsonb AS request_metadata
|
||||||
|
"#
|
||||||
|
}
|
||||||
|
|
||||||
fn claim_due_sql() -> String {
|
fn claim_due_sql() -> String {
|
||||||
|
let columns = select_video_task_claim_columns();
|
||||||
format!(
|
format!(
|
||||||
"WITH due AS (
|
"WITH due AS (
|
||||||
SELECT id
|
SELECT id
|
||||||
@@ -94,12 +183,13 @@ SET next_poll_at = TO_TIMESTAMP($4),
|
|||||||
updated_at = TO_TIMESTAMP($5)
|
updated_at = TO_TIMESTAMP($5)
|
||||||
WHERE id IN (SELECT id FROM due)
|
WHERE id IN (SELECT id FROM due)
|
||||||
RETURNING
|
RETURNING
|
||||||
{SELECT_VIDEO_TASK_COLUMNS}
|
{columns}
|
||||||
"
|
"
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
fn upsert_sql() -> String {
|
fn upsert_sql() -> String {
|
||||||
|
let columns = select_video_task_full_columns();
|
||||||
format!(
|
format!(
|
||||||
"INSERT INTO video_tasks (
|
"INSERT INTO video_tasks (
|
||||||
id,
|
id,
|
||||||
@@ -216,12 +306,13 @@ ON CONFLICT (id) DO UPDATE SET
|
|||||||
completed_at = EXCLUDED.completed_at,
|
completed_at = EXCLUDED.completed_at,
|
||||||
updated_at = EXCLUDED.updated_at
|
updated_at = EXCLUDED.updated_at
|
||||||
RETURNING
|
RETURNING
|
||||||
{SELECT_VIDEO_TASK_COLUMNS}
|
{columns}
|
||||||
"
|
"
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
fn update_if_active_sql() -> String {
|
fn update_if_active_sql() -> String {
|
||||||
|
let columns = select_video_task_full_columns();
|
||||||
format!(
|
format!(
|
||||||
"UPDATE video_tasks SET
|
"UPDATE video_tasks SET
|
||||||
short_id = $2,
|
short_id = $2,
|
||||||
@@ -263,7 +354,7 @@ fn update_if_active_sql() -> String {
|
|||||||
WHERE id = $1
|
WHERE id = $1
|
||||||
AND status = ANY($38)
|
AND status = ANY($38)
|
||||||
RETURNING
|
RETURNING
|
||||||
{SELECT_VIDEO_TASK_COLUMNS}
|
{columns}
|
||||||
"
|
"
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
@@ -343,16 +434,16 @@ impl SqlxVideoTaskRepository {
|
|||||||
|
|
||||||
let active_statuses = vec!["pending", "submitted", "queued", "processing"];
|
let active_statuses = vec!["pending", "submitted", "queued", "processing"];
|
||||||
let sql = list_active_sql();
|
let sql = list_active_sql();
|
||||||
let rows = sqlx::query(&sql)
|
collect_query_rows(
|
||||||
.bind(active_statuses)
|
sqlx::query(&sql)
|
||||||
.bind(i64::try_from(limit).map_err(|_| {
|
.bind(active_statuses)
|
||||||
DataLayerError::UnexpectedValue(format!("invalid active task limit: {limit}"))
|
.bind(i64::try_from(limit).map_err(|_| {
|
||||||
})?)
|
DataLayerError::UnexpectedValue(format!("invalid active task limit: {limit}"))
|
||||||
.fetch_all(&self.pool)
|
})?)
|
||||||
.await
|
.fetch(&self.pool),
|
||||||
.map_postgres_err()?;
|
map_video_task_row,
|
||||||
|
)
|
||||||
rows.iter().map(map_video_task_row).collect()
|
.await
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn list_due(
|
pub async fn list_due(
|
||||||
@@ -366,17 +457,17 @@ impl SqlxVideoTaskRepository {
|
|||||||
|
|
||||||
let active_statuses = vec!["submitted", "queued", "processing"];
|
let active_statuses = vec!["submitted", "queued", "processing"];
|
||||||
let sql = list_due_sql();
|
let sql = list_due_sql();
|
||||||
let rows = sqlx::query(&sql)
|
collect_query_rows(
|
||||||
.bind(active_statuses)
|
sqlx::query(&sql)
|
||||||
.bind(now_unix_secs as f64)
|
.bind(active_statuses)
|
||||||
.bind(i64::try_from(limit).map_err(|_| {
|
.bind(now_unix_secs as f64)
|
||||||
DataLayerError::UnexpectedValue(format!("invalid due task limit: {limit}"))
|
.bind(i64::try_from(limit).map_err(|_| {
|
||||||
})?)
|
DataLayerError::UnexpectedValue(format!("invalid due task limit: {limit}"))
|
||||||
.fetch_all(&self.pool)
|
})?)
|
||||||
.await
|
.fetch(&self.pool),
|
||||||
.map_postgres_err()?;
|
map_video_task_row,
|
||||||
|
)
|
||||||
rows.iter().map(map_video_task_row).collect()
|
.await
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn list_page(
|
pub async fn list_page(
|
||||||
@@ -395,7 +486,7 @@ impl SqlxVideoTaskRepository {
|
|||||||
.map_err(|_| DataLayerError::UnexpectedValue(format!("invalid limit: {limit}")))?;
|
.map_err(|_| DataLayerError::UnexpectedValue(format!("invalid limit: {limit}")))?;
|
||||||
|
|
||||||
let mut builder = QueryBuilder::<Postgres>::new("SELECT\n");
|
let mut builder = QueryBuilder::<Postgres>::new("SELECT\n");
|
||||||
builder.push(SELECT_VIDEO_TASK_COLUMNS);
|
builder.push(select_video_task_full_columns());
|
||||||
builder.push("\nFROM video_tasks");
|
builder.push("\nFROM video_tasks");
|
||||||
push_video_task_filter(&mut builder, filter, None);
|
push_video_task_filter(&mut builder, filter, None);
|
||||||
builder.push("\nORDER BY created_at DESC, updated_at DESC");
|
builder.push("\nORDER BY created_at DESC, updated_at DESC");
|
||||||
@@ -404,12 +495,37 @@ impl SqlxVideoTaskRepository {
|
|||||||
builder.push("\nLIMIT ");
|
builder.push("\nLIMIT ");
|
||||||
builder.push_bind(limit);
|
builder.push_bind(limit);
|
||||||
|
|
||||||
let rows = builder
|
let query = builder.build();
|
||||||
.build()
|
collect_query_rows(query.fetch(&self.pool), map_video_task_row).await
|
||||||
.fetch_all(&self.pool)
|
}
|
||||||
.await
|
|
||||||
.map_postgres_err()?;
|
pub async fn list_page_summary(
|
||||||
rows.iter().map(map_video_task_row).collect()
|
&self,
|
||||||
|
filter: &VideoTaskQueryFilter,
|
||||||
|
offset: usize,
|
||||||
|
limit: usize,
|
||||||
|
) -> Result<Vec<StoredVideoTask>, DataLayerError> {
|
||||||
|
if limit == 0 {
|
||||||
|
return Ok(Vec::new());
|
||||||
|
}
|
||||||
|
|
||||||
|
let offset = i64::try_from(offset)
|
||||||
|
.map_err(|_| DataLayerError::UnexpectedValue(format!("invalid offset: {offset}")))?;
|
||||||
|
let limit = i64::try_from(limit)
|
||||||
|
.map_err(|_| DataLayerError::UnexpectedValue(format!("invalid limit: {limit}")))?;
|
||||||
|
|
||||||
|
let mut builder = QueryBuilder::<Postgres>::new("SELECT\n");
|
||||||
|
builder.push(select_video_task_page_summary_columns());
|
||||||
|
builder.push("\nFROM video_tasks");
|
||||||
|
push_video_task_filter(&mut builder, filter, None);
|
||||||
|
builder.push("\nORDER BY created_at DESC, updated_at DESC");
|
||||||
|
builder.push("\nOFFSET ");
|
||||||
|
builder.push_bind(offset);
|
||||||
|
builder.push("\nLIMIT ");
|
||||||
|
builder.push_bind(limit);
|
||||||
|
|
||||||
|
let query = builder.build();
|
||||||
|
collect_query_rows(query.fetch(&self.pool), map_video_task_row).await
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn count(&self, filter: &VideoTaskQueryFilter) -> Result<u64, DataLayerError> {
|
pub async fn count(&self, filter: &VideoTaskQueryFilter) -> Result<u64, DataLayerError> {
|
||||||
@@ -436,29 +552,29 @@ impl SqlxVideoTaskRepository {
|
|||||||
push_video_task_filter(&mut builder, filter, None);
|
push_video_task_filter(&mut builder, filter, None);
|
||||||
builder.push("\nGROUP BY status\nORDER BY status ASC");
|
builder.push("\nGROUP BY status\nORDER BY status ASC");
|
||||||
|
|
||||||
let rows = builder
|
let query = builder.build();
|
||||||
.build()
|
let mut rows = query.fetch(&self.pool);
|
||||||
.fetch_all(&self.pool)
|
let mut items = Vec::new();
|
||||||
.await
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||||
.map_postgres_err()?;
|
let entry = {
|
||||||
rows.into_iter()
|
|
||||||
.map(|row| {
|
|
||||||
let status = VideoTaskStatus::from_database(
|
let status = VideoTaskStatus::from_database(
|
||||||
row.try_get::<String, _>("status")
|
row.try_get::<String, _>("status")
|
||||||
.map_postgres_err()?
|
.map_postgres_err()?
|
||||||
.as_str(),
|
.as_str(),
|
||||||
)?;
|
)?;
|
||||||
let total = row.try_get::<i64, _>("total").map_postgres_err()?;
|
let total = row.try_get::<i64, _>("total").map_postgres_err()?;
|
||||||
Ok(VideoTaskStatusCount {
|
VideoTaskStatusCount {
|
||||||
status,
|
status,
|
||||||
count: u64::try_from(total).map_err(|_| {
|
count: u64::try_from(total).map_err(|_| {
|
||||||
DataLayerError::UnexpectedValue(format!(
|
DataLayerError::UnexpectedValue(format!(
|
||||||
"invalid status count result: {total}"
|
"invalid status count result: {total}"
|
||||||
))
|
))
|
||||||
})?,
|
})?,
|
||||||
})
|
}
|
||||||
})
|
};
|
||||||
.collect()
|
items.push(entry);
|
||||||
|
}
|
||||||
|
Ok(items)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn count_distinct_users(
|
pub async fn count_distinct_users(
|
||||||
@@ -504,25 +620,25 @@ impl SqlxVideoTaskRepository {
|
|||||||
builder.push("\nGROUP BY model\nORDER BY total DESC, model ASC\nLIMIT ");
|
builder.push("\nGROUP BY model\nORDER BY total DESC, model ASC\nLIMIT ");
|
||||||
builder.push_bind(limit);
|
builder.push_bind(limit);
|
||||||
|
|
||||||
let rows = builder
|
let query = builder.build();
|
||||||
.build()
|
let mut rows = query.fetch(&self.pool);
|
||||||
.fetch_all(&self.pool)
|
let mut items = Vec::new();
|
||||||
.await
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||||
.map_postgres_err()?;
|
let entry = {
|
||||||
rows.into_iter()
|
|
||||||
.map(|row| {
|
|
||||||
let model = row.try_get::<String, _>("model").map_postgres_err()?;
|
let model = row.try_get::<String, _>("model").map_postgres_err()?;
|
||||||
let total = row.try_get::<i64, _>("total").map_postgres_err()?;
|
let total = row.try_get::<i64, _>("total").map_postgres_err()?;
|
||||||
Ok(VideoTaskModelCount {
|
VideoTaskModelCount {
|
||||||
model,
|
model,
|
||||||
count: u64::try_from(total).map_err(|_| {
|
count: u64::try_from(total).map_err(|_| {
|
||||||
DataLayerError::UnexpectedValue(format!(
|
DataLayerError::UnexpectedValue(format!(
|
||||||
"invalid model count result: {total}"
|
"invalid model count result: {total}"
|
||||||
))
|
))
|
||||||
})?,
|
})?,
|
||||||
})
|
}
|
||||||
})
|
};
|
||||||
.collect()
|
items.push(entry);
|
||||||
|
}
|
||||||
|
Ok(items)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn count_created_since(
|
pub async fn count_created_since(
|
||||||
@@ -693,20 +809,17 @@ impl SqlxVideoTaskRepository {
|
|||||||
let limit = i64::try_from(limit)
|
let limit = i64::try_from(limit)
|
||||||
.map_err(|_| DataLayerError::UnexpectedValue(format!("invalid limit: {limit}")))?;
|
.map_err(|_| DataLayerError::UnexpectedValue(format!("invalid limit: {limit}")))?;
|
||||||
let sql = claim_due_sql();
|
let sql = claim_due_sql();
|
||||||
let rows = sqlx::query(&sql)
|
let mut tasks = collect_query_rows(
|
||||||
.bind(vec!["submitted", "queued", "processing"])
|
sqlx::query(&sql)
|
||||||
.bind(now_unix_secs as f64)
|
.bind(vec!["submitted", "queued", "processing"])
|
||||||
.bind(limit)
|
.bind(now_unix_secs as f64)
|
||||||
.bind(claim_until_unix_secs as f64)
|
.bind(limit)
|
||||||
.bind(now_unix_secs as f64)
|
.bind(claim_until_unix_secs as f64)
|
||||||
.fetch_all(&self.pool)
|
.bind(now_unix_secs as f64)
|
||||||
.await
|
.fetch(&self.pool),
|
||||||
.map_postgres_err()?;
|
map_video_task_row,
|
||||||
|
)
|
||||||
let mut tasks = rows
|
.await?;
|
||||||
.iter()
|
|
||||||
.map(map_video_task_row)
|
|
||||||
.collect::<Result<Vec<_>, _>>()?;
|
|
||||||
tasks.sort_by(|left, right| {
|
tasks.sort_by(|left, right| {
|
||||||
left.next_poll_at_unix_secs
|
left.next_poll_at_unix_secs
|
||||||
.cmp(&right.next_poll_at_unix_secs)
|
.cmp(&right.next_poll_at_unix_secs)
|
||||||
@@ -746,6 +859,15 @@ impl VideoTaskReadRepository for SqlxVideoTaskRepository {
|
|||||||
Self::list_page(self, filter, offset, limit).await
|
Self::list_page(self, filter, offset, limit).await
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn list_page_summary(
|
||||||
|
&self,
|
||||||
|
filter: &VideoTaskQueryFilter,
|
||||||
|
offset: usize,
|
||||||
|
limit: usize,
|
||||||
|
) -> Result<Vec<StoredVideoTask>, DataLayerError> {
|
||||||
|
Self::list_page_summary(self, filter, offset, limit).await
|
||||||
|
}
|
||||||
|
|
||||||
async fn count(&self, filter: &VideoTaskQueryFilter) -> Result<u64, DataLayerError> {
|
async fn count(&self, filter: &VideoTaskQueryFilter) -> Result<u64, DataLayerError> {
|
||||||
Self::count(self, filter).await
|
Self::count(self, filter).await
|
||||||
}
|
}
|
||||||
@@ -881,7 +1003,21 @@ fn map_status_for_database(status: VideoTaskStatus) -> &'static str {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn map_video_task_row(row: &sqlx::postgres::PgRow) -> Result<StoredVideoTask, DataLayerError> {
|
async fn collect_query_rows<T, S>(
|
||||||
|
mut rows: S,
|
||||||
|
map_row: fn(&PgRow) -> Result<T, DataLayerError>,
|
||||||
|
) -> Result<Vec<T>, DataLayerError>
|
||||||
|
where
|
||||||
|
S: TryStream<Ok = PgRow, Error = sqlx::Error> + Unpin,
|
||||||
|
{
|
||||||
|
let mut items = Vec::new();
|
||||||
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||||
|
items.push(map_row(&row)?);
|
||||||
|
}
|
||||||
|
Ok(items)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn map_video_task_row(row: &PgRow) -> Result<StoredVideoTask, DataLayerError> {
|
||||||
let status = VideoTaskStatus::from_database(
|
let status = VideoTaskStatus::from_database(
|
||||||
row.try_get::<String, _>("status")
|
row.try_get::<String, _>("status")
|
||||||
.map_postgres_err()?
|
.map_postgres_err()?
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use chrono::Utc;
|
use chrono::Utc;
|
||||||
|
use futures_util::{stream::TryStream, TryStreamExt};
|
||||||
use sqlx::{postgres::PgRow, PgPool, Row};
|
use sqlx::{postgres::PgRow, PgPool, Row};
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
@@ -24,7 +25,6 @@ use crate::{
|
|||||||
postgres::PostgresTransactionRunner,
|
postgres::PostgresTransactionRunner,
|
||||||
DataLayerError,
|
DataLayerError,
|
||||||
};
|
};
|
||||||
use std::collections::BTreeMap;
|
|
||||||
|
|
||||||
const FIND_BY_WALLET_ID_SQL: &str = r#"
|
const FIND_BY_WALLET_ID_SQL: &str = r#"
|
||||||
SELECT
|
SELECT
|
||||||
@@ -638,21 +638,13 @@ impl WalletReadRepository for SqlxWalletRepository {
|
|||||||
if user_ids.is_empty() {
|
if user_ids.is_empty() {
|
||||||
return Ok(Vec::new());
|
return Ok(Vec::new());
|
||||||
}
|
}
|
||||||
let mut ids_map = BTreeMap::new();
|
collect_query_rows(
|
||||||
for (index, id) in user_ids.iter().enumerate() {
|
sqlx::query(LIST_BY_USER_IDS_SQL)
|
||||||
ids_map.entry(id).or_insert_with(Vec::new).push(index);
|
.bind(user_ids)
|
||||||
}
|
.fetch(&self.pool),
|
||||||
let rows = sqlx::query(LIST_BY_USER_IDS_SQL)
|
map_wallet_row,
|
||||||
.bind(user_ids)
|
)
|
||||||
.fetch_all(&self.pool)
|
.await
|
||||||
.await
|
|
||||||
.map_postgres_err()?;
|
|
||||||
let mut wallets = Vec::with_capacity(rows.len());
|
|
||||||
for row in rows {
|
|
||||||
let wallet = map_wallet_row(&row)?;
|
|
||||||
wallets.push(wallet);
|
|
||||||
}
|
|
||||||
Ok(wallets)
|
|
||||||
}
|
}
|
||||||
async fn list_wallets_by_api_key_ids(
|
async fn list_wallets_by_api_key_ids(
|
||||||
&self,
|
&self,
|
||||||
@@ -661,17 +653,13 @@ impl WalletReadRepository for SqlxWalletRepository {
|
|||||||
if api_key_ids.is_empty() {
|
if api_key_ids.is_empty() {
|
||||||
return Ok(Vec::new());
|
return Ok(Vec::new());
|
||||||
}
|
}
|
||||||
let rows = sqlx::query(LIST_BY_API_KEY_IDS_SQL)
|
collect_query_rows(
|
||||||
.bind(api_key_ids)
|
sqlx::query(LIST_BY_API_KEY_IDS_SQL)
|
||||||
.fetch_all(&self.pool)
|
.bind(api_key_ids)
|
||||||
.await
|
.fetch(&self.pool),
|
||||||
.map_postgres_err()?;
|
map_wallet_row,
|
||||||
let mut wallets = Vec::with_capacity(rows.len());
|
)
|
||||||
for row in rows {
|
.await
|
||||||
let wallet = map_wallet_row(&row)?;
|
|
||||||
wallets.push(wallet);
|
|
||||||
}
|
|
||||||
Ok(wallets)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn list_admin_wallets(
|
async fn list_admin_wallets(
|
||||||
@@ -686,18 +674,16 @@ impl WalletReadRepository for SqlxWalletRepository {
|
|||||||
.await
|
.await
|
||||||
.map_postgres_err()?,
|
.map_postgres_err()?,
|
||||||
)?;
|
)?;
|
||||||
let rows = sqlx::query(LIST_ADMIN_WALLETS_SQL)
|
let items = collect_query_rows(
|
||||||
.bind(query.status.as_deref())
|
sqlx::query(LIST_ADMIN_WALLETS_SQL)
|
||||||
.bind(query.owner_type.as_deref())
|
.bind(query.status.as_deref())
|
||||||
.bind(as_i64(query.offset, "wallet offset")?)
|
.bind(query.owner_type.as_deref())
|
||||||
.bind(as_i64(query.limit, "wallet limit")?)
|
.bind(as_i64(query.offset, "wallet offset")?)
|
||||||
.fetch_all(&self.pool)
|
.bind(as_i64(query.limit, "wallet limit")?)
|
||||||
.await
|
.fetch(&self.pool),
|
||||||
.map_postgres_err()?;
|
map_admin_wallet_list_item_row,
|
||||||
let items = rows
|
)
|
||||||
.iter()
|
.await?;
|
||||||
.map(map_admin_wallet_list_item_row)
|
|
||||||
.collect::<Result<Vec<_>, _>>()?;
|
|
||||||
Ok(StoredAdminWalletListPage { items, total })
|
Ok(StoredAdminWalletListPage { items, total })
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -714,19 +700,17 @@ impl WalletReadRepository for SqlxWalletRepository {
|
|||||||
.await
|
.await
|
||||||
.map_postgres_err()?,
|
.map_postgres_err()?,
|
||||||
)?;
|
)?;
|
||||||
let rows = sqlx::query(LIST_ADMIN_WALLET_LEDGER_SQL)
|
let items = collect_query_rows(
|
||||||
.bind(query.category.as_deref())
|
sqlx::query(LIST_ADMIN_WALLET_LEDGER_SQL)
|
||||||
.bind(query.reason_code.as_deref())
|
.bind(query.category.as_deref())
|
||||||
.bind(query.owner_type.as_deref())
|
.bind(query.reason_code.as_deref())
|
||||||
.bind(as_i64(query.offset, "wallet ledger offset")?)
|
.bind(query.owner_type.as_deref())
|
||||||
.bind(as_i64(query.limit, "wallet ledger limit")?)
|
.bind(as_i64(query.offset, "wallet ledger offset")?)
|
||||||
.fetch_all(&self.pool)
|
.bind(as_i64(query.limit, "wallet ledger limit")?)
|
||||||
.await
|
.fetch(&self.pool),
|
||||||
.map_postgres_err()?;
|
map_admin_wallet_ledger_item_row,
|
||||||
let items = rows
|
)
|
||||||
.iter()
|
.await?;
|
||||||
.map(map_admin_wallet_ledger_item_row)
|
|
||||||
.collect::<Result<Vec<_>, _>>()?;
|
|
||||||
Ok(StoredAdminWalletLedgerPage { items, total })
|
Ok(StoredAdminWalletLedgerPage { items, total })
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -741,17 +725,15 @@ impl WalletReadRepository for SqlxWalletRepository {
|
|||||||
.await
|
.await
|
||||||
.map_postgres_err()?,
|
.map_postgres_err()?,
|
||||||
)?;
|
)?;
|
||||||
let rows = sqlx::query(LIST_ADMIN_WALLET_REFUND_REQUESTS_SQL)
|
let items = collect_query_rows(
|
||||||
.bind(query.status.as_deref())
|
sqlx::query(LIST_ADMIN_WALLET_REFUND_REQUESTS_SQL)
|
||||||
.bind(as_i64(query.offset, "wallet refund request offset")?)
|
.bind(query.status.as_deref())
|
||||||
.bind(as_i64(query.limit, "wallet refund request limit")?)
|
.bind(as_i64(query.offset, "wallet refund request offset")?)
|
||||||
.fetch_all(&self.pool)
|
.bind(as_i64(query.limit, "wallet refund request limit")?)
|
||||||
.await
|
.fetch(&self.pool),
|
||||||
.map_postgres_err()?;
|
map_admin_wallet_refund_request_item_row,
|
||||||
let items = rows
|
)
|
||||||
.iter()
|
.await?;
|
||||||
.map(map_admin_wallet_refund_request_item_row)
|
|
||||||
.collect::<Result<Vec<_>, _>>()?;
|
|
||||||
Ok(StoredAdminWalletRefundRequestPage { items, total })
|
Ok(StoredAdminWalletRefundRequestPage { items, total })
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -768,17 +750,15 @@ impl WalletReadRepository for SqlxWalletRepository {
|
|||||||
.await
|
.await
|
||||||
.map_postgres_err()?,
|
.map_postgres_err()?,
|
||||||
)?;
|
)?;
|
||||||
let rows = sqlx::query(LIST_ADMIN_WALLET_TRANSACTIONS_SQL)
|
let items = collect_query_rows(
|
||||||
.bind(wallet_id)
|
sqlx::query(LIST_ADMIN_WALLET_TRANSACTIONS_SQL)
|
||||||
.bind(as_i64(offset, "wallet transaction offset")?)
|
.bind(wallet_id)
|
||||||
.bind(as_i64(limit, "wallet transaction limit")?)
|
.bind(as_i64(offset, "wallet transaction offset")?)
|
||||||
.fetch_all(&self.pool)
|
.bind(as_i64(limit, "wallet transaction limit")?)
|
||||||
.await
|
.fetch(&self.pool),
|
||||||
.map_postgres_err()?;
|
map_admin_wallet_transaction_row,
|
||||||
let items = rows
|
)
|
||||||
.iter()
|
.await?;
|
||||||
.map(map_admin_wallet_transaction_row)
|
|
||||||
.collect::<Result<Vec<_>, _>>()?;
|
|
||||||
Ok(StoredAdminWalletTransactionPage { items, total })
|
Ok(StoredAdminWalletTransactionPage { items, total })
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -810,17 +790,15 @@ impl WalletReadRepository for SqlxWalletRepository {
|
|||||||
.await
|
.await
|
||||||
.map_postgres_err()?,
|
.map_postgres_err()?,
|
||||||
)?;
|
)?;
|
||||||
let rows = sqlx::query(LIST_WALLET_DAILY_USAGE_HISTORY_SQL)
|
let items = collect_query_rows(
|
||||||
.bind(wallet_id)
|
sqlx::query(LIST_WALLET_DAILY_USAGE_HISTORY_SQL)
|
||||||
.bind(billing_timezone)
|
.bind(wallet_id)
|
||||||
.bind(as_i64(limit, "wallet daily usage history limit")?)
|
.bind(billing_timezone)
|
||||||
.fetch_all(&self.pool)
|
.bind(as_i64(limit, "wallet daily usage history limit")?)
|
||||||
.await
|
.fetch(&self.pool),
|
||||||
.map_postgres_err()?;
|
map_wallet_daily_usage_row,
|
||||||
let items = rows
|
)
|
||||||
.iter()
|
.await?;
|
||||||
.map(map_wallet_daily_usage_row)
|
|
||||||
.collect::<Result<Vec<_>, _>>()?;
|
|
||||||
Ok(StoredWalletDailyUsageLedgerPage { items, total })
|
Ok(StoredWalletDailyUsageLedgerPage { items, total })
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -837,17 +815,15 @@ impl WalletReadRepository for SqlxWalletRepository {
|
|||||||
.await
|
.await
|
||||||
.map_postgres_err()?,
|
.map_postgres_err()?,
|
||||||
)?;
|
)?;
|
||||||
let rows = sqlx::query(LIST_ADMIN_WALLET_REFUNDS_SQL)
|
let items = collect_query_rows(
|
||||||
.bind(wallet_id)
|
sqlx::query(LIST_ADMIN_WALLET_REFUNDS_SQL)
|
||||||
.bind(as_i64(offset, "wallet refund offset")?)
|
.bind(wallet_id)
|
||||||
.bind(as_i64(limit, "wallet refund limit")?)
|
.bind(as_i64(offset, "wallet refund offset")?)
|
||||||
.fetch_all(&self.pool)
|
.bind(as_i64(limit, "wallet refund limit")?)
|
||||||
.await
|
.fetch(&self.pool),
|
||||||
.map_postgres_err()?;
|
map_admin_wallet_refund_row,
|
||||||
let items = rows
|
)
|
||||||
.iter()
|
.await?;
|
||||||
.map(map_admin_wallet_refund_row)
|
|
||||||
.collect::<Result<Vec<_>, _>>()?;
|
|
||||||
Ok(StoredAdminWalletRefundPage { items, total })
|
Ok(StoredAdminWalletRefundPage { items, total })
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -863,18 +839,16 @@ impl WalletReadRepository for SqlxWalletRepository {
|
|||||||
.await
|
.await
|
||||||
.map_postgres_err()?,
|
.map_postgres_err()?,
|
||||||
)?;
|
)?;
|
||||||
let rows = sqlx::query(LIST_ADMIN_PAYMENT_ORDERS_SQL)
|
let items = collect_query_rows(
|
||||||
.bind(query.payment_method.as_deref())
|
sqlx::query(LIST_ADMIN_PAYMENT_ORDERS_SQL)
|
||||||
.bind(query.status.as_deref())
|
.bind(query.payment_method.as_deref())
|
||||||
.bind(as_i64(query.offset, "payment order offset")?)
|
.bind(query.status.as_deref())
|
||||||
.bind(as_i64(query.limit, "payment order limit")?)
|
.bind(as_i64(query.offset, "payment order offset")?)
|
||||||
.fetch_all(&self.pool)
|
.bind(as_i64(query.limit, "payment order limit")?)
|
||||||
.await
|
.fetch(&self.pool),
|
||||||
.map_postgres_err()?;
|
map_admin_payment_order_row,
|
||||||
let items = rows
|
)
|
||||||
.iter()
|
.await?;
|
||||||
.map(map_admin_payment_order_row)
|
|
||||||
.collect::<Result<Vec<_>, _>>()?;
|
|
||||||
Ok(StoredAdminPaymentOrderPage { items, total })
|
Ok(StoredAdminPaymentOrderPage { items, total })
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -903,17 +877,15 @@ impl WalletReadRepository for SqlxWalletRepository {
|
|||||||
.await
|
.await
|
||||||
.map_postgres_err()?,
|
.map_postgres_err()?,
|
||||||
)?;
|
)?;
|
||||||
let rows = sqlx::query(LIST_WALLET_PAYMENT_ORDERS_BY_USER_SQL)
|
let items = collect_query_rows(
|
||||||
.bind(user_id)
|
sqlx::query(LIST_WALLET_PAYMENT_ORDERS_BY_USER_SQL)
|
||||||
.bind(as_i64(offset, "wallet payment order offset")?)
|
.bind(user_id)
|
||||||
.bind(as_i64(limit, "wallet payment order limit")?)
|
.bind(as_i64(offset, "wallet payment order offset")?)
|
||||||
.fetch_all(&self.pool)
|
.bind(as_i64(limit, "wallet payment order limit")?)
|
||||||
.await
|
.fetch(&self.pool),
|
||||||
.map_postgres_err()?;
|
map_admin_payment_order_row,
|
||||||
let items = rows
|
)
|
||||||
.iter()
|
.await?;
|
||||||
.map(map_admin_payment_order_row)
|
|
||||||
.collect::<Result<Vec<_>, _>>()?;
|
|
||||||
Ok(StoredAdminPaymentOrderPage { items, total })
|
Ok(StoredAdminPaymentOrderPage { items, total })
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -958,17 +930,15 @@ impl WalletReadRepository for SqlxWalletRepository {
|
|||||||
.await
|
.await
|
||||||
.map_postgres_err()?,
|
.map_postgres_err()?,
|
||||||
)?;
|
)?;
|
||||||
let rows = sqlx::query(LIST_ADMIN_PAYMENT_CALLBACKS_SQL)
|
let items = collect_query_rows(
|
||||||
.bind(payment_method)
|
sqlx::query(LIST_ADMIN_PAYMENT_CALLBACKS_SQL)
|
||||||
.bind(as_i64(offset, "payment callback offset")?)
|
.bind(payment_method)
|
||||||
.bind(as_i64(limit, "payment callback limit")?)
|
.bind(as_i64(offset, "payment callback offset")?)
|
||||||
.fetch_all(&self.pool)
|
.bind(as_i64(limit, "payment callback limit")?)
|
||||||
.await
|
.fetch(&self.pool),
|
||||||
.map_postgres_err()?;
|
map_admin_payment_callback_row,
|
||||||
let items = rows
|
)
|
||||||
.iter()
|
.await?;
|
||||||
.map(map_admin_payment_callback_row)
|
|
||||||
.collect::<Result<Vec<_>, _>>()?;
|
|
||||||
Ok(StoredAdminPaymentCallbackPage { items, total })
|
Ok(StoredAdminPaymentCallbackPage { items, total })
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -3556,6 +3526,20 @@ fn read_count(row: PgRow) -> Result<u64, DataLayerError> {
|
|||||||
Ok(total.max(0) as u64)
|
Ok(total.max(0) as u64)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn collect_query_rows<T, S>(
|
||||||
|
mut rows: S,
|
||||||
|
map_row: fn(&PgRow) -> Result<T, DataLayerError>,
|
||||||
|
) -> Result<Vec<T>, DataLayerError>
|
||||||
|
where
|
||||||
|
S: TryStream<Ok = PgRow, Error = sqlx::Error> + Unpin,
|
||||||
|
{
|
||||||
|
let mut items = Vec::new();
|
||||||
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||||
|
items.push(map_row(&row)?);
|
||||||
|
}
|
||||||
|
Ok(items)
|
||||||
|
}
|
||||||
|
|
||||||
fn map_admin_wallet_list_item_row(
|
fn map_admin_wallet_list_item_row(
|
||||||
row: &PgRow,
|
row: &PgRow,
|
||||||
) -> Result<StoredAdminWalletListItem, DataLayerError> {
|
) -> Result<StoredAdminWalletListItem, DataLayerError> {
|
||||||
|
|||||||
Reference in New Issue
Block a user