Merge PR #672: 支持账号批量配置与可用模型管理

This commit is contained in:
elky
2026-07-18 22:06:32 +08:00
37 changed files with 2612 additions and 377 deletions
@@ -1,6 +1,10 @@
use async_trait::async_trait;
use futures_util::TryStreamExt;
use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row};
use sqlx::{
postgres::{PgArguments, PgRow},
query::Query,
PgPool, Postgres, QueryBuilder, Row,
};
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogReadRepository,
@@ -334,6 +338,136 @@ FROM provider_api_keys
WHERE provider_id IN (
"#;
const KEY_UPDATE_SQL: &str = r#"
UPDATE provider_api_keys
SET
provider_id = $2,
api_formats = $3,
auth_type_by_format = $40,
allow_auth_channel_mismatch_formats = $41,
auth_type = $4,
api_key = $5,
auth_config = $6,
name = $7,
note = $8,
rate_multipliers = $9,
internal_priority = $10,
global_priority_by_format = $11,
rpm_limit = $12,
concurrent_limit = $13,
learned_rpm_limit = $14,
allowed_models = $15,
capabilities = $16,
cache_ttl_minutes = $17,
max_probe_interval_minutes = $18,
auto_fetch_models = $19,
locked_models = $20,
model_include_patterns = $21,
model_exclude_patterns = $22,
proxy = $23,
fingerprint = $24,
upstream_metadata = $25,
expires_at = CASE
WHEN $39::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($39::double precision)
END,
oauth_invalid_at = CASE
WHEN $26::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($26::double precision)
END,
oauth_invalid_reason = $27,
status_snapshot = $28,
concurrent_429_count = COALESCE($29, 0),
rpm_429_count = COALESCE($30, 0),
last_429_at = CASE
WHEN $31::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($31::double precision)
END,
last_429_type = $32,
adjustment_history = $33,
utilization_samples = $34,
last_probe_increase_at = CASE
WHEN $35::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($35::double precision)
END,
last_rpm_peak = $36,
is_active = $37,
updated_at = CASE
WHEN $38::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($38::double precision)
END,
last_models_fetch_at = CASE
WHEN $42::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($42::double precision)
END,
last_models_fetch_error = $43
WHERE id = $1
"#;
fn validate_key_for_update(key: &StoredProviderCatalogKey) -> Result<(), DataLayerError> {
if key.id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog key.id is empty".to_string(),
));
}
if key.provider_id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog key.provider_id is empty".to_string(),
));
}
Ok(())
}
fn key_update_query(key: &StoredProviderCatalogKey) -> Query<'_, Postgres, PgArguments> {
sqlx::query(KEY_UPDATE_SQL)
.bind(&key.id)
.bind(&key.provider_id)
.bind(&key.api_formats)
.bind(&key.auth_type)
.bind(&key.encrypted_api_key)
.bind(&key.encrypted_auth_config)
.bind(&key.name)
.bind(&key.note)
.bind(&key.rate_multipliers)
.bind(key.internal_priority)
.bind(&key.global_priority_by_format)
.bind(key.rpm_limit.map(|value| value as i32))
.bind(key.concurrent_limit)
.bind(key.learned_rpm_limit.map(|value| value as i32))
.bind(&key.allowed_models)
.bind(&key.capabilities)
.bind(key.cache_ttl_minutes)
.bind(key.max_probe_interval_minutes)
.bind(key.auto_fetch_models)
.bind(&key.locked_models)
.bind(&key.model_include_patterns)
.bind(&key.model_exclude_patterns)
.bind(&key.proxy)
.bind(&key.fingerprint)
.bind(&key.upstream_metadata)
.bind(key.oauth_invalid_at_unix_secs.map(|value| value as f64))
.bind(&key.oauth_invalid_reason)
.bind(&key.status_snapshot)
.bind(key.concurrent_429_count.map(|value| value as i32))
.bind(key.rpm_429_count.map(|value| value as i32))
.bind(key.last_429_at_unix_secs.map(|value| value as f64))
.bind(&key.last_429_type)
.bind(&key.adjustment_history)
.bind(&key.utilization_samples)
.bind(
key.last_probe_increase_at_unix_secs
.map(|value| value as f64),
)
.bind(key.last_rpm_peak.map(|value| value as i32))
.bind(key.is_active)
.bind(key.updated_at_unix_secs.map(|value| value as f64))
.bind(key.expires_at_unix_secs.map(|value| value as f64))
.bind(&key.auth_type_by_format)
.bind(&key.allow_auth_channel_mismatch_formats)
.bind(key.last_models_fetch_at_unix_secs.map(|value| value as f64))
.bind(&key.last_models_fetch_error)
}
#[derive(Debug, Clone)]
pub struct SqlxProviderCatalogReadRepository {
pool: PgPool,
@@ -1652,134 +1786,12 @@ WHERE id = $1
&self,
key: &StoredProviderCatalogKey,
) -> Result<StoredProviderCatalogKey, DataLayerError> {
if key.id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog key.id is empty".to_string(),
));
}
if key.provider_id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog key.provider_id is empty".to_string(),
));
}
let rows_affected = sqlx::query(
r#"
UPDATE provider_api_keys
SET
provider_id = $2,
api_formats = $3,
auth_type_by_format = $40,
allow_auth_channel_mismatch_formats = $41,
auth_type = $4,
api_key = $5,
auth_config = $6,
name = $7,
note = $8,
rate_multipliers = $9,
internal_priority = $10,
global_priority_by_format = $11,
rpm_limit = $12,
concurrent_limit = $13,
learned_rpm_limit = $14,
allowed_models = $15,
capabilities = $16,
cache_ttl_minutes = $17,
max_probe_interval_minutes = $18,
auto_fetch_models = $19,
locked_models = $20,
model_include_patterns = $21,
model_exclude_patterns = $22,
proxy = $23,
fingerprint = $24,
upstream_metadata = $25,
expires_at = CASE
WHEN $39::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($39::double precision)
END,
oauth_invalid_at = CASE
WHEN $26::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($26::double precision)
END,
oauth_invalid_reason = $27,
status_snapshot = $28,
concurrent_429_count = COALESCE($29, 0),
rpm_429_count = COALESCE($30, 0),
last_429_at = CASE
WHEN $31::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($31::double precision)
END,
last_429_type = $32,
adjustment_history = $33,
utilization_samples = $34,
last_probe_increase_at = CASE
WHEN $35::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($35::double precision)
END,
last_rpm_peak = $36,
is_active = $37,
updated_at = CASE
WHEN $38::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($38::double precision)
END,
last_models_fetch_at = CASE
WHEN $42::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($42::double precision)
END,
last_models_fetch_error = $43
WHERE id = $1
"#,
)
.bind(&key.id)
.bind(&key.provider_id)
.bind(&key.api_formats)
.bind(&key.auth_type)
.bind(&key.encrypted_api_key)
.bind(&key.encrypted_auth_config)
.bind(&key.name)
.bind(&key.note)
.bind(&key.rate_multipliers)
.bind(key.internal_priority)
.bind(&key.global_priority_by_format)
.bind(key.rpm_limit.map(|value| value as i32))
.bind(key.concurrent_limit)
.bind(key.learned_rpm_limit.map(|value| value as i32))
.bind(&key.allowed_models)
.bind(&key.capabilities)
.bind(key.cache_ttl_minutes)
.bind(key.max_probe_interval_minutes)
.bind(key.auto_fetch_models)
.bind(&key.locked_models)
.bind(&key.model_include_patterns)
.bind(&key.model_exclude_patterns)
.bind(&key.proxy)
.bind(&key.fingerprint)
.bind(&key.upstream_metadata)
.bind(key.oauth_invalid_at_unix_secs.map(|value| value as f64))
.bind(&key.oauth_invalid_reason)
.bind(&key.status_snapshot)
.bind(key.concurrent_429_count.map(|value| value as i32))
.bind(key.rpm_429_count.map(|value| value as i32))
.bind(key.last_429_at_unix_secs.map(|value| value as f64))
.bind(&key.last_429_type)
.bind(&key.adjustment_history)
.bind(&key.utilization_samples)
.bind(
key.last_probe_increase_at_unix_secs
.map(|value| value as f64),
)
.bind(key.last_rpm_peak.map(|value| value as i32))
.bind(key.is_active)
.bind(key.updated_at_unix_secs.map(|value| value as f64))
.bind(key.expires_at_unix_secs.map(|value| value as f64))
.bind(&key.auth_type_by_format)
.bind(&key.allow_auth_channel_mismatch_formats)
.bind(key.last_models_fetch_at_unix_secs.map(|value| value as f64))
.bind(&key.last_models_fetch_error)
.execute(&self.pool)
.await
.map_postgres_err()?
.rows_affected();
validate_key_for_update(key)?;
let rows_affected = key_update_query(key)
.execute(&self.pool)
.await
.map_postgres_err()?
.rows_affected();
if rows_affected == 0 {
return Err(DataLayerError::UnexpectedValue(format!(
@@ -1800,6 +1812,35 @@ WHERE id = $1
})
}
pub async fn update_keys(
&self,
keys: &[StoredProviderCatalogKey],
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
if keys.is_empty() {
return Ok(Vec::new());
}
for key in keys {
validate_key_for_update(key)?;
}
let mut transaction = self.pool.begin().await.map_postgres_err()?;
for key in keys {
let rows_affected = key_update_query(key)
.execute(&mut *transaction)
.await
.map_postgres_err()?
.rows_affected();
if rows_affected == 0 {
return Err(DataLayerError::UnexpectedValue(format!(
"provider catalog key {} not found",
key.id
)));
}
}
transaction.commit().await.map_postgres_err()?;
Ok(keys.to_vec())
}
pub async fn delete_key(&self, key_id: &str) -> Result<bool, DataLayerError> {
if key_id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
@@ -2178,6 +2219,13 @@ impl ProviderCatalogWriteRepository for SqlxProviderCatalogReadRepository {
Self::update_key(self, key).await
}
async fn update_keys(
&self,
keys: &[StoredProviderCatalogKey],
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
Self::update_keys(self, keys).await
}
async fn update_key_upstream_metadata(
&self,
key_id: &str,