i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-10 05:00:19 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
elky
2026-07-18 22:06:32 +08:00
parent 8fbda84acb 3b0d87b0fd
commit 95053f9502
37 changed files with 2612 additions and 377 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-data/adapters/mysql/src/provider_catalog.rs
+148 -82
View File
@@ -1,5 +1,9 @@
use async_trait::async_trait;
use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row};
use sqlx::{
mysql::{MySqlArguments, MySqlRow},
query::Query,
MySql, QueryBuilder, Row,
};
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, ProviderCatalogSnapshot,
@@ -669,87 +673,7 @@ WHERE id = ?
) -> Result<StoredProviderCatalogKey, DataLayerError> {
validate_key(key)?;
let updated_at = key.updated_at_unix_secs.unwrap_or_else(current_unix_secs) as i64;
let rows_affected = sqlx::query(key_update_sql())
.bind(&key.provider_id)
.bind(&key.name)
.bind(&key.encrypted_api_key)
.bind(&key.auth_type)
.bind(optional_json_to_string(
&key.capabilities,
"provider_api_keys.capabilities",
)?)
.bind(key.is_active)
.bind(optional_json_to_string(
&key.api_formats,
"provider_api_keys.api_formats",
)?)
.bind(optional_json_to_string(
&key.auth_type_by_format,
"provider_api_keys.auth_type_by_format",
)?)
.bind(optional_json_to_string(
&key.allow_auth_channel_mismatch_formats,
"provider_api_keys.allow_auth_channel_mismatch_formats",
)?)
.bind(&key.encrypted_auth_config)
.bind(&key.note)
.bind(key.internal_priority)
.bind(optional_json_to_string(
&key.rate_multipliers,
"provider_api_keys.rate_multipliers",
)?)
.bind(optional_json_to_string(
&key.global_priority_by_format,
"provider_api_keys.global_priority_by_format",
)?)
.bind(optional_json_to_string(
&key.allowed_models,
"provider_api_keys.allowed_models",
)?)
.bind(optional_i64_from_u64(
key.expires_at_unix_secs,
"provider_api_keys.expires_at",
)?)
.bind(key.cache_ttl_minutes)
.bind(key.max_probe_interval_minutes)
.bind(optional_json_to_string(
&key.proxy,
"provider_api_keys.proxy",
)?)
.bind(optional_json_to_string(
&key.fingerprint,
"provider_api_keys.fingerprint",
)?)
.bind(optional_i64_from_u32(key.rpm_limit))
.bind(key.concurrent_limit)
.bind(optional_i64_from_u32(key.learned_rpm_limit))
.bind(optional_i64_from_u32(key.concurrent_429_count).unwrap_or(0))
.bind(optional_i64_from_u32(key.rpm_429_count).unwrap_or(0))
.bind(optional_i64_from_u64(
key.last_429_at_unix_secs,
"provider_api_keys.last_429_at",
)?)
.bind(&key.last_429_type)
.bind(optional_json_to_string(
&key.adjustment_history,
"provider_api_keys.adjustment_history",
)?)
.bind(optional_json_to_string(
&key.utilization_samples,
"provider_api_keys.utilization_samples",
)?)
.bind(optional_i64_from_u64(
key.last_probe_increase_at_unix_secs,
"provider_api_keys.last_probe_increase_at",
)?)
.bind(optional_i64_from_u32(key.last_rpm_peak))
.bind(optional_i64_from_u64(
key.last_models_fetch_at_unix_secs,
"provider_api_keys.last_models_fetch_at",
)?)
.bind(&key.last_models_fetch_error)
.bind(updated_at)
.bind(&key.id)
let rows_affected = key_update_query(key, updated_at)?
.execute(&self.pool)
.await
.map_sql_err()?
@@ -764,6 +688,37 @@ WHERE id = ?
self.reload_key(&key.id, "updated").await
}
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(key)?;
}
let updated_at = current_unix_secs() as i64;
let mut transaction = self.pool.begin().await.map_sql_err()?;
for key in keys {
let key_updated_at = key.updated_at_unix_secs.unwrap_or(updated_at as u64) as i64;
let rows_affected = key_update_query(key, key_updated_at)?
.execute(&mut *transaction)
.await
.map_sql_err()?
.rows_affected();
if rows_affected == 0 {
return Err(DataLayerError::UnexpectedValue(format!(
"provider catalog key {} not found",
key.id
)));
}
}
transaction.commit().await.map_sql_err()?;
Ok(keys.to_vec())
}
pub async fn delete_key(&self, key_id: &str) -> Result<bool, DataLayerError> {
validate_non_empty(key_id, "provider catalog key_id")?;
let rows_affected = sqlx::query("DELETE FROM provider_api_keys WHERE id = ?")
@@ -1251,6 +1206,13 @@ impl ProviderCatalogWriteRepository for MysqlProviderCatalogReadRepository {
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,
@@ -1526,11 +1488,115 @@ SET
last_rpm_peak = ?,
last_models_fetch_at = ?,
last_models_fetch_error = ?,
auto_fetch_models = ?,
locked_models = ?,
model_include_patterns = ?,
model_exclude_patterns = ?,
updated_at = ?
WHERE id = ?
"#
}
fn key_update_query(
key: &StoredProviderCatalogKey,
updated_at: i64,
) -> Result<Query<'_, MySql, MySqlArguments>, DataLayerError> {
Ok(sqlx::query(key_update_sql())
.bind(&key.provider_id)
.bind(&key.name)
.bind(&key.encrypted_api_key)
.bind(&key.auth_type)
.bind(optional_json_to_string(
&key.capabilities,
"provider_api_keys.capabilities",
)?)
.bind(key.is_active)
.bind(optional_json_to_string(
&key.api_formats,
"provider_api_keys.api_formats",
)?)
.bind(optional_json_to_string(
&key.auth_type_by_format,
"provider_api_keys.auth_type_by_format",
)?)
.bind(optional_json_to_string(
&key.allow_auth_channel_mismatch_formats,
"provider_api_keys.allow_auth_channel_mismatch_formats",
)?)
.bind(&key.encrypted_auth_config)
.bind(&key.note)
.bind(key.internal_priority)
.bind(optional_json_to_string(
&key.rate_multipliers,
"provider_api_keys.rate_multipliers",
)?)
.bind(optional_json_to_string(
&key.global_priority_by_format,
"provider_api_keys.global_priority_by_format",
)?)
.bind(optional_json_to_string(
&key.allowed_models,
"provider_api_keys.allowed_models",
)?)
.bind(optional_i64_from_u64(
key.expires_at_unix_secs,
"provider_api_keys.expires_at",
)?)
.bind(key.cache_ttl_minutes)
.bind(key.max_probe_interval_minutes)
.bind(optional_json_to_string(
&key.proxy,
"provider_api_keys.proxy",
)?)
.bind(optional_json_to_string(
&key.fingerprint,
"provider_api_keys.fingerprint",
)?)
.bind(optional_i64_from_u32(key.rpm_limit))
.bind(key.concurrent_limit)
.bind(optional_i64_from_u32(key.learned_rpm_limit))
.bind(optional_i64_from_u32(key.concurrent_429_count).unwrap_or(0))
.bind(optional_i64_from_u32(key.rpm_429_count).unwrap_or(0))
.bind(optional_i64_from_u64(
key.last_429_at_unix_secs,
"provider_api_keys.last_429_at",
)?)
.bind(&key.last_429_type)
.bind(optional_json_to_string(
&key.adjustment_history,
"provider_api_keys.adjustment_history",
)?)
.bind(optional_json_to_string(
&key.utilization_samples,
"provider_api_keys.utilization_samples",
)?)
.bind(optional_i64_from_u64(
key.last_probe_increase_at_unix_secs,
"provider_api_keys.last_probe_increase_at",
)?)
.bind(optional_i64_from_u32(key.last_rpm_peak))
.bind(optional_i64_from_u64(
key.last_models_fetch_at_unix_secs,
"provider_api_keys.last_models_fetch_at",
)?)
.bind(&key.last_models_fetch_error)
.bind(key.auto_fetch_models)
.bind(optional_json_to_string(
&key.locked_models,
"provider_api_keys.locked_models",
)?)
.bind(optional_json_to_string(
&key.model_include_patterns,
"provider_api_keys.model_include_patterns",
)?)
.bind(optional_json_to_string(
&key.model_exclude_patterns,
"provider_api_keys.model_exclude_patterns",
)?)
.bind(updated_at)
.bind(&key.id))
}
fn optional_json_from_string(
value: Option<String>,
field_name: &str,
crates/aether-data/adapters/postgres/src/provider_catalog.rs
+177 -129
View File
@@ -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,
crates/aether-data/adapters/sqlite/src/provider_catalog.rs
+208 -82
View File
@@ -1,5 +1,9 @@
use async_trait::async_trait;
use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite};
use sqlx::{
query::Query,
sqlite::{SqliteArguments, SqliteRow},
QueryBuilder, Row, Sqlite,
};
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogReadRepository,
@@ -1093,87 +1097,7 @@ WHERE id = ?
) -> Result<StoredProviderCatalogKey, DataLayerError> {
validate_key(key)?;
let updated_at = key.updated_at_unix_secs.unwrap_or_else(current_unix_secs) as i64;
let rows_affected = sqlx::query(key_update_sql())
.bind(&key.provider_id)
.bind(&key.name)
.bind(&key.encrypted_api_key)
.bind(&key.auth_type)
.bind(optional_json_to_string(
&key.capabilities,
"provider_api_keys.capabilities",
)?)
.bind(key.is_active)
.bind(optional_json_to_string(
&key.api_formats,
"provider_api_keys.api_formats",
)?)
.bind(optional_json_to_string(
&key.auth_type_by_format,
"provider_api_keys.auth_type_by_format",
)?)
.bind(optional_json_to_string(
&key.allow_auth_channel_mismatch_formats,
"provider_api_keys.allow_auth_channel_mismatch_formats",
)?)
.bind(&key.encrypted_auth_config)
.bind(&key.note)
.bind(key.internal_priority)
.bind(optional_json_to_string(
&key.rate_multipliers,
"provider_api_keys.rate_multipliers",
)?)
.bind(optional_json_to_string(
&key.global_priority_by_format,
"provider_api_keys.global_priority_by_format",
)?)
.bind(optional_json_to_string(
&key.allowed_models,
"provider_api_keys.allowed_models",
)?)
.bind(optional_i64_from_u64(
key.expires_at_unix_secs,
"provider_api_keys.expires_at",
)?)
.bind(key.cache_ttl_minutes)
.bind(key.max_probe_interval_minutes)
.bind(optional_json_to_string(
&key.proxy,
"provider_api_keys.proxy",
)?)
.bind(optional_json_to_string(
&key.fingerprint,
"provider_api_keys.fingerprint",
)?)
.bind(optional_i64_from_u32(key.rpm_limit))
.bind(key.concurrent_limit)
.bind(optional_i64_from_u32(key.learned_rpm_limit))
.bind(optional_i64_from_u32(key.concurrent_429_count).unwrap_or(0))
.bind(optional_i64_from_u32(key.rpm_429_count).unwrap_or(0))
.bind(optional_i64_from_u64(
key.last_429_at_unix_secs,
"provider_api_keys.last_429_at",
)?)
.bind(&key.last_429_type)
.bind(optional_json_to_string(
&key.adjustment_history,
"provider_api_keys.adjustment_history",
)?)
.bind(optional_json_to_string(
&key.utilization_samples,
"provider_api_keys.utilization_samples",
)?)
.bind(optional_i64_from_u64(
key.last_probe_increase_at_unix_secs,
"provider_api_keys.last_probe_increase_at",
)?)
.bind(optional_i64_from_u32(key.last_rpm_peak))
.bind(optional_i64_from_u64(
key.last_models_fetch_at_unix_secs,
"provider_api_keys.last_models_fetch_at",
)?)
.bind(&key.last_models_fetch_error)
.bind(updated_at)
.bind(&key.id)
let rows_affected = key_update_query(key, updated_at)?
.execute(&self.pool)
.await
.map_sql_err()?
@@ -1188,6 +1112,37 @@ WHERE id = ?
self.reload_key(&key.id, "updated").await
}
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(key)?;
}
let updated_at = current_unix_secs() as i64;
let mut transaction = self.pool.begin().await.map_sql_err()?;
for key in keys {
let key_updated_at = key.updated_at_unix_secs.unwrap_or(updated_at as u64) as i64;
let rows_affected = key_update_query(key, key_updated_at)?
.execute(&mut *transaction)
.await
.map_sql_err()?
.rows_affected();
if rows_affected == 0 {
return Err(DataLayerError::UnexpectedValue(format!(
"provider catalog key {} not found",
key.id
)));
}
}
transaction.commit().await.map_sql_err()?;
Ok(keys.to_vec())
}
pub async fn delete_key(&self, key_id: &str) -> Result<bool, DataLayerError> {
validate_non_empty(key_id, "provider catalog key_id")?;
let rows_affected = sqlx::query("DELETE FROM provider_api_keys WHERE id = ?")
@@ -1661,6 +1616,13 @@ impl ProviderCatalogWriteRepository for SqliteProviderCatalogReadRepository {
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,
@@ -1978,11 +1940,115 @@ SET
last_rpm_peak = ?,
last_models_fetch_at = ?,
last_models_fetch_error = ?,
auto_fetch_models = ?,
locked_models = ?,
model_include_patterns = ?,
model_exclude_patterns = ?,
updated_at = ?
WHERE id = ?
"#
}
fn key_update_query(
key: &StoredProviderCatalogKey,
updated_at: i64,
) -> Result<Query<'_, Sqlite, SqliteArguments<'_>>, DataLayerError> {
Ok(sqlx::query(key_update_sql())
.bind(&key.provider_id)
.bind(&key.name)
.bind(&key.encrypted_api_key)
.bind(&key.auth_type)
.bind(optional_json_to_string(
&key.capabilities,
"provider_api_keys.capabilities",
)?)
.bind(key.is_active)
.bind(optional_json_to_string(
&key.api_formats,
"provider_api_keys.api_formats",
)?)
.bind(optional_json_to_string(
&key.auth_type_by_format,
"provider_api_keys.auth_type_by_format",
)?)
.bind(optional_json_to_string(
&key.allow_auth_channel_mismatch_formats,
"provider_api_keys.allow_auth_channel_mismatch_formats",
)?)
.bind(&key.encrypted_auth_config)
.bind(&key.note)
.bind(key.internal_priority)
.bind(optional_json_to_string(
&key.rate_multipliers,
"provider_api_keys.rate_multipliers",
)?)
.bind(optional_json_to_string(
&key.global_priority_by_format,
"provider_api_keys.global_priority_by_format",
)?)
.bind(optional_json_to_string(
&key.allowed_models,
"provider_api_keys.allowed_models",
)?)
.bind(optional_i64_from_u64(
key.expires_at_unix_secs,
"provider_api_keys.expires_at",
)?)
.bind(key.cache_ttl_minutes)
.bind(key.max_probe_interval_minutes)
.bind(optional_json_to_string(
&key.proxy,
"provider_api_keys.proxy",
)?)
.bind(optional_json_to_string(
&key.fingerprint,
"provider_api_keys.fingerprint",
)?)
.bind(optional_i64_from_u32(key.rpm_limit))
.bind(key.concurrent_limit)
.bind(optional_i64_from_u32(key.learned_rpm_limit))
.bind(optional_i64_from_u32(key.concurrent_429_count).unwrap_or(0))
.bind(optional_i64_from_u32(key.rpm_429_count).unwrap_or(0))
.bind(optional_i64_from_u64(
key.last_429_at_unix_secs,
"provider_api_keys.last_429_at",
)?)
.bind(&key.last_429_type)
.bind(optional_json_to_string(
&key.adjustment_history,
"provider_api_keys.adjustment_history",
)?)
.bind(optional_json_to_string(
&key.utilization_samples,
"provider_api_keys.utilization_samples",
)?)
.bind(optional_i64_from_u64(
key.last_probe_increase_at_unix_secs,
"provider_api_keys.last_probe_increase_at",
)?)
.bind(optional_i64_from_u32(key.last_rpm_peak))
.bind(optional_i64_from_u64(
key.last_models_fetch_at_unix_secs,
"provider_api_keys.last_models_fetch_at",
)?)
.bind(&key.last_models_fetch_error)
.bind(key.auto_fetch_models)
.bind(optional_json_to_string(
&key.locked_models,
"provider_api_keys.locked_models",
)?)
.bind(optional_json_to_string(
&key.model_include_patterns,
"provider_api_keys.model_include_patterns",
)?)
.bind(optional_json_to_string(
&key.model_exclude_patterns,
"provider_api_keys.model_exclude_patterns",
)?)
.bind(updated_at)
.bind(&key.id))
}
fn optional_json_from_string(
value: Option<String>,
field_name: &str,
@@ -2548,6 +2614,13 @@ mod tests {
created_key.last_models_fetch_error.as_deref(),
Some("stale models fetch error")
);
let mut second_key = key.clone();
second_key.id = "key-write-2".to_string();
second_key.name = "Secondary Key".to_string();
let created_second_key = repository
.create_key(&second_key)
.await
.expect("second key should create");
let mut updated_key = created_key.clone();
updated_key.name = "Updated Key".to_string();
@@ -2567,6 +2640,55 @@ mod tests {
);
assert_eq!(updated_key.last_models_fetch_error, None);
let mut batch_first = updated_key.clone();
batch_first.auto_fetch_models = true;
batch_first.allowed_models = Some(json!(["gpt-4.1", "gpt-4.1-mini"]));
batch_first.locked_models = Some(json!(["gpt-4.1"]));
batch_first.model_include_patterns = Some(json!(["gpt-*"]));
batch_first.model_exclude_patterns = Some(json!(["*-preview"]));
let mut batch_second = created_second_key;
batch_second.auto_fetch_models = true;
batch_second.allowed_models = batch_first.allowed_models.clone();
batch_second.locked_models = batch_first.locked_models.clone();
batch_second.model_include_patterns = batch_first.model_include_patterns.clone();
batch_second.model_exclude_patterns = batch_first.model_exclude_patterns.clone();
let batch_updated = repository
.update_keys(&[batch_first, batch_second])
.await
.expect("keys should update in one transaction");
assert_eq!(batch_updated.len(), 2);
assert!(batch_updated.iter().all(|key| key.auto_fetch_models));
assert!(batch_updated
.iter()
.all(|key| key.locked_models == Some(json!(["gpt-4.1"]))));
assert!(batch_updated
.iter()
.all(|key| key.model_include_patterns == Some(json!(["gpt-*"]))));
assert!(batch_updated
.iter()
.all(|key| key.model_exclude_patterns == Some(json!(["*-preview"]))));
let mut valid_change = batch_updated
.iter()
.find(|key| key.id == "key-write-1")
.expect("first key should be returned")
.clone();
valid_change.name = "Must Roll Back".to_string();
let mut missing_change = valid_change.clone();
missing_change.id = "missing-key".to_string();
assert!(repository
.update_keys(&[valid_change, missing_change])
.await
.is_err());
let rolled_back = repository
.list_keys_by_ids(&["key-write-1".to_string()])
.await
.expect("first key should reload")
.pop()
.expect("first key should exist");
assert_eq!(rolled_back.name, "Updated Key");
assert!(repository
.update_key_upstream_metadata(
"key-write-1",
@@ -2664,6 +2786,10 @@ mod tests {
.delete_key("key-write-1")
.await
.expect("key should delete"));
assert!(repository
.delete_key("key-write-2")
.await
.expect("second key should delete"));
assert!(repository
.delete_endpoint("endpoint-write-1")
.await
crates/aether-data/contracts/src/repository/provider_catalog/types.rs
+5
View File
@@ -684,6 +684,11 @@ pub trait ProviderCatalogWriteRepository: Send + Sync {
key: &StoredProviderCatalogKey,
) -> Result<StoredProviderCatalogKey, crate::DataLayerError>;
async fn update_keys(
&self,
keys: &[StoredProviderCatalogKey],
) -> Result<Vec<StoredProviderCatalogKey>, crate::DataLayerError>;
async fn update_key_upstream_metadata(
&self,
key_id: &str,
crates/aether-data/runtime/src/repository/provider_catalog/memory.rs
+22
View File
@@ -537,6 +537,28 @@ impl ProviderCatalogWriteRepository for InMemoryProviderCatalogReadRepository {
Ok(stored.clone())
}
async fn update_keys(
&self,
keys: &[StoredProviderCatalogKey],
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
let mut index = self
.index
.write()
.expect("provider catalog repository lock");
for key in keys {
if !index.keys.contains_key(&key.id) {
return Err(DataLayerError::UnexpectedValue(format!(
"provider catalog key {} not found",
key.id
)));
}
}
for key in keys {
index.keys.insert(key.id.clone(), key.clone());
}
Ok(keys.to_vec())
}
async fn update_key_upstream_metadata(
&self,
key_id: &str,