refactor(runtime): 优化管理端摘要查询与维护聚合链路

- 为 provider catalog key 和 video task 列表增加 summary/page 查询与排序能力,减少列表场景读取重字段
- 将多处 SQL 结果读取改为流式收集,降低 `fetch_all` 的内存占用
- 把日/小时统计、钱包日用量等维护任务改为数据库侧 `CTE + upsert` 聚合
- 修复视频任务轮询更新时从本地 snapshot 回填稀疏字段,避免 `prompt` 和请求体信息丢失
This commit is contained in:
AAEE86
2026-04-12 20:45:47 +08:00
committed by fawney19
parent 6aa16ec792
commit b2d85d70ca
55 changed files with 1954 additions and 1382 deletions

View File

@@ -4,9 +4,9 @@ use std::sync::RwLock;
use async_trait::async_trait;
use super::{
ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyPage,
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogReadRepository,
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
};
use crate::DataLayerError;
@@ -146,6 +146,13 @@ impl ProviderCatalogReadRepository for InMemoryProviderCatalogReadRepository {
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(
&self,
query: &ProviderCatalogKeyListQuery,
@@ -170,12 +177,28 @@ impl ProviderCatalogReadRepository for InMemoryProviderCatalogReadRepository {
})
.cloned()
.collect::<Vec<_>>();
keys.sort_by(|left, right| {
left.internal_priority
.cmp(&right.internal_priority)
.then(left.name.cmp(&right.name))
.then(left.id.cmp(&right.id))
});
match query.order {
ProviderCatalogKeyListOrder::Name => {
keys.sort_by(|left, right| {
left.internal_priority
.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 items = keys
.into_iter()
@@ -415,8 +438,9 @@ impl ProviderCatalogWriteRepository for InMemoryProviderCatalogReadRepository {
mod tests {
use super::InMemoryProviderCatalogReadRepository;
use crate::repository::provider_catalog::{
ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogReadRepository,
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogProvider,
};
fn sample_provider(id: &str) -> StoredProviderCatalogProvider {
@@ -583,6 +607,7 @@ mod tests {
is_active: Some(true),
offset: 0,
limit: 10,
order: ProviderCatalogKeyListOrder::Name,
})
.await
.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]
async fn summarizes_provider_key_stats() {
let mut inactive = sample_key("key-2", "provider-1");

View File

@@ -3,9 +3,9 @@ mod sql;
#[allow(unused_imports)]
pub(crate) use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyPage,
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogReadRepository,
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
};
pub use memory::InMemoryProviderCatalogReadRepository;
pub use sql::SqlxProviderCatalogReadRepository;

View File

@@ -1,10 +1,11 @@
use async_trait::async_trait;
use futures_util::TryStreamExt;
use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row};
use super::{
ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyPage,
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogReadRepository,
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
};
use crate::{
error::{postgres_error, SqlxResultExt},
@@ -240,6 +241,61 @@ FROM provider_api_keys
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#"
SELECT
provider_id,
@@ -271,24 +327,26 @@ impl SqlxProviderCatalogReadRepository {
return Ok(Vec::new());
}
let rows = build_list_query(
LIST_PROVIDERS_BY_IDS_PREFIX,
provider_ids,
" ORDER BY name ASC",
collect_query_rows(
build_list_query(
LIST_PROVIDERS_BY_IDS_PREFIX,
provider_ids,
" ORDER BY name ASC",
)
.build()
.fetch(&self.pool),
map_provider_row,
)
.build()
.fetch_all(&self.pool)
.await
.map_postgres_err()?;
rows.iter().map(map_provider_row).collect()
}
pub async fn list_providers(
&self,
active_only: bool,
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
let rows = sqlx::query(
r#"
collect_query_rows(
sqlx::query(
r#"
SELECT
id,
name,
@@ -317,12 +375,12 @@ FROM providers
WHERE ($1::boolean = false OR is_active = true)
ORDER BY provider_priority ASC, name ASC
"#,
)
.bind(active_only)
.fetch(&self.pool),
map_provider_row,
)
.bind(active_only)
.fetch_all(&self.pool)
.await
.map_postgres_err()?;
rows.iter().map(map_provider_row).collect()
}
pub async fn list_endpoints_by_ids(
@@ -333,28 +391,35 @@ ORDER BY provider_priority ASC, name ASC
return Ok(Vec::new());
}
let rows = match build_list_query(
LIST_ENDPOINTS_BY_IDS_PREFIX,
endpoint_ids,
" 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,
let rows = match collect_query_rows(
build_list_query(
LIST_ENDPOINTS_BY_IDS_PREFIX,
endpoint_ids,
" ORDER BY api_format ASC, id ASC",
)
.build()
.fetch_all(&self.pool)
.await
.map_postgres_err()?,
.fetch(&self.pool),
map_endpoint_row,
)
.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)),
};
rows.iter().map(map_endpoint_row).collect()
Ok(rows)
}
pub async fn list_endpoints_by_provider_ids(
@@ -365,28 +430,35 @@ ORDER BY provider_priority ASC, name ASC
return Ok(Vec::new());
}
let rows = match build_list_query(
LIST_ENDPOINTS_BY_PROVIDER_IDS_PREFIX,
provider_ids,
" 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,
let rows = match collect_query_rows(
build_list_query(
LIST_ENDPOINTS_BY_PROVIDER_IDS_PREFIX,
provider_ids,
" ORDER BY provider_id ASC, api_format ASC, id ASC",
)
.build()
.fetch_all(&self.pool)
.await
.map_postgres_err()?,
.fetch(&self.pool),
map_endpoint_row,
)
.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)),
};
rows.iter().map(map_endpoint_row).collect()
Ok(rows)
}
pub async fn list_keys_by_ids(
@@ -397,16 +469,17 @@ ORDER BY provider_priority ASC, name ASC
return Ok(Vec::new());
}
let rows = build_list_query(
LIST_KEYS_BY_IDS_PREFIX,
key_ids,
" ORDER BY name ASC, id ASC",
collect_query_rows(
build_list_query(
LIST_KEYS_BY_IDS_PREFIX,
key_ids,
" ORDER BY name ASC, id ASC",
)
.build()
.fetch(&self.pool),
map_key_row,
)
.build()
.fetch_all(&self.pool)
.await
.map_postgres_err()?;
rows.iter().map(map_key_row).collect()
}
pub async fn list_keys_by_provider_ids(
@@ -417,16 +490,38 @@ ORDER BY provider_priority ASC, name ASC
return Ok(Vec::new());
}
let rows = build_list_query(
LIST_KEYS_BY_PROVIDER_IDS_PREFIX,
provider_ids,
" ORDER BY provider_id ASC, name ASC, id ASC",
collect_query_rows(
build_list_query(
LIST_KEYS_BY_PROVIDER_IDS_PREFIX,
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
.map_postgres_err()?;
rows.iter().map(map_key_row).collect()
}
pub async fn list_keys_page(
@@ -457,6 +552,12 @@ ORDER BY provider_priority ASC, name ASC
.map(str::trim)
.filter(|value| !value.is_empty())
.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(
r#"
@@ -475,7 +576,7 @@ WHERE provider_id = $1
.map_postgres_err()?;
let total = row_get::<i64>(&count_row, "total")?.max(0) as usize;
let rows = sqlx::query(
let sql = format!(
r#"
SELECT
id,
@@ -531,23 +632,22 @@ FROM provider_api_keys
WHERE provider_id = $1
AND ($2::TEXT IS NULL OR LOWER(name) LIKE $2 OR LOWER(id) LIKE $2)
AND ($3::BOOLEAN IS NULL OR is_active = $3)
ORDER BY internal_priority ASC, name ASC, id ASC
ORDER BY {order_by}
OFFSET $4
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)
.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<_>, _>>()?;
.await?;
Ok(StoredProviderCatalogKeyPage { items, total })
}
@@ -560,16 +660,17 @@ LIMIT $5
return Ok(Vec::new());
}
let rows = build_list_query(
LIST_KEY_STATS_BY_PROVIDER_IDS_PREFIX,
provider_ids,
"\nGROUP BY provider_id\nORDER BY provider_id ASC",
collect_query_rows(
build_list_query(
LIST_KEY_STATS_BY_PROVIDER_IDS_PREFIX,
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
.map_postgres_err()?;
rows.iter().map(map_key_stats_row).collect()
}
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
}
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(
&self,
query: &ProviderCatalogKeyListQuery,
@@ -2051,6 +2159,29 @@ fn is_missing_endpoint_health_score_column(error: &sqlx::Error) -> bool {
.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> {
StoredProviderCatalogKeyStats::new(
row_get(row, "provider_id")?,