Files
Aether/crates/aether-data/adapters/postgres/src/provider_catalog.rs
T

4079 lines
129 KiB
Rust
Raw Normal View History

use std::collections::BTreeMap;
use async_trait::async_trait;
use futures_util::TryStreamExt;
use sqlx::{
postgres::{PgArguments, PgRow},
query::Query,
PgPool, Postgres, QueryBuilder, Row,
};
use aether_data_contracts::repository::provider_catalog::{
2026-08-14 09:28:07 +08:00
ProviderCatalogKeyAdaptiveStateUpdate, ProviderCatalogKeyAdminCasUpdate,
ProviderCatalogKeyCredentialsCasUpdate, ProviderCatalogKeyHealthStateUpdate,
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery,
ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate,
ProviderCatalogProviderConfigCasUpdate, ProviderCatalogProxyCasUpdate,
ProviderCatalogReadRepository, ProviderCatalogUpstreamMetadataNamespaceUpdate,
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
};
use aether_data_contracts::DataLayerError;
use crate::error::{postgres_error, SqlxResultExt};
use aether_data_query::{
push_ci_contains_any, push_eq, push_in, push_limit_offset, push_optional_eq, SqlDialect,
WhereClause,
};
const LIST_PROVIDERS_BY_IDS_PREFIX: &str = r#"
SELECT
id,
name,
description,
website,
provider_type,
CAST(billing_type AS TEXT) AS billing_type,
CAST(monthly_quota_usd AS DOUBLE PRECISION) AS monthly_quota_usd,
CAST(monthly_used_usd AS DOUBLE PRECISION) AS monthly_used_usd,
quota_reset_day,
CAST(EXTRACT(EPOCH FROM quota_last_reset_at) AS BIGINT) AS quota_last_reset_at_unix_secs,
CAST(EXTRACT(EPOCH FROM quota_expires_at) AS BIGINT) AS quota_expires_at_unix_secs,
provider_priority,
is_active,
keep_priority_on_conversion,
enable_format_conversion,
concurrent_limit,
max_retries,
proxy,
request_timeout,
stream_first_byte_timeout,
config,
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_ms,
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
FROM providers
WHERE id IN (
"#;
const LIST_ENDPOINTS_BY_IDS_PREFIX: &str = r#"
SELECT
id,
provider_id,
api_format,
api_family,
endpoint_kind,
is_active,
health_score,
base_url,
header_rules,
body_rules,
max_retries,
custom_path,
config,
format_acceptance_config,
proxy,
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_ms,
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
FROM provider_endpoints
WHERE id IN (
"#;
const LIST_ENDPOINTS_BY_IDS_PREFIX_LEGACY: &str = r#"
SELECT
id,
provider_id,
api_format,
api_family,
endpoint_kind,
is_active,
base_url,
header_rules,
body_rules,
max_retries,
custom_path,
config,
format_acceptance_config,
proxy,
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_ms,
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
FROM provider_endpoints
WHERE id IN (
"#;
const LIST_ENDPOINTS_BY_PROVIDER_IDS_PREFIX: &str = r#"
SELECT
id,
provider_id,
api_format,
api_family,
endpoint_kind,
is_active,
health_score,
base_url,
header_rules,
body_rules,
max_retries,
custom_path,
config,
format_acceptance_config,
proxy,
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_ms,
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
FROM provider_endpoints
WHERE provider_id IN (
"#;
const LIST_ENDPOINTS_BY_PROVIDER_IDS_PREFIX_LEGACY: &str = r#"
SELECT
id,
provider_id,
api_format,
api_family,
endpoint_kind,
is_active,
base_url,
header_rules,
body_rules,
max_retries,
custom_path,
config,
format_acceptance_config,
proxy,
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_ms,
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
FROM provider_endpoints
WHERE provider_id IN (
"#;
const LIST_KEYS_BY_IDS_PREFIX: &str = r#"
SELECT
id,
provider_id,
name,
auth_type,
capabilities,
is_active,
api_formats,
2026-04-29 15:46:50 +08:00
auth_type_by_format,
allow_auth_channel_mismatch_formats,
2026-05-06 00:20:06 +08:00
COALESCE(api_key, encrypted_key) AS api_key,
auth_config,
note,
internal_priority,
rate_multipliers,
global_priority_by_format,
allowed_models,
EXTRACT(EPOCH FROM expires_at)::bigint AS expires_at_unix_secs,
cache_ttl_minutes,
max_probe_interval_minutes,
proxy,
fingerprint,
rpm_limit,
concurrent_limit,
learned_rpm_limit,
concurrent_429_count,
rpm_429_count,
EXTRACT(EPOCH FROM last_429_at)::bigint AS last_429_at_unix_secs,
last_429_type,
adjustment_history,
utilization_samples,
EXTRACT(EPOCH FROM last_probe_increase_at)::bigint AS last_probe_increase_at_unix_secs,
last_rpm_peak,
request_count,
total_tokens,
CAST(total_cost_usd AS DOUBLE PRECISION) AS total_cost_usd,
success_count,
error_count,
total_response_time_ms,
EXTRACT(EPOCH FROM last_used_at)::bigint AS last_used_at_unix_secs,
auto_fetch_models,
EXTRACT(EPOCH FROM last_models_fetch_at)::bigint AS last_models_fetch_at_unix_secs,
last_models_fetch_error,
locked_models,
model_include_patterns,
model_exclude_patterns,
upstream_metadata,
EXTRACT(EPOCH FROM oauth_invalid_at)::bigint AS oauth_invalid_at_unix_secs,
oauth_invalid_reason,
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 id IN (
"#;
const LIST_KEYS_BY_PROVIDER_IDS_PREFIX: &str = r#"
SELECT
id,
provider_id,
name,
auth_type,
capabilities,
is_active,
api_formats,
2026-04-29 15:46:50 +08:00
auth_type_by_format,
allow_auth_channel_mismatch_formats,
2026-05-06 00:20:06 +08:00
COALESCE(api_key, encrypted_key) AS api_key,
auth_config,
note,
internal_priority,
rate_multipliers,
global_priority_by_format,
allowed_models,
EXTRACT(EPOCH FROM expires_at)::bigint AS expires_at_unix_secs,
cache_ttl_minutes,
max_probe_interval_minutes,
proxy,
fingerprint,
rpm_limit,
concurrent_limit,
learned_rpm_limit,
concurrent_429_count,
rpm_429_count,
EXTRACT(EPOCH FROM last_429_at)::bigint AS last_429_at_unix_secs,
last_429_type,
adjustment_history,
utilization_samples,
EXTRACT(EPOCH FROM last_probe_increase_at)::bigint AS last_probe_increase_at_unix_secs,
last_rpm_peak,
request_count,
total_tokens,
CAST(total_cost_usd AS DOUBLE PRECISION) AS total_cost_usd,
success_count,
error_count,
total_response_time_ms,
EXTRACT(EPOCH FROM last_used_at)::bigint AS last_used_at_unix_secs,
auto_fetch_models,
EXTRACT(EPOCH FROM last_models_fetch_at)::bigint AS last_models_fetch_at_unix_secs,
last_models_fetch_error,
locked_models,
model_include_patterns,
model_exclude_patterns,
upstream_metadata,
EXTRACT(EPOCH FROM oauth_invalid_at)::bigint AS oauth_invalid_at_unix_secs,
oauth_invalid_reason,
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_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,
NULL::jsonb AS capabilities,
is_active,
api_formats,
2026-04-29 15:46:50 +08:00
NULL::jsonb AS auth_type_by_format,
NULL::jsonb AS allow_auth_channel_mismatch_formats,
'summary' AS api_key,
CASE
WHEN auth_config IS NULL THEN NULL
ELSE '{}'::text
END AS auth_config,
NULL::text AS note,
NULL::integer AS internal_priority,
NULL::jsonb AS rate_multipliers,
NULL::jsonb AS global_priority_by_format,
NULL::jsonb AS allowed_models,
NULL::bigint AS expires_at_unix_secs,
NULL::integer AS cache_ttl_minutes,
NULL::integer AS max_probe_interval_minutes,
NULL::jsonb AS proxy,
NULL::jsonb AS fingerprint,
NULL::integer AS rpm_limit,
NULL::integer AS concurrent_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,
NULL::integer AS last_rpm_peak,
2026-05-19 00:58:09 +08:00
NULL::bigint AS request_count,
0::bigint AS total_tokens,
0::double precision AS total_cost_usd,
2026-05-19 00:58:09 +08:00
NULL::bigint AS success_count,
NULL::bigint AS error_count,
NULL::bigint AS total_response_time_ms,
NULL::bigint AS last_used_at_unix_secs,
FALSE AS 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,
NULL::bigint AS oauth_invalid_at_unix_secs,
NULL::text AS oauth_invalid_reason,
NULL::jsonb AS status_snapshot,
NULL::bigint AS created_at_unix_ms,
NULL::bigint AS updated_at_unix_secs,
health_by_format,
NULL::jsonb AS 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,
COUNT(*)::BIGINT AS total_keys,
COUNT(*) FILTER (WHERE is_active = TRUE)::BIGINT AS active_keys
FROM provider_api_keys
WHERE provider_id IN (
"#;
const LIST_KEY_MAINTENANCE_SUMMARIES_BY_PROVIDER_IDS_PREFIX: &str = r#"
SELECT
id,
provider_id,
is_active,
upstream_metadata
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 = $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,
allowed_models = $14,
capabilities = $15,
cache_ttl_minutes = $16,
max_probe_interval_minutes = $17,
auto_fetch_models = $18,
locked_models = $19,
model_include_patterns = $20,
model_exclude_patterns = $21,
proxy = $22,
fingerprint = $23,
expires_at = CASE
WHEN $24::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($24::double precision)
END,
is_active = $25,
updated_at = CASE
WHEN $26::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($26::double precision)
END,
auth_type_by_format = $27,
allow_auth_channel_mismatch_formats = $28
WHERE id = $1
2026-08-14 09:28:07 +08:00
AND provider_id = $2
AND auth_type = $4
AND api_key IS NOT DISTINCT FROM $5
AND auth_config IS NOT DISTINCT FROM $6
"#;
const KEY_RUNTIME_HEALTH_CAS_SQL: &str = r#"
UPDATE provider_api_keys
SET
health_by_format = $2,
circuit_breaker_by_format = $3,
updated_at = NOW()
WHERE id = $1
AND health_by_format::jsonb IS NOT DISTINCT FROM $4::jsonb
AND circuit_breaker_by_format::jsonb IS NOT DISTINCT FROM $5::jsonb
AND ($6::text IS NULL OR auth_config IS NOT DISTINCT FROM $6)
"#;
const KEY_RUNTIME_METADATA_NAMESPACE_LOCK_SQL: &str = r#"
SELECT
jsonb_typeof(COALESCE(upstream_metadata, '{}'::jsonb)) = 'object'
AS metadata_is_object,
COALESCE(upstream_metadata, '{}'::jsonb) ? $2 AS namespace_exists,
COALESCE(upstream_metadata, '{}'::jsonb) -> $2 AS namespace_value
FROM provider_api_keys
WHERE id = $1
FOR UPDATE
"#;
const KEY_RUNTIME_METADATA_UPDATE_SQL: &str = r#"
UPDATE provider_api_keys
SET
upstream_metadata = COALESCE(upstream_metadata, '{}'::jsonb)
|| jsonb_build_object($2, $3::jsonb),
status_snapshot = (COALESCE(status_snapshot::jsonb, '{}'::jsonb) || $4::jsonb)::json,
updated_at = CASE
WHEN $5::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($5::double precision)
END
WHERE id = $1
2026-08-14 09:28:07 +08:00
AND jsonb_typeof(COALESCE(upstream_metadata, '{}'::jsonb)) = 'object'
"#;
fn runtime_metadata_namespace_matches(
metadata_is_object: bool,
namespace_exists: bool,
current: Option<&serde_json::Value>,
expected: Option<&serde_json::Value>,
) -> bool {
metadata_is_object
&& match expected {
Some(expected) => namespace_exists && current == Some(expected),
None => !namespace_exists,
}
}
async fn lock_runtime_metadata_namespace_matches(
tx: &mut sqlx::Transaction<'_, Postgres>,
key_id: &str,
namespace: &str,
expected: Option<&serde_json::Value>,
) -> Result<bool, DataLayerError> {
let Some(row) = sqlx::query(KEY_RUNTIME_METADATA_NAMESPACE_LOCK_SQL)
.bind(key_id)
.bind(namespace)
.fetch_optional(&mut **tx)
.await
.map_postgres_err()?
else {
return Ok(false);
};
let metadata_is_object = row
.try_get::<bool, _>("metadata_is_object")
.map_postgres_err()?;
let namespace_exists = row
.try_get::<bool, _>("namespace_exists")
.map_postgres_err()?;
let current = row
.try_get::<Option<serde_json::Value>, _>("namespace_value")
.map_postgres_err()?;
// PostgreSQL jsonb retains decimal lexemes that serde_json's default
// Number representation rounds to f64. Re-read and compare while holding
// the row lock instead of binding that rounded value back into a jsonb
// equality predicate, which would report a false CAS conflict.
Ok(runtime_metadata_namespace_matches(
metadata_is_object,
namespace_exists,
current.as_ref(),
expected,
))
}
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 adaptive_status_snapshot_patch(
patch: &serde_json::Value,
) -> Result<serde_json::Value, DataLayerError> {
const OWNED_FIELDS: [&str; 6] = [
"observation_count",
"header_observation_count",
"latest_upstream_limit",
"learning_confidence",
"enforcement_active",
"known_boundary",
];
let object = patch.as_object().ok_or_else(|| {
DataLayerError::InvalidInput(
"provider catalog adaptive status snapshot patch must be an object".to_string(),
)
})?;
Ok(serde_json::Value::Object(
OWNED_FIELDS
.into_iter()
.filter_map(|field| {
object
.get(field)
.cloned()
.map(|value| (field.to_string(), value))
})
.collect(),
))
}
fn validate_runtime_metadata_update(
update: &ProviderCatalogKeyRuntimeMetadataUpdate,
) -> Result<(), DataLayerError> {
if update.key_id.trim().is_empty() || update.namespace.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog key_id and runtime metadata namespace are required".to_string(),
));
}
if !update.status_snapshot_patch.is_object() {
return Err(DataLayerError::InvalidInput(
"provider catalog runtime status snapshot patch must be an object".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.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.expires_at_unix_secs.map(|value| value as f64))
.bind(key.is_active)
.bind(key.updated_at_unix_secs.map(|value| value as f64))
.bind(&key.auth_type_by_format)
.bind(&key.allow_auth_channel_mismatch_formats)
}
2026-08-14 09:28:07 +08:00
fn push_admin_key_assignments<'args>(
builder: &mut QueryBuilder<'args, Postgres>,
key: &'args StoredProviderCatalogKey,
) {
builder
.push_bind(&key.provider_id)
.push(", api_formats = ")
.push_bind(&key.api_formats)
.push(", auth_type = ")
.push_bind(&key.auth_type)
.push(", api_key = ")
.push_bind(&key.encrypted_api_key)
.push(", auth_config = ")
.push_bind(&key.encrypted_auth_config)
.push(", name = ")
.push_bind(&key.name)
.push(", note = ")
.push_bind(&key.note)
.push(", rate_multipliers = ")
.push_bind(&key.rate_multipliers)
.push(", internal_priority = ")
.push_bind(key.internal_priority)
.push(", global_priority_by_format = ")
.push_bind(&key.global_priority_by_format)
.push(", rpm_limit = ")
.push_bind(key.rpm_limit.map(|value| value as i32))
.push(", concurrent_limit = ")
.push_bind(key.concurrent_limit)
.push(", allowed_models = ")
.push_bind(&key.allowed_models)
.push(", capabilities = ")
.push_bind(&key.capabilities)
.push(", cache_ttl_minutes = ")
.push_bind(key.cache_ttl_minutes)
.push(", max_probe_interval_minutes = ")
.push_bind(key.max_probe_interval_minutes)
.push(", auto_fetch_models = ")
.push_bind(key.auto_fetch_models)
.push(", locked_models = ")
.push_bind(&key.locked_models)
.push(", model_include_patterns = ")
.push_bind(&key.model_include_patterns)
.push(", model_exclude_patterns = ")
.push_bind(&key.model_exclude_patterns)
.push(", proxy = ")
.push_bind(&key.proxy)
.push(", fingerprint = ")
.push_bind(&key.fingerprint)
.push(", expires_at = CASE WHEN ")
.push_bind(key.expires_at_unix_secs.map(|value| value as f64))
.push("::double precision IS NULL THEN NULL ELSE TO_TIMESTAMP(")
.push_bind(key.expires_at_unix_secs.map(|value| value as f64))
.push("::double precision) END, is_active = ")
.push_bind(key.is_active)
.push(", updated_at = CASE WHEN ")
.push_bind(key.updated_at_unix_secs.map(|value| value as f64))
.push("::double precision IS NULL THEN NOW() ELSE TO_TIMESTAMP(")
.push_bind(key.updated_at_unix_secs.map(|value| value as f64))
.push("::double precision) END, auth_type_by_format = ")
.push_bind(&key.auth_type_by_format)
.push(", allow_auth_channel_mismatch_formats = ")
.push_bind(&key.allow_auth_channel_mismatch_formats);
}
fn validate_admin_key_cas_update(
update: &ProviderCatalogKeyAdminCasUpdate,
) -> Result<(), DataLayerError> {
validate_key_for_update(&update.key)?;
let expected = &update.expected_credential;
if update.key.name.trim().is_empty()
|| update.key.auth_type.trim().is_empty()
|| expected.auth_type.trim().is_empty()
|| expected.provider_id.trim().is_empty()
|| expected.provider_type.trim().is_empty()
|| expected
.encrypted_api_key
.as_deref()
.is_some_and(|value| value.trim().is_empty())
|| update
.expected_encrypted_auth_config
.as_deref()
.is_some_and(|value| value.trim().is_empty())
{
return Err(DataLayerError::InvalidInput(
"provider catalog admin credential fence contains empty fields".to_string(),
));
}
let Some(rotation) = update.codex_rotation.as_ref() else {
return Ok(());
};
let valid_rotation = expected.provider_type.eq_ignore_ascii_case("codex")
&& rotation.as_object().is_some_and(|object| {
object.len() == 1
&& object
.get("credential_generation")
.and_then(serde_json::Value::as_str)
.is_some_and(|generation| !generation.trim().is_empty())
});
if !valid_rotation {
return Err(DataLayerError::InvalidInput(
"provider catalog Codex rotation must contain only credential_generation".to_string(),
));
}
Ok(())
}
#[derive(Debug, Clone)]
pub struct SqlxProviderCatalogReadRepository {
pool: PgPool,
}
impl SqlxProviderCatalogReadRepository {
pub fn new(pool: PgPool) -> Self {
Self { pool }
}
pub fn pool(&self) -> &PgPool {
&self.pool
}
pub async fn list_providers_by_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
if provider_ids.is_empty() {
return Ok(Vec::new());
}
collect_query_rows(
build_list_query(
LIST_PROVIDERS_BY_IDS_PREFIX,
provider_ids,
" ORDER BY name ASC",
)
.build()
.fetch(&self.pool),
map_provider_row,
)
.await
}
pub async fn list_providers(
&self,
active_only: bool,
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
let mut builder =
QueryBuilder::<Postgres>::new(select_prefix_for_in(LIST_PROVIDERS_BY_IDS_PREFIX));
let mut where_clause = WhereClause::new();
if active_only {
push_eq(&mut builder, &mut where_clause, "is_active", true);
}
builder.push(" ORDER BY provider_priority ASC, name ASC");
collect_query_rows(builder.build().fetch(&self.pool), map_provider_row).await
}
pub async fn list_endpoints_by_ids(
&self,
endpoint_ids: &[String],
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
if endpoint_ids.is_empty() {
return Ok(Vec::new());
}
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(&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(error),
};
Ok(rows)
}
pub async fn list_endpoints_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
if provider_ids.is_empty() {
return Ok(Vec::new());
}
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(&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(error),
};
Ok(rows)
}
pub async fn list_keys_by_ids(
&self,
key_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
if key_ids.is_empty() {
return Ok(Vec::new());
}
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,
)
.await
}
pub async fn list_keys_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_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,
)
.await
}
pub async fn list_key_maintenance_summaries_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKeyMaintenanceSummary>, DataLayerError> {
if provider_ids.is_empty() {
return Ok(Vec::new());
}
collect_query_rows(
build_list_query(
LIST_KEY_MAINTENANCE_SUMMARIES_BY_PROVIDER_IDS_PREFIX,
provider_ids,
" ORDER BY provider_id ASC, id ASC",
)
.build()
.fetch(&self.pool),
map_key_maintenance_summary_row,
)
.await
}
pub async fn list_keys_page(
&self,
query: &ProviderCatalogKeyListQuery,
) -> Result<StoredProviderCatalogKeyPage, DataLayerError> {
if query.provider_id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog provider_id is empty".to_string(),
));
}
let offset = i64::try_from(query.offset).map_err(|_| {
DataLayerError::InvalidInput(format!(
"invalid provider catalog key offset: {}",
query.offset
))
})?;
let limit = i64::try_from(query.limit).map_err(|_| {
DataLayerError::InvalidInput(format!(
"invalid provider catalog key limit: {}",
query.limit
))
})?;
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"
}
ProviderCatalogKeyListOrder::CreatedAtAsc => {
"created_at ASC NULLS LAST, name ASC, id ASC"
}
ProviderCatalogKeyListOrder::CreatedAtDesc => {
"created_at DESC NULLS LAST, name ASC, id ASC"
}
ProviderCatalogKeyListOrder::LastUsedAtAsc => {
"last_used_at ASC NULLS LAST, name ASC, id ASC"
}
ProviderCatalogKeyListOrder::LastUsedAtDesc => {
"last_used_at DESC NULLS LAST, name ASC, id ASC"
}
};
let mut count_builder = QueryBuilder::<Postgres>::new(
"SELECT COUNT(*)::BIGINT AS total FROM provider_api_keys",
);
let mut count_where = WhereClause::new();
apply_key_page_filters(&mut count_builder, &mut count_where, query);
let total = count_builder
.build_query_scalar::<i64>()
.fetch_one(&self.pool)
.await
.map_postgres_err()?
.max(0) as usize;
let mut list_builder =
QueryBuilder::<Postgres>::new(select_prefix_for_in(LIST_KEYS_BY_IDS_PREFIX));
let mut list_where = WhereClause::new();
apply_key_page_filters(&mut list_builder, &mut list_where, query);
list_builder.push(" ORDER BY ").push(order_by);
push_limit_offset(&mut list_builder, limit, offset);
let items = collect_query_rows(list_builder.build().fetch(&self.pool), map_key_row).await?;
Ok(StoredProviderCatalogKeyPage { items, total })
}
pub async fn list_key_stats_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKeyStats>, DataLayerError> {
if provider_ids.is_empty() {
return Ok(Vec::new());
}
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,
)
.await
}
pub async fn update_key_oauth_runtime_state(
&self,
key_id: &str,
oauth_invalid_at_unix_secs: Option<u64>,
oauth_invalid_reason: Option<&str>,
updated_at_unix_secs: Option<u64>,
) -> Result<bool, DataLayerError> {
if key_id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog key_id is empty".to_string(),
));
}
let rows_affected = sqlx::query(
r#"
UPDATE provider_api_keys
SET
oauth_invalid_at = CASE
WHEN $2::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($2::double precision)
END,
oauth_invalid_reason = $3,
updated_at = CASE
WHEN $4::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($4::double precision)
END
WHERE id = $1
"#,
)
.bind(key_id)
.bind(oauth_invalid_at_unix_secs.map(|value| value as f64))
.bind(oauth_invalid_reason)
.bind(updated_at_unix_secs.map(|value| value as f64))
.execute(&self.pool)
.await
.map_postgres_err()?
.rows_affected();
Ok(rows_affected > 0)
}
pub async fn compare_and_update_key_oauth_runtime_state(
&self,
update: &ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
) -> Result<bool, DataLayerError> {
if update.key_id.trim().is_empty()
|| update.encrypted_auth_config.trim().is_empty()
|| update
.encrypted_api_key_update
.as_deref()
.is_some_and(|value| value.trim().is_empty())
|| update.expected_credential.as_ref().is_some_and(|expected| {
expected
.encrypted_api_key
.as_deref()
.is_some_and(|value| value.trim().is_empty())
|| expected.auth_type.trim().is_empty()
|| expected.provider_id.trim().is_empty()
|| expected.provider_type.trim().is_empty()
})
2026-08-14 09:28:07 +08:00
|| update
.expected_upstream_metadata_namespace
.as_ref()
.is_some_and(|expected| expected.namespace.trim().is_empty())
|| update
.upstream_metadata_namespace_to_remove
.as_deref()
.is_some_and(|namespace| namespace.trim().is_empty())
|| !update.status_snapshot_patch.is_object()
|| update
.upstream_metadata_patch
.as_ref()
.is_some_and(|patch| !patch.is_object())
2026-08-14 09:28:07 +08:00
|| update
.upstream_metadata_namespace_to_remove
.as_ref()
.is_some_and(|namespace| {
update
.upstream_metadata_patch
.as_ref()
.and_then(serde_json::Value::as_object)
.is_some_and(|patch| patch.contains_key(namespace))
})
{
return Err(DataLayerError::InvalidInput(
"provider catalog OAuth runtime CAS requires key_id, auth_config, and object status patch"
.to_string(),
));
}
let mut tx = self.pool.begin().await.map_postgres_err()?;
if let Some(expected) = update.expected_upstream_metadata_namespace.as_ref() {
let matches = lock_runtime_metadata_namespace_matches(
&mut tx,
&update.key_id,
&expected.namespace,
expected.expected_value.as_ref(),
)
.await?;
if !matches {
tx.rollback().await.map_postgres_err()?;
return Ok(false);
}
}
let rows_affected = sqlx::query(
r#"
UPDATE provider_api_keys
SET
oauth_invalid_at = CASE
WHEN $2::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($2::double precision)
END,
oauth_invalid_reason = $3,
auth_config = $4,
api_key = CASE
WHEN $5::text IS NULL THEN api_key
ELSE $5
END,
expires_at = CASE
WHEN $6::boolean IS FALSE THEN expires_at
WHEN $7::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($7::double precision)
END,
upstream_metadata = CASE
2026-08-14 09:28:07 +08:00
WHEN $8::jsonb IS NULL AND $9::text IS NULL THEN upstream_metadata
WHEN $9::text IS NULL THEN COALESCE(upstream_metadata, '{}'::jsonb) || $8::jsonb
ELSE (COALESCE(upstream_metadata, '{}'::jsonb) || COALESCE($8::jsonb, '{}'::jsonb)) - $9
END,
2026-08-14 09:28:07 +08:00
status_snapshot = (COALESCE(status_snapshot::jsonb, '{}'::jsonb) || $10::jsonb)::json,
error_count = CASE WHEN $11::boolean THEN 0 ELSE error_count END,
updated_at = CASE
2026-08-14 09:28:07 +08:00
WHEN $12::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($12::double precision)
END
WHERE id = $1
2026-08-14 09:28:07 +08:00
AND auth_config IS NOT DISTINCT FROM $13
AND (
2026-08-14 09:28:07 +08:00
($8::jsonb IS NULL AND $9::text IS NULL)
OR jsonb_typeof(COALESCE(upstream_metadata, '{}'::jsonb)) = 'object'
)
AND ($14::boolean IS FALSE OR api_key IS NOT DISTINCT FROM $15)
AND ($16::text IS NULL OR auth_type = $16)
AND ($17::text IS NULL OR provider_id = $17)
AND (
$18::text IS NULL
OR EXISTS (
SELECT 1
FROM providers
WHERE providers.id = provider_api_keys.provider_id
2026-08-14 09:28:07 +08:00
AND providers.provider_type = $18
)
)
"#,
)
.bind(&update.key_id)
.bind(update.oauth_invalid_at_unix_secs.map(|value| value as f64))
.bind(update.oauth_invalid_reason.as_deref())
.bind(&update.encrypted_auth_config)
.bind(update.encrypted_api_key_update.as_deref())
.bind(update.expires_at_unix_secs_update.is_some())
.bind(
update
.expires_at_unix_secs_update
.flatten()
.map(|value| value as f64),
)
.bind(update.upstream_metadata_patch.as_ref())
2026-08-14 09:28:07 +08:00
.bind(update.upstream_metadata_namespace_to_remove.as_deref())
.bind(&update.status_snapshot_patch)
.bind(update.reset_error_count)
.bind(update.updated_at_unix_secs.map(|value| value as f64))
.bind(update.expected_encrypted_auth_config.as_deref())
.bind(update.expected_credential.is_some())
.bind(
update
.expected_credential
.as_ref()
.and_then(|expected| expected.encrypted_api_key.as_deref()),
)
.bind(
update
.expected_credential
.as_ref()
.map(|expected| expected.auth_type.as_str()),
)
.bind(
update
.expected_credential
.as_ref()
.map(|expected| expected.provider_id.as_str()),
)
.bind(
update
.expected_credential
.as_ref()
.map(|expected| expected.provider_type.as_str()),
)
.execute(&mut *tx)
.await
.map_postgres_err()?
.rows_affected();
if rows_affected == 0 {
tx.rollback().await.map_postgres_err()?;
return Ok(false);
}
tx.commit().await.map_postgres_err()?;
Ok(true)
}
pub async fn create_provider(
&self,
provider: &StoredProviderCatalogProvider,
shift_existing_priorities_from: Option<i32>,
) -> Result<StoredProviderCatalogProvider, DataLayerError> {
if provider.id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog provider.id is empty".to_string(),
));
}
if provider.name.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog provider.name is empty".to_string(),
));
}
if provider.provider_type.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog provider.provider_type is empty".to_string(),
));
}
if provider
.billing_type
.as_deref()
.map(str::trim)
.is_some_and(str::is_empty)
{
return Err(DataLayerError::InvalidInput(
"provider catalog provider.billing_type is empty".to_string(),
));
}
let mut tx = self.pool.begin().await.map_postgres_err()?;
if let Some(target_priority) = shift_existing_priorities_from {
sqlx::query(
r#"
UPDATE providers
SET provider_priority = provider_priority + 1
WHERE provider_priority IS NOT NULL
AND provider_priority >= $1
"#,
)
.bind(target_priority)
.execute(&mut *tx)
.await
.map_postgres_err()?;
}
sqlx::query(
r#"
INSERT INTO providers (
id,
name,
description,
website,
provider_type,
billing_type,
monthly_quota_usd,
monthly_used_usd,
quota_reset_day,
quota_last_reset_at,
quota_expires_at,
provider_priority,
is_active,
keep_priority_on_conversion,
enable_format_conversion,
concurrent_limit,
max_retries,
proxy,
request_timeout,
stream_first_byte_timeout,
config,
created_at,
updated_at
) VALUES (
$1,
$2,
$3,
$4,
$5,
CAST($6 AS providerbillingtype),
$7,
$8,
$9,
CASE
WHEN $10::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($10::double precision)
END,
CASE
WHEN $11::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($11::double precision)
END,
$12,
$13,
$14,
$15,
$16,
$17,
$18,
$19,
$20,
$21,
CASE
WHEN $22::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($22::double precision)
END,
CASE
WHEN $23::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($23::double precision)
END
)
"#,
)
.bind(&provider.id)
.bind(&provider.name)
.bind(&provider.description)
.bind(&provider.website)
.bind(&provider.provider_type)
.bind(
provider
.billing_type
.clone()
.unwrap_or_else(|| "pay_as_you_go".to_string()),
)
.bind(provider.monthly_quota_usd)
.bind(provider.monthly_used_usd)
.bind(provider.quota_reset_day.map(|value| value as i32))
.bind(
provider
.quota_last_reset_at_unix_secs
.map(|value| value as f64),
)
.bind(
provider
.quota_expires_at_unix_secs
.map(|value| value as f64),
)
.bind(provider.provider_priority)
.bind(provider.is_active)
.bind(provider.keep_priority_on_conversion)
.bind(provider.enable_format_conversion)
.bind(provider.concurrent_limit)
.bind(provider.max_retries)
.bind(&provider.proxy)
.bind(provider.request_timeout_secs)
.bind(provider.stream_first_byte_timeout_secs)
.bind(&provider.config)
.bind(provider.created_at_unix_ms.map(|value| value as f64))
.bind(provider.updated_at_unix_secs.map(|value| value as f64))
.execute(&mut *tx)
.await
.map_postgres_err()?;
tx.commit().await.map_err(postgres_error)?;
self.list_providers_by_ids(std::slice::from_ref(&provider.id))
.await?
.into_iter()
.next()
.ok_or_else(|| {
DataLayerError::UnexpectedValue(format!(
"created provider catalog provider {} could not be reloaded",
provider.id
))
})
}
pub async fn update_provider(
&self,
provider: &StoredProviderCatalogProvider,
) -> Result<StoredProviderCatalogProvider, DataLayerError> {
if provider.id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog provider.id is empty".to_string(),
));
}
if provider.name.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog provider.name is empty".to_string(),
));
}
if provider.provider_type.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog provider.provider_type is empty".to_string(),
));
}
if provider
.billing_type
.as_deref()
.map(str::trim)
.is_some_and(str::is_empty)
{
return Err(DataLayerError::InvalidInput(
"provider catalog provider.billing_type is empty".to_string(),
));
}
let rows_affected = sqlx::query(
r#"
UPDATE providers
SET
name = $2,
description = $3,
website = $4,
provider_type = $5,
billing_type = CAST($6 AS providerbillingtype),
monthly_quota_usd = $7,
monthly_used_usd = COALESCE($8, monthly_used_usd),
quota_reset_day = $9,
quota_last_reset_at = CASE
WHEN $10::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($10::double precision)
END,
quota_expires_at = CASE
WHEN $11::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($11::double precision)
END,
provider_priority = $12,
is_active = $13,
keep_priority_on_conversion = $14,
enable_format_conversion = $15,
concurrent_limit = $16,
max_retries = $17,
proxy = $18,
request_timeout = $19,
stream_first_byte_timeout = $20,
config = $21,
updated_at = CASE
WHEN $22::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($22::double precision)
END
WHERE id = $1
"#,
)
.bind(&provider.id)
.bind(&provider.name)
.bind(&provider.description)
.bind(&provider.website)
.bind(&provider.provider_type)
.bind(
provider
.billing_type
.clone()
.unwrap_or_else(|| "pay_as_you_go".to_string()),
)
.bind(provider.monthly_quota_usd)
.bind(provider.monthly_used_usd)
.bind(provider.quota_reset_day.map(|value| value as i32))
.bind(
provider
.quota_last_reset_at_unix_secs
.map(|value| value as f64),
)
.bind(
provider
.quota_expires_at_unix_secs
.map(|value| value as f64),
)
.bind(provider.provider_priority)
.bind(provider.is_active)
.bind(provider.keep_priority_on_conversion)
.bind(provider.enable_format_conversion)
.bind(provider.concurrent_limit)
.bind(provider.max_retries)
.bind(&provider.proxy)
.bind(provider.request_timeout_secs)
.bind(provider.stream_first_byte_timeout_secs)
.bind(&provider.config)
.bind(provider.updated_at_unix_secs.map(|value| value as f64))
.execute(&self.pool)
.await
.map_postgres_err()?
.rows_affected();
if rows_affected == 0 {
return Err(DataLayerError::UnexpectedValue(format!(
"provider catalog provider {} not found",
provider.id
)));
}
self.list_providers_by_ids(std::slice::from_ref(&provider.id))
.await?
.into_iter()
.next()
.ok_or_else(|| {
DataLayerError::UnexpectedValue(format!(
"updated provider catalog provider {} could not be reloaded",
provider.id
))
})
}
pub async fn compare_and_swap_provider_config(
&self,
update: &ProviderCatalogProviderConfigCasUpdate,
) -> Result<bool, DataLayerError> {
if update.provider_id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog provider_id is empty".to_string(),
));
}
let rows_affected = sqlx::query(
r#"
UPDATE providers
SET config = $3, updated_at = NOW()
WHERE id = $1
AND config::jsonb IS NOT DISTINCT FROM $2::jsonb
"#,
)
.bind(&update.provider_id)
.bind(&update.expected_config)
.bind(&update.config)
.execute(&self.pool)
.await
.map_postgres_err()?
.rows_affected();
Ok(rows_affected == 1)
}
pub async fn compare_and_swap_provider_proxy(
&self,
update: &ProviderCatalogProxyCasUpdate,
) -> Result<bool, DataLayerError> {
if update.record_id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog provider_id is empty".to_string(),
));
}
let rows_affected = sqlx::query(
r#"
UPDATE providers
SET proxy = $3, updated_at = NOW()
WHERE id = $1
AND proxy::jsonb IS NOT DISTINCT FROM $2::jsonb
"#,
)
.bind(&update.record_id)
.bind(&update.expected_proxy)
.bind(&update.proxy)
.execute(&self.pool)
.await
.map_postgres_err()?
.rows_affected();
Ok(rows_affected == 1)
}
pub async fn delete_provider(&self, provider_id: &str) -> Result<bool, DataLayerError> {
if provider_id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog provider_id is empty".to_string(),
));
}
let rows_affected = sqlx::query(
r#"
DELETE FROM providers
WHERE id = $1
"#,
)
.bind(provider_id)
.execute(&self.pool)
.await
.map_postgres_err()?
.rows_affected();
Ok(rows_affected > 0)
}
pub async fn cleanup_deleted_provider_refs(
&self,
provider_id: &str,
2026-06-12 16:25:11 +08:00
provider_deleted: bool,
endpoint_ids: &[String],
key_ids: &[String],
) -> Result<(), DataLayerError> {
if provider_id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog provider_id is empty".to_string(),
));
}
let mut tx = self.pool.begin().await.map_postgres_err()?;
2026-06-12 16:25:11 +08:00
if provider_deleted {
sqlx::query(
"UPDATE user_preferences SET default_provider_id = NULL WHERE default_provider_id = $1",
)
.bind(provider_id)
.execute(&mut *tx)
.await
.map_postgres_err()?;
2026-06-12 16:25:11 +08:00
sqlx::query("UPDATE video_tasks SET provider_id = NULL WHERE provider_id = $1")
.bind(provider_id)
.execute(&mut *tx)
.await
.map_postgres_err()?;
sqlx::query("DELETE FROM request_candidates WHERE provider_id = $1")
.bind(provider_id)
.execute(&mut *tx)
.await
.map_postgres_err()?;
}
for endpoint_id in endpoint_ids {
sqlx::query("UPDATE video_tasks SET endpoint_id = NULL WHERE endpoint_id = $1")
.bind(endpoint_id)
.execute(&mut *tx)
.await
.map_postgres_err()?;
sqlx::query("DELETE FROM request_candidates WHERE endpoint_id = $1")
.bind(endpoint_id)
.execute(&mut *tx)
.await
.map_postgres_err()?;
}
for key_id in key_ids {
sqlx::query("DELETE FROM gemini_file_mappings WHERE key_id = $1")
.bind(key_id)
.execute(&mut *tx)
.await
.map_postgres_err()?;
sqlx::query("UPDATE video_tasks SET key_id = NULL WHERE key_id = $1")
.bind(key_id)
.execute(&mut *tx)
.await
.map_postgres_err()?;
}
2026-06-12 16:25:11 +08:00
if provider_deleted {
sqlx::query("DELETE FROM api_key_provider_mappings WHERE provider_id = $1")
.bind(provider_id)
.execute(&mut *tx)
.await
.map_postgres_err()?;
sqlx::query("DELETE FROM provider_usage_tracking WHERE provider_id = $1")
.bind(provider_id)
.execute(&mut *tx)
.await
.map_postgres_err()?;
}
tx.commit().await.map_err(postgres_error)?;
Ok(())
}
pub async fn clear_key_oauth_invalid_marker(
&self,
key_id: &str,
) -> Result<bool, DataLayerError> {
if key_id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog key_id is empty".to_string(),
));
}
let rows_affected = sqlx::query(
r#"
UPDATE provider_api_keys
SET
oauth_invalid_at = NULL,
oauth_invalid_reason = NULL,
updated_at = NOW()
WHERE id = $1
"#,
)
.bind(key_id)
.execute(&self.pool)
.await
.map_postgres_err()?
.rows_affected();
Ok(rows_affected > 0)
}
pub async fn create_key(
&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(),
));
}
sqlx::query(
r#"
INSERT INTO provider_api_keys (
id,
provider_id,
api_formats,
2026-04-29 15:46:50 +08:00
auth_type_by_format,
auth_type,
api_key,
auth_config,
name,
note,
rate_multipliers,
internal_priority,
global_priority_by_format,
rpm_limit,
concurrent_limit,
learned_rpm_limit,
allowed_models,
capabilities,
cache_ttl_minutes,
max_probe_interval_minutes,
auto_fetch_models,
locked_models,
model_include_patterns,
model_exclude_patterns,
proxy,
fingerprint,
upstream_metadata,
expires_at,
oauth_invalid_at,
oauth_invalid_reason,
status_snapshot,
concurrent_429_count,
rpm_429_count,
last_429_at,
last_429_type,
adjustment_history,
utilization_samples,
last_probe_increase_at,
last_rpm_peak,
request_count,
total_tokens,
total_cost_usd,
success_count,
error_count,
total_response_time_ms,
last_used_at,
last_models_fetch_at,
last_models_fetch_error,
health_by_format,
circuit_breaker_by_format,
is_active,
created_at,
updated_at,
allow_auth_channel_mismatch_formats
) VALUES (
$1,
$2,
$3,
$4,
$5,
$6,
$7,
$8,
$9,
$10,
$11,
$12,
$13,
$14,
$15,
$16,
$17,
$18,
$19,
$20,
$21,
$22,
$23,
$24,
2026-04-29 15:46:50 +08:00
$25,
2026-05-04 12:21:54 +08:00
$26,
CASE
2026-04-29 15:46:50 +08:00
WHEN $27::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($27::double precision)
END,
CASE
2026-05-04 12:21:54 +08:00
WHEN $28::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($28::double precision)
END,
$29,
$30,
COALESCE($31, 0),
COALESCE($32, 0),
CASE
WHEN $33::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($33::double precision)
END,
$34,
2026-04-29 15:46:50 +08:00
$35,
2026-05-04 12:21:54 +08:00
$36,
CASE
2026-05-04 12:21:54 +08:00
WHEN $37::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($37::double precision)
END,
COALESCE($38, 0),
COALESCE($39, 0),
COALESCE($40, 0),
COALESCE($41, 0),
COALESCE($42, 0),
2026-04-29 15:46:50 +08:00
COALESCE($43, 0),
2026-05-04 12:21:54 +08:00
COALESCE($44, 0),
2026-04-29 15:46:50 +08:00
CASE
WHEN $45::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($45::double precision)
END,
2026-05-04 12:21:54 +08:00
CASE
WHEN $46::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($46::double precision)
END,
$47,
$48,
2026-04-29 15:46:50 +08:00
$49,
2026-05-04 12:21:54 +08:00
$50,
2026-04-29 15:46:50 +08:00
CASE
WHEN $51::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($51::double precision)
END,
2026-05-04 12:21:54 +08:00
CASE
WHEN $52::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($52::double precision)
END,
$53
)
"#,
)
.bind(&key.id)
.bind(&key.provider_id)
.bind(&key.api_formats)
2026-04-29 15:46:50 +08:00
.bind(&key.auth_type_by_format)
.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.expires_at_unix_secs.map(|value| value as f64))
.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))
2026-05-19 00:58:09 +08:00
.bind(key.request_count.map(i64::from))
.bind(Some(i64::try_from(key.total_tokens).map_err(|_| {
DataLayerError::InvalidInput(format!(
"provider catalog key.total_tokens exceeds i64: {}",
key.total_tokens
))
})?))
.bind(key.total_cost_usd)
2026-05-19 00:58:09 +08:00
.bind(key.success_count.map(i64::from))
.bind(key.error_count.map(i64::from))
2026-07-09 14:34:09 +08:00
.bind(
key.total_response_time_ms
.map(|value| {
i64::try_from(value).map_err(|_| {
DataLayerError::InvalidInput(format!(
"provider catalog key.total_response_time_ms exceeds i64: {value}"
))
})
})
.transpose()?,
)
.bind(key.last_used_at_unix_secs.map(|value| value as f64))
.bind(key.last_models_fetch_at_unix_secs.map(|value| value as f64))
.bind(&key.last_models_fetch_error)
.bind(&key.health_by_format)
.bind(&key.circuit_breaker_by_format)
.bind(key.is_active)
.bind(key.created_at_unix_ms.map(|value| value as f64))
.bind(key.updated_at_unix_secs.map(|value| value as f64))
.bind(&key.allow_auth_channel_mismatch_formats)
.execute(&self.pool)
.await
.map_postgres_err()?;
self.list_keys_by_ids(std::slice::from_ref(&key.id))
.await?
.into_iter()
.next()
.ok_or_else(|| {
DataLayerError::UnexpectedValue(format!(
"created provider catalog key {} could not be reloaded",
key.id
))
})
}
pub async fn create_endpoint(
&self,
endpoint: &StoredProviderCatalogEndpoint,
) -> Result<StoredProviderCatalogEndpoint, DataLayerError> {
if endpoint.id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog endpoint.id is empty".to_string(),
));
}
if endpoint.provider_id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog endpoint.provider_id is empty".to_string(),
));
}
match sqlx::query(
r#"
INSERT INTO provider_endpoints (
id,
provider_id,
api_format,
api_family,
endpoint_kind,
is_active,
health_score,
base_url,
header_rules,
body_rules,
max_retries,
custom_path,
config,
format_acceptance_config,
proxy,
created_at,
updated_at
) VALUES (
$1,
$2,
$3,
$4,
$5,
$6,
$7,
$8,
$9,
$10,
$11,
$12,
$13,
$14,
$15,
CASE
WHEN $16::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($16::double precision)
END,
CASE
WHEN $17::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($17::double precision)
END
)
"#,
)
.bind(&endpoint.id)
.bind(&endpoint.provider_id)
.bind(&endpoint.api_format)
.bind(&endpoint.api_family)
.bind(&endpoint.endpoint_kind)
.bind(endpoint.is_active)
.bind(endpoint.health_score)
.bind(&endpoint.base_url)
.bind(&endpoint.header_rules)
.bind(&endpoint.body_rules)
.bind(endpoint.max_retries)
.bind(&endpoint.custom_path)
.bind(&endpoint.config)
.bind(&endpoint.format_acceptance_config)
.bind(&endpoint.proxy)
.bind(endpoint.created_at_unix_ms.map(|value| value as f64))
.bind(endpoint.updated_at_unix_secs.map(|value| value as f64))
.execute(&self.pool)
.await
{
Ok(_) => {}
Err(error) if is_missing_endpoint_health_score_column(&error) => {
sqlx::query(
r#"
INSERT INTO provider_endpoints (
id,
provider_id,
api_format,
api_family,
endpoint_kind,
is_active,
base_url,
header_rules,
body_rules,
max_retries,
custom_path,
config,
format_acceptance_config,
proxy,
created_at,
updated_at
) VALUES (
$1,
$2,
$3,
$4,
$5,
$6,
$7,
$8,
$9,
$10,
$11,
$12,
$13,
$14,
CASE
WHEN $15::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($15::double precision)
END,
CASE
WHEN $16::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($16::double precision)
END
)
"#,
)
.bind(&endpoint.id)
.bind(&endpoint.provider_id)
.bind(&endpoint.api_format)
.bind(&endpoint.api_family)
.bind(&endpoint.endpoint_kind)
.bind(endpoint.is_active)
.bind(&endpoint.base_url)
.bind(&endpoint.header_rules)
.bind(&endpoint.body_rules)
.bind(endpoint.max_retries)
.bind(&endpoint.custom_path)
.bind(&endpoint.config)
.bind(&endpoint.format_acceptance_config)
.bind(&endpoint.proxy)
.bind(endpoint.created_at_unix_ms.map(|value| value as f64))
.bind(endpoint.updated_at_unix_secs.map(|value| value as f64))
.execute(&self.pool)
.await
.map_postgres_err()?;
}
Err(error) => return Err(postgres_error(error)),
}
self.list_endpoints_by_ids(std::slice::from_ref(&endpoint.id))
.await?
.into_iter()
.next()
.ok_or_else(|| {
DataLayerError::UnexpectedValue(format!(
"created provider catalog endpoint {} could not be reloaded",
endpoint.id
))
})
}
pub async fn update_endpoint(
&self,
endpoint: &StoredProviderCatalogEndpoint,
) -> Result<StoredProviderCatalogEndpoint, DataLayerError> {
if endpoint.id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog endpoint.id is empty".to_string(),
));
}
if endpoint.provider_id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog endpoint.provider_id is empty".to_string(),
));
}
let rows_affected = match sqlx::query(
r#"
UPDATE provider_endpoints
SET
provider_id = $2,
api_format = $3,
api_family = $4,
endpoint_kind = $5,
is_active = $6,
health_score = $7,
base_url = $8,
header_rules = $9,
body_rules = $10,
max_retries = $11,
custom_path = $12,
config = $13,
format_acceptance_config = $14,
proxy = $15,
updated_at = CASE
WHEN $16::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($16::double precision)
END
WHERE id = $1
"#,
)
.bind(&endpoint.id)
.bind(&endpoint.provider_id)
.bind(&endpoint.api_format)
.bind(&endpoint.api_family)
.bind(&endpoint.endpoint_kind)
.bind(endpoint.is_active)
.bind(endpoint.health_score)
.bind(&endpoint.base_url)
.bind(&endpoint.header_rules)
.bind(&endpoint.body_rules)
.bind(endpoint.max_retries)
.bind(&endpoint.custom_path)
.bind(&endpoint.config)
.bind(&endpoint.format_acceptance_config)
.bind(&endpoint.proxy)
.bind(endpoint.updated_at_unix_secs.map(|value| value as f64))
.execute(&self.pool)
.await
{
Ok(result) => result.rows_affected(),
Err(error) if is_missing_endpoint_health_score_column(&error) => sqlx::query(
r#"
UPDATE provider_endpoints
SET
provider_id = $2,
api_format = $3,
api_family = $4,
endpoint_kind = $5,
is_active = $6,
base_url = $7,
header_rules = $8,
body_rules = $9,
max_retries = $10,
custom_path = $11,
config = $12,
format_acceptance_config = $13,
proxy = $14,
updated_at = CASE
WHEN $15::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($15::double precision)
END
WHERE id = $1
"#,
)
.bind(&endpoint.id)
.bind(&endpoint.provider_id)
.bind(&endpoint.api_format)
.bind(&endpoint.api_family)
.bind(&endpoint.endpoint_kind)
.bind(endpoint.is_active)
.bind(&endpoint.base_url)
.bind(&endpoint.header_rules)
.bind(&endpoint.body_rules)
.bind(endpoint.max_retries)
.bind(&endpoint.custom_path)
.bind(&endpoint.config)
.bind(&endpoint.format_acceptance_config)
.bind(&endpoint.proxy)
.bind(endpoint.updated_at_unix_secs.map(|value| value as f64))
.execute(&self.pool)
.await
.map_postgres_err()?
.rows_affected(),
Err(error) => return Err(postgres_error(error)),
};
if rows_affected == 0 {
return Err(DataLayerError::UnexpectedValue(format!(
"provider catalog endpoint {} not found",
endpoint.id
)));
}
self.list_endpoints_by_ids(std::slice::from_ref(&endpoint.id))
.await?
.into_iter()
.next()
.ok_or_else(|| {
DataLayerError::UnexpectedValue(format!(
"updated provider catalog endpoint {} could not be reloaded",
endpoint.id
))
})
}
pub async fn compare_and_swap_endpoint_proxy(
&self,
update: &ProviderCatalogProxyCasUpdate,
) -> Result<bool, DataLayerError> {
if update.record_id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog endpoint_id is empty".to_string(),
));
}
let rows_affected = sqlx::query(
r#"
UPDATE provider_endpoints
SET proxy = $3, updated_at = NOW()
WHERE id = $1
AND proxy::jsonb IS NOT DISTINCT FROM $2::jsonb
"#,
)
.bind(&update.record_id)
.bind(&update.expected_proxy)
.bind(&update.proxy)
.execute(&self.pool)
.await
.map_postgres_err()?
.rows_affected();
Ok(rows_affected == 1)
}
pub async fn delete_endpoint(&self, endpoint_id: &str) -> Result<bool, DataLayerError> {
if endpoint_id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog endpoint_id is empty".to_string(),
));
}
let rows_affected = sqlx::query(
r#"
DELETE FROM provider_endpoints
WHERE id = $1
"#,
)
.bind(endpoint_id)
.execute(&self.pool)
.await
.map_postgres_err()?
.rows_affected();
Ok(rows_affected > 0)
}
pub async fn update_key(
&self,
key: &StoredProviderCatalogKey,
) -> Result<StoredProviderCatalogKey, DataLayerError> {
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!(
"provider catalog key {} not found",
key.id
)));
}
self.list_keys_by_ids(std::slice::from_ref(&key.id))
.await?
.into_iter()
.next()
.ok_or_else(|| {
DataLayerError::UnexpectedValue(format!(
"updated provider catalog key {} could not be reloaded",
key.id
))
})
}
pub async fn compare_and_swap_key_proxy(
&self,
update: &ProviderCatalogProxyCasUpdate,
) -> Result<bool, DataLayerError> {
if update.record_id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog key_id is empty".to_string(),
));
}
let rows_affected = sqlx::query(
r#"
UPDATE provider_api_keys
SET proxy = $3, updated_at = NOW()
WHERE id = $1
AND proxy::jsonb IS NOT DISTINCT FROM $2::jsonb
"#,
)
.bind(&update.record_id)
.bind(&update.expected_proxy)
.bind(&update.proxy)
.execute(&self.pool)
.await
.map_postgres_err()?
.rows_affected();
Ok(rows_affected == 1)
}
pub async fn compare_and_swap_key_credentials(
&self,
update: &ProviderCatalogKeyCredentialsCasUpdate,
) -> Result<bool, DataLayerError> {
if update.key_id.trim().is_empty() || update.expected_provider_id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog key credential CAS requires key_id and provider_id".to_string(),
));
}
let rows_affected = sqlx::query(
r#"
UPDATE provider_api_keys
SET api_key = $5, encrypted_key = NULL, auth_config = $6
WHERE id = $1
AND provider_id = $2
AND COALESCE(api_key, encrypted_key) IS NOT DISTINCT FROM $3
AND auth_config IS NOT DISTINCT FROM $4
"#,
)
.bind(&update.key_id)
.bind(&update.expected_provider_id)
.bind(update.expected_encrypted_api_key.as_deref())
.bind(update.expected_encrypted_auth_config.as_deref())
.bind(update.encrypted_api_key.as_deref())
.bind(update.encrypted_auth_config.as_deref())
.execute(&self.pool)
.await
.map_postgres_err()?
.rows_affected();
Ok(rows_affected == 1)
}
2026-08-14 09:28:07 +08:00
pub async fn compare_and_update_key_admin_state(
&self,
update: &ProviderCatalogKeyAdminCasUpdate,
) -> Result<bool, DataLayerError> {
validate_admin_key_cas_update(update)?;
let key = &update.key;
let mut builder =
QueryBuilder::<Postgres>::new("UPDATE provider_api_keys SET provider_id = ");
push_admin_key_assignments(&mut builder, key);
if update.reset_oauth_runtime {
builder.push(", oauth_invalid_at = NULL, oauth_invalid_reason = NULL, error_count = 0");
}
if let Some(rotation) = update.codex_rotation.as_ref() {
builder
.push(", upstream_metadata = jsonb_set(COALESCE(upstream_metadata, '{}'::jsonb), '{codex}', ")
.push_bind(rotation)
.push("::jsonb, true)");
}
if update.codex_rotation.is_some() || update.reset_oauth_runtime {
builder.push(", status_snapshot = ");
match (update.codex_rotation.is_some(), update.reset_oauth_runtime) {
(true, true) => builder.push("jsonb_set(jsonb_set(COALESCE(status_snapshot::jsonb, '{}'::jsonb), '{quota}', 'null'::jsonb, true), '{oauth}', 'null'::jsonb, true)::json"),
(true, false) => builder.push("jsonb_set(COALESCE(status_snapshot::jsonb, '{}'::jsonb), '{quota}', 'null'::jsonb, true)::json"),
(false, true) => builder.push("jsonb_set(COALESCE(status_snapshot::jsonb, '{}'::jsonb), '{oauth}', 'null'::jsonb, true)::json"),
(false, false) => unreachable!(),
};
}
builder
.push(" WHERE id = ")
.push_bind(&key.id)
.push(" AND api_key IS NOT DISTINCT FROM ")
.push_bind(update.expected_credential.encrypted_api_key.as_deref())
.push(" AND auth_config IS NOT DISTINCT FROM ")
.push_bind(update.expected_encrypted_auth_config.as_deref())
.push(" AND auth_type = ")
.push_bind(&update.expected_credential.auth_type)
.push(" AND provider_id = ")
.push_bind(&update.expected_credential.provider_id)
.push(
" AND EXISTS (SELECT 1 FROM providers WHERE providers.id = provider_api_keys.provider_id AND providers.provider_type = ",
)
.push_bind(&update.expected_credential.provider_type)
.push(")");
if update.codex_rotation.is_some() {
builder
.push(" AND jsonb_typeof(COALESCE(upstream_metadata, '{}'::jsonb)) = 'object'")
.push(" AND NOT (api_key IS NOT DISTINCT FROM ")
.push_bind(key.encrypted_api_key.as_deref())
.push(" AND auth_config IS NOT DISTINCT FROM ")
.push_bind(key.encrypted_auth_config.as_deref())
.push(" AND auth_type = ")
.push_bind(&key.auth_type)
.push(" AND provider_id = ")
.push_bind(&key.provider_id)
.push(")");
}
if update.codex_rotation.is_some() || update.reset_oauth_runtime {
builder.push(
" AND jsonb_typeof(COALESCE(status_snapshot::jsonb, '{}'::jsonb)) = 'object'",
);
}
let rows_affected = builder
.build()
.execute(&self.pool)
.await
.map_postgres_err()?
.rows_affected();
Ok(rows_affected > 0)
}
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()?;
let key_ids = keys.iter().map(|key| key.id.clone()).collect::<Vec<_>>();
let mut reloaded = self
.list_keys_by_ids(&key_ids)
.await?
.into_iter()
.map(|key| (key.id.clone(), key))
.collect::<BTreeMap<_, _>>();
keys.iter()
.map(|key| {
reloaded.remove(&key.id).ok_or_else(|| {
DataLayerError::UnexpectedValue(format!(
"updated provider catalog key {} could not be reloaded",
key.id
))
})
})
.collect()
}
pub async fn delete_key(&self, key_id: &str) -> Result<bool, DataLayerError> {
if key_id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog key_id is empty".to_string(),
));
}
let rows_affected = sqlx::query(
r#"
DELETE FROM provider_api_keys
WHERE id = $1
"#,
)
.bind(key_id)
.execute(&self.pool)
.await
.map_postgres_err()?
.rows_affected();
Ok(rows_affected > 0)
}
pub async fn compare_and_delete_key_oauth_credential(
&self,
delete: &ProviderCatalogKeyOAuthCredentialCasDelete,
) -> Result<bool, DataLayerError> {
let expected = &delete.expected_credential;
if delete.key_id.trim().is_empty()
|| expected
.encrypted_api_key
.as_deref()
.is_some_and(|value| value.trim().is_empty())
|| expected.auth_type.trim().is_empty()
|| expected.provider_id.trim().is_empty()
|| expected.provider_type.trim().is_empty()
2026-08-14 09:28:07 +08:00
|| delete
.expected_upstream_metadata_namespace
.as_ref()
.is_some_and(|expected| expected.namespace.trim().is_empty())
{
return Err(DataLayerError::InvalidInput(
"provider catalog OAuth credential CAS delete contains empty fields".to_string(),
));
}
let rows_affected = sqlx::query(
r#"
DELETE FROM provider_api_keys
WHERE id = $1
AND auth_config IS NOT DISTINCT FROM $2
AND api_key IS NOT DISTINCT FROM $3
AND auth_type = $4
AND provider_id = $5
AND EXISTS (
SELECT 1
FROM providers
WHERE providers.id = provider_api_keys.provider_id
AND providers.provider_type = $6
)
2026-08-14 09:28:07 +08:00
AND (
$7::boolean IS FALSE
OR (
jsonb_typeof(COALESCE(upstream_metadata, '{}'::jsonb)) = 'object'
AND (COALESCE(upstream_metadata, '{}'::jsonb) -> $8)
IS NOT DISTINCT FROM $9::jsonb
)
)
"#,
)
.bind(&delete.key_id)
.bind(delete.expected_encrypted_auth_config.as_deref())
.bind(expected.encrypted_api_key.as_deref())
.bind(&expected.auth_type)
.bind(&expected.provider_id)
.bind(&expected.provider_type)
2026-08-14 09:28:07 +08:00
.bind(delete.expected_upstream_metadata_namespace.is_some())
.bind(
delete
.expected_upstream_metadata_namespace
.as_ref()
.map(|expected| expected.namespace.as_str()),
)
.bind(
delete
.expected_upstream_metadata_namespace
.as_ref()
.and_then(|expected| expected.expected_value.as_ref()),
)
.execute(&self.pool)
.await
.map_postgres_err()?
.rows_affected();
Ok(rows_affected > 0)
}
pub async fn update_key_upstream_metadata(
&self,
key_id: &str,
upstream_metadata: Option<&serde_json::Value>,
updated_at_unix_secs: Option<u64>,
) -> Result<bool, DataLayerError> {
if key_id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog key_id is empty".to_string(),
));
}
let rows_affected = sqlx::query(
r#"
UPDATE provider_api_keys
SET
upstream_metadata = $2,
updated_at = CASE
WHEN $3::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($3::double precision)
END
WHERE id = $1
"#,
)
.bind(key_id)
.bind(upstream_metadata)
.bind(updated_at_unix_secs.map(|value| value as f64))
.execute(&self.pool)
.await
.map_postgres_err()?
.rows_affected();
Ok(rows_affected > 0)
}
pub async fn upsert_key_upstream_metadata_namespace(
&self,
key_id: &str,
namespace: &str,
value: &serde_json::Value,
updated_at_unix_secs: Option<u64>,
) -> Result<bool, DataLayerError> {
if key_id.trim().is_empty() || namespace.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog key_id and upstream metadata namespace are required".to_string(),
));
}
let rows_affected = sqlx::query(
r#"
UPDATE provider_api_keys
SET
upstream_metadata = COALESCE(upstream_metadata, '{}'::jsonb) || jsonb_build_object($2, $3::jsonb),
updated_at = CASE
WHEN $4::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($4::double precision)
END
WHERE id = $1
"#,
)
.bind(key_id)
.bind(namespace)
.bind(value)
.bind(updated_at_unix_secs.map(|value| value as f64))
.execute(&self.pool)
.await
.map_postgres_err()?
.rows_affected();
Ok(rows_affected > 0)
}
pub async fn update_key_model_fetch_state(
&self,
key_id: &str,
allowed_models: Option<&serde_json::Value>,
last_models_fetch_at_unix_secs: Option<u64>,
last_models_fetch_error: Option<&str>,
updated_at_unix_secs: Option<u64>,
) -> Result<bool, DataLayerError> {
if key_id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog key_id is empty".to_string(),
));
}
let rows_affected = sqlx::query(
r#"
UPDATE provider_api_keys
SET
allowed_models = $2,
last_models_fetch_at = CASE
WHEN $3::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($3::double precision)
END,
last_models_fetch_error = $4,
updated_at = CASE
WHEN $5::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($5::double precision)
END
WHERE id = $1
"#,
)
.bind(key_id)
.bind(allowed_models)
.bind(last_models_fetch_at_unix_secs.map(|value| value as f64))
.bind(last_models_fetch_error)
.bind(updated_at_unix_secs.map(|value| value as f64))
.execute(&self.pool)
.await
.map_postgres_err()?
.rows_affected();
Ok(rows_affected > 0)
}
pub async fn update_key_model_fetch_success(
&self,
key_id: &str,
allowed_models: Option<&serde_json::Value>,
last_models_fetch_at_unix_secs: u64,
upstream_metadata_updates: &[ProviderCatalogUpstreamMetadataNamespaceUpdate],
updated_at_unix_secs: Option<u64>,
) -> Result<bool, DataLayerError> {
if key_id.trim().is_empty()
|| upstream_metadata_updates
.iter()
.any(|update| update.namespace.trim().is_empty())
{
return Err(DataLayerError::InvalidInput(
"provider catalog key_id and upstream metadata namespaces are required".to_string(),
));
}
let mut tx = self.pool.begin().await.map_postgres_err()?;
let rows_affected = sqlx::query(
r#"
UPDATE provider_api_keys
SET
allowed_models = $2,
last_models_fetch_at = TO_TIMESTAMP($3::double precision),
last_models_fetch_error = NULL,
updated_at = CASE
WHEN $4::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($4::double precision)
END
WHERE id = $1
"#,
)
.bind(key_id)
.bind(allowed_models)
.bind(last_models_fetch_at_unix_secs as f64)
.bind(updated_at_unix_secs.map(|value| value as f64))
.execute(&mut *tx)
.await
.map_postgres_err()?
.rows_affected();
if rows_affected == 0 {
tx.rollback().await.map_postgres_err()?;
return Ok(false);
}
for update in upstream_metadata_updates {
sqlx::query(
r#"
UPDATE provider_api_keys
SET upstream_metadata = COALESCE(upstream_metadata, '{}'::jsonb)
|| jsonb_build_object($2, $3::jsonb)
WHERE id = $1
"#,
)
.bind(key_id)
.bind(&update.namespace)
.bind(&update.value)
.execute(&mut *tx)
.await
.map_postgres_err()?;
}
tx.commit().await.map_postgres_err()?;
Ok(true)
}
pub async fn update_key_health_state(
&self,
key_id: &str,
is_active: bool,
health_by_format: Option<&serde_json::Value>,
circuit_breaker_by_format: Option<&serde_json::Value>,
) -> Result<bool, DataLayerError> {
if key_id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog key_id is empty".to_string(),
));
}
let rows_affected = sqlx::query(
r#"
UPDATE provider_api_keys
SET
is_active = $2,
health_by_format = $3,
circuit_breaker_by_format = $4,
updated_at = NOW()
WHERE id = $1
"#,
)
.bind(key_id)
.bind(is_active)
.bind(health_by_format)
.bind(circuit_breaker_by_format)
.execute(&self.pool)
.await
.map_postgres_err()?
.rows_affected();
Ok(rows_affected > 0)
}
pub async fn reset_key_error_count(&self, key_id: &str) -> Result<bool, DataLayerError> {
if key_id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog key_id is empty".to_string(),
));
}
let rows_affected = sqlx::query(
r#"
UPDATE provider_api_keys
SET error_count = 0, updated_at = NOW()
WHERE id = $1
"#,
)
.bind(key_id)
.execute(&self.pool)
.await
.map_postgres_err()?
.rows_affected();
Ok(rows_affected > 0)
}
pub async fn compare_and_update_key_adaptive_state(
&self,
update: &ProviderCatalogKeyAdaptiveStateUpdate,
) -> Result<bool, DataLayerError> {
if update.key_id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog key_id is empty".to_string(),
));
}
let status_snapshot_patch = adaptive_status_snapshot_patch(&update.status_snapshot_patch)?;
let expected = update.expected.canonicalized();
let next = update.next.canonicalized();
let rows_affected = sqlx::query(
r#"
UPDATE provider_api_keys
SET
learned_rpm_limit = $2,
rpm_429_count = $3,
last_429_at = CASE
WHEN $4::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($4::double precision)
END,
last_429_type = $5,
adjustment_history = $6,
utilization_samples = $7,
last_probe_increase_at = CASE
WHEN $8::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($8::double precision)
END,
last_rpm_peak = $9,
concurrent_429_count = $10,
status_snapshot = (COALESCE(status_snapshot::jsonb, '{}'::jsonb) || $11::jsonb)::json,
updated_at = CASE
WHEN $12::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($12::double precision)
END
WHERE id = $1
AND learned_rpm_limit IS NOT DISTINCT FROM $13
AND rpm_429_count IS NOT DISTINCT FROM $14
AND CAST(EXTRACT(EPOCH FROM last_429_at) AS BIGINT) IS NOT DISTINCT FROM $15
AND last_429_type IS NOT DISTINCT FROM $16
AND adjustment_history::jsonb IS NOT DISTINCT FROM $17::jsonb
AND utilization_samples::jsonb IS NOT DISTINCT FROM $18::jsonb
AND CAST(EXTRACT(EPOCH FROM last_probe_increase_at) AS BIGINT) IS NOT DISTINCT FROM $19
AND last_rpm_peak IS NOT DISTINCT FROM $20
AND concurrent_429_count IS NOT DISTINCT FROM $21
AND ($22::text IS NULL OR auth_config IS NOT DISTINCT FROM $22)
"#,
)
.bind(&update.key_id)
.bind(next.learned_rpm_limit.map(|value| value as i32))
.bind(next.rpm_429_count.map(|value| value as i32))
.bind(next.last_429_at_unix_secs.map(|value| value as f64))
.bind(&next.last_429_type)
.bind(&next.adjustment_history)
.bind(&next.utilization_samples)
.bind(
next.last_probe_increase_at_unix_secs
.map(|value| value as f64),
)
.bind(next.last_rpm_peak.map(|value| value as i32))
.bind(next.concurrent_429_count.map(|value| value as i32))
.bind(&status_snapshot_patch)
.bind(update.updated_at_unix_secs.map(|value| value as f64))
.bind(expected.learned_rpm_limit.map(|value| value as i32))
.bind(expected.rpm_429_count.map(|value| value as i32))
.bind(
expected
.last_429_at_unix_secs
.and_then(|value| i64::try_from(value).ok()),
)
.bind(&expected.last_429_type)
.bind(&expected.adjustment_history)
.bind(&expected.utilization_samples)
.bind(
expected
.last_probe_increase_at_unix_secs
.and_then(|value| i64::try_from(value).ok()),
)
.bind(expected.last_rpm_peak.map(|value| value as i32))
.bind(expected.concurrent_429_count.map(|value| value as i32))
.bind(update.expected_encrypted_auth_config.as_deref())
.execute(&self.pool)
.await
.map_postgres_err()?
.rows_affected();
Ok(rows_affected > 0)
}
pub async fn update_key_runtime_metadata(
&self,
update: &ProviderCatalogKeyRuntimeMetadataUpdate,
) -> Result<bool, DataLayerError> {
validate_runtime_metadata_update(update)?;
let mut tx = self.pool.begin().await.map_postgres_err()?;
if !lock_runtime_metadata_namespace_matches(
&mut tx,
&update.key_id,
&update.namespace,
update.expected_upstream_metadata_value.as_ref(),
)
.await?
{
tx.rollback().await.map_postgres_err()?;
return Ok(false);
}
let rows_affected = sqlx::query(KEY_RUNTIME_METADATA_UPDATE_SQL)
.bind(&update.key_id)
.bind(&update.namespace)
.bind(&update.upstream_metadata_value)
.bind(&update.status_snapshot_patch)
.bind(update.updated_at_unix_secs.map(|value| value as f64))
.execute(&mut *tx)
.await
.map_postgres_err()?
.rows_affected();
if rows_affected == 0 {
tx.rollback().await.map_postgres_err()?;
return Ok(false);
}
tx.commit().await.map_postgres_err()?;
Ok(true)
}
pub async fn update_key_status_snapshot(
&self,
update: &ProviderCatalogKeyStatusSnapshotUpdate,
) -> Result<bool, DataLayerError> {
if update.key_id.trim().is_empty() || !update.status_snapshot_patch.is_object() {
return Err(DataLayerError::InvalidInput(
"provider catalog key_id and object status snapshot patch are required".to_string(),
));
}
let rows_affected = sqlx::query(
r#"
UPDATE provider_api_keys
SET
status_snapshot = (COALESCE(status_snapshot::jsonb, '{}'::jsonb) || $2::jsonb)::json,
updated_at = CASE
WHEN $3::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($3::double precision)
END
WHERE id = $1
"#,
)
.bind(&update.key_id)
.bind(&update.status_snapshot_patch)
.bind(update.updated_at_unix_secs.map(|value| value as f64))
.execute(&self.pool)
.await
.map_postgres_err()?
.rows_affected();
Ok(rows_affected > 0)
}
pub async fn compare_and_update_key_health_state(
&self,
update: &ProviderCatalogKeyHealthStateUpdate,
) -> Result<bool, DataLayerError> {
if update.key_id.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"provider catalog key_id is empty".to_string(),
));
}
let rows_affected = sqlx::query(KEY_RUNTIME_HEALTH_CAS_SQL)
.bind(&update.key_id)
.bind(&update.health_by_format)
.bind(&update.circuit_breaker_by_format)
.bind(&update.expected_health_by_format)
.bind(&update.expected_circuit_breaker_by_format)
.bind(update.expected_encrypted_auth_config.as_deref())
.execute(&self.pool)
.await
.map_postgres_err()?
.rows_affected();
Ok(rows_affected > 0)
}
}
#[async_trait]
impl ProviderCatalogReadRepository for SqlxProviderCatalogReadRepository {
async fn list_providers(
&self,
active_only: bool,
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
Self::list_providers(self, active_only).await
}
async fn list_providers_by_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
Self::list_providers_by_ids(self, provider_ids).await
}
async fn list_endpoints_by_ids(
&self,
endpoint_ids: &[String],
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
Self::list_endpoints_by_ids(self, endpoint_ids).await
}
async fn list_endpoints_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
Self::list_endpoints_by_provider_ids(self, provider_ids).await
}
async fn list_keys_by_ids(
&self,
key_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
Self::list_keys_by_ids(self, key_ids).await
}
async fn list_keys_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
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_key_maintenance_summaries_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKeyMaintenanceSummary>, DataLayerError> {
Self::list_key_maintenance_summaries_by_provider_ids(self, provider_ids).await
}
async fn list_keys_page(
&self,
query: &ProviderCatalogKeyListQuery,
) -> Result<StoredProviderCatalogKeyPage, DataLayerError> {
Self::list_keys_page(self, query).await
}
async fn list_key_stats_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKeyStats>, DataLayerError> {
Self::list_key_stats_by_provider_ids(self, provider_ids).await
}
}
#[async_trait]
impl ProviderCatalogWriteRepository for SqlxProviderCatalogReadRepository {
async fn create_provider(
&self,
provider: &StoredProviderCatalogProvider,
shift_existing_priorities_from: Option<i32>,
) -> Result<StoredProviderCatalogProvider, DataLayerError> {
Self::create_provider(self, provider, shift_existing_priorities_from).await
}
async fn update_provider(
&self,
provider: &StoredProviderCatalogProvider,
) -> Result<StoredProviderCatalogProvider, DataLayerError> {
Self::update_provider(self, provider).await
}
async fn compare_and_swap_provider_config(
&self,
update: &ProviderCatalogProviderConfigCasUpdate,
) -> Result<bool, DataLayerError> {
Self::compare_and_swap_provider_config(self, update).await
}
async fn compare_and_swap_provider_proxy(
&self,
update: &ProviderCatalogProxyCasUpdate,
) -> Result<bool, DataLayerError> {
Self::compare_and_swap_provider_proxy(self, update).await
}
async fn delete_provider(&self, provider_id: &str) -> Result<bool, DataLayerError> {
Self::delete_provider(self, provider_id).await
}
async fn cleanup_deleted_provider_refs(
&self,
provider_id: &str,
2026-06-12 16:25:11 +08:00
provider_deleted: bool,
endpoint_ids: &[String],
key_ids: &[String],
) -> Result<(), DataLayerError> {
2026-06-12 16:25:11 +08:00
Self::cleanup_deleted_provider_refs(
self,
provider_id,
provider_deleted,
endpoint_ids,
key_ids,
)
.await
}
async fn create_endpoint(
&self,
endpoint: &StoredProviderCatalogEndpoint,
) -> Result<StoredProviderCatalogEndpoint, DataLayerError> {
self.create_endpoint(endpoint).await
}
async fn update_endpoint(
&self,
endpoint: &StoredProviderCatalogEndpoint,
) -> Result<StoredProviderCatalogEndpoint, DataLayerError> {
self.update_endpoint(endpoint).await
}
async fn compare_and_swap_endpoint_proxy(
&self,
update: &ProviderCatalogProxyCasUpdate,
) -> Result<bool, DataLayerError> {
Self::compare_and_swap_endpoint_proxy(self, update).await
}
async fn delete_endpoint(&self, endpoint_id: &str) -> Result<bool, DataLayerError> {
Self::delete_endpoint(self, endpoint_id).await
}
async fn create_key(
&self,
key: &StoredProviderCatalogKey,
) -> Result<StoredProviderCatalogKey, DataLayerError> {
Self::create_key(self, key).await
}
async fn update_key(
&self,
key: &StoredProviderCatalogKey,
) -> Result<StoredProviderCatalogKey, DataLayerError> {
Self::update_key(self, key).await
}
async fn compare_and_swap_key_proxy(
&self,
update: &ProviderCatalogProxyCasUpdate,
) -> Result<bool, DataLayerError> {
Self::compare_and_swap_key_proxy(self, update).await
}
async fn compare_and_swap_key_credentials(
&self,
update: &ProviderCatalogKeyCredentialsCasUpdate,
) -> Result<bool, DataLayerError> {
Self::compare_and_swap_key_credentials(self, update).await
}
2026-08-14 09:28:07 +08:00
async fn compare_and_update_key_admin_state(
&self,
update: &ProviderCatalogKeyAdminCasUpdate,
) -> Result<bool, DataLayerError> {
Self::compare_and_update_key_admin_state(self, update).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,
upstream_metadata: Option<&serde_json::Value>,
updated_at_unix_secs: Option<u64>,
) -> Result<bool, DataLayerError> {
Self::update_key_upstream_metadata(self, key_id, upstream_metadata, updated_at_unix_secs)
.await
}
async fn upsert_key_upstream_metadata_namespace(
&self,
key_id: &str,
namespace: &str,
value: &serde_json::Value,
updated_at_unix_secs: Option<u64>,
) -> Result<bool, DataLayerError> {
Self::upsert_key_upstream_metadata_namespace(
self,
key_id,
namespace,
value,
updated_at_unix_secs,
)
.await
}
async fn update_key_model_fetch_state(
&self,
key_id: &str,
allowed_models: Option<&serde_json::Value>,
last_models_fetch_at_unix_secs: Option<u64>,
last_models_fetch_error: Option<&str>,
updated_at_unix_secs: Option<u64>,
) -> Result<bool, DataLayerError> {
Self::update_key_model_fetch_state(
self,
key_id,
allowed_models,
last_models_fetch_at_unix_secs,
last_models_fetch_error,
updated_at_unix_secs,
)
.await
}
async fn update_key_model_fetch_success(
&self,
key_id: &str,
allowed_models: Option<&serde_json::Value>,
last_models_fetch_at_unix_secs: u64,
upstream_metadata_updates: &[ProviderCatalogUpstreamMetadataNamespaceUpdate],
updated_at_unix_secs: Option<u64>,
) -> Result<bool, DataLayerError> {
Self::update_key_model_fetch_success(
self,
key_id,
allowed_models,
last_models_fetch_at_unix_secs,
upstream_metadata_updates,
updated_at_unix_secs,
)
.await
}
async fn delete_key(&self, key_id: &str) -> Result<bool, DataLayerError> {
Self::delete_key(self, key_id).await
}
async fn compare_and_delete_key_oauth_credential(
&self,
delete: &ProviderCatalogKeyOAuthCredentialCasDelete,
) -> Result<bool, DataLayerError> {
Self::compare_and_delete_key_oauth_credential(self, delete).await
}
async fn clear_key_oauth_invalid_marker(&self, key_id: &str) -> Result<bool, DataLayerError> {
Self::clear_key_oauth_invalid_marker(self, key_id).await
}
async fn update_key_oauth_runtime_state(
&self,
key_id: &str,
oauth_invalid_at_unix_secs: Option<u64>,
oauth_invalid_reason: Option<&str>,
updated_at_unix_secs: Option<u64>,
) -> Result<bool, DataLayerError> {
Self::update_key_oauth_runtime_state(
self,
key_id,
oauth_invalid_at_unix_secs,
oauth_invalid_reason,
updated_at_unix_secs,
)
.await
}
async fn compare_and_update_key_oauth_runtime_state(
&self,
update: &ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
) -> Result<bool, DataLayerError> {
Self::compare_and_update_key_oauth_runtime_state(self, update).await
}
async fn update_key_health_state(
&self,
key_id: &str,
is_active: bool,
health_by_format: Option<&serde_json::Value>,
circuit_breaker_by_format: Option<&serde_json::Value>,
) -> Result<bool, DataLayerError> {
Self::update_key_health_state(
self,
key_id,
is_active,
health_by_format,
circuit_breaker_by_format,
)
.await
}
async fn reset_key_error_count(&self, key_id: &str) -> Result<bool, DataLayerError> {
Self::reset_key_error_count(self, key_id).await
}
async fn compare_and_update_key_adaptive_state(
&self,
update: &ProviderCatalogKeyAdaptiveStateUpdate,
) -> Result<bool, DataLayerError> {
Self::compare_and_update_key_adaptive_state(self, update).await
}
async fn update_key_runtime_metadata(
&self,
update: &ProviderCatalogKeyRuntimeMetadataUpdate,
) -> Result<bool, DataLayerError> {
Self::update_key_runtime_metadata(self, update).await
}
async fn update_key_status_snapshot(
&self,
update: &ProviderCatalogKeyStatusSnapshotUpdate,
) -> Result<bool, DataLayerError> {
Self::update_key_status_snapshot(self, update).await
}
async fn compare_and_update_key_health_state(
&self,
update: &ProviderCatalogKeyHealthStateUpdate,
) -> Result<bool, DataLayerError> {
Self::compare_and_update_key_health_state(self, update).await
}
}
fn build_list_query<'a>(
prefix: &'static str,
ids: &'a [String],
suffix: &'static str,
) -> QueryBuilder<'a, Postgres> {
let mut builder = QueryBuilder::<Postgres>::new(select_prefix_for_in(prefix));
let mut where_clause = WhereClause::new();
push_in(
&mut builder,
&mut where_clause,
in_column_for_prefix(prefix),
ids,
);
builder.push(suffix);
builder
}
fn select_prefix_for_in(prefix: &'static str) -> &'static str {
prefix
.rsplit_once("\nWHERE ")
.map(|(select_prefix, _)| select_prefix)
.expect("provider catalog IN query prefix must contain WHERE")
}
fn in_column_for_prefix(prefix: &'static str) -> &'static str {
prefix
.rsplit_once("\nWHERE ")
.and_then(|(_, predicate)| predicate.trim().strip_suffix("IN ("))
.map(str::trim)
.expect("provider catalog IN query prefix must end with IN (")
}
fn apply_key_page_filters<'a>(
builder: &mut QueryBuilder<'a, Postgres>,
where_clause: &mut WhereClause,
query: &'a ProviderCatalogKeyListQuery,
) {
push_eq(
builder,
where_clause,
"provider_id",
query.provider_id.clone(),
);
if let Some(search) = query.search.as_deref() {
push_ci_contains_any(
builder,
where_clause,
SqlDialect::Postgres,
&["name", "id"],
search,
);
}
push_optional_eq(builder, where_clause, "is_active", query.is_active);
}
fn row_get<T>(row: &PgRow, column: &str) -> Result<T, DataLayerError>
where
for<'r> T: sqlx::Decode<'r, sqlx::Postgres> + sqlx::Type<sqlx::Postgres>,
{
row.try_get(column).map_postgres_err()
}
fn optional_u64(value: Option<i64>, field_name: &str) -> Result<Option<u64>, DataLayerError> {
value
.map(|value| {
u64::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!("invalid {field_name}: {value}"))
})
})
.transpose()
}
fn map_provider_row(row: &PgRow) -> Result<StoredProviderCatalogProvider, DataLayerError> {
let quota_reset_day = row_get::<Option<i32>>(row, "quota_reset_day")?
.map(|value| {
u64::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid providers.quota_reset_day: {value}"
))
})
})
.transpose()?;
let created_at_unix_ms = row_get::<Option<i64>>(row, "created_at_unix_ms")?
.map(|value| {
u64::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid providers.created_at_unix_ms: {value}"
))
})
})
.transpose()?;
let updated_at_unix_secs = row_get::<Option<i64>>(row, "updated_at_unix_secs")?
.map(|value| {
u64::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid providers.updated_at_unix_secs: {value}"
))
})
})
.transpose()?;
Ok(StoredProviderCatalogProvider::new(
row_get(row, "id")?,
row_get(row, "name")?,
row_get(row, "website")?,
row_get(row, "provider_type")?,
)?
.with_description(row_get(row, "description")?)
.with_billing_fields(
row_get(row, "billing_type")?,
row_get(row, "monthly_quota_usd")?,
row_get(row, "monthly_used_usd")?,
quota_reset_day,
optional_u64(
row_get(row, "quota_last_reset_at_unix_secs")?,
"providers.quota_last_reset_at",
)?,
optional_u64(
row_get(row, "quota_expires_at_unix_secs")?,
"providers.quota_expires_at",
)?,
)
.with_routing_fields(row_get(row, "provider_priority")?)
.with_transport_fields(
row_get(row, "is_active")?,
row_get(row, "keep_priority_on_conversion")?,
row_get(row, "enable_format_conversion")?,
row_get(row, "concurrent_limit")?,
row_get(row, "max_retries")?,
row_get(row, "proxy")?,
row_get(row, "request_timeout")?,
row_get(row, "stream_first_byte_timeout")?,
row_get(row, "config")?,
)
.with_timestamps(created_at_unix_ms, updated_at_unix_secs))
}
fn map_endpoint_row(row: &PgRow) -> Result<StoredProviderCatalogEndpoint, DataLayerError> {
let created_at_unix_ms = row_get::<Option<i64>>(row, "created_at_unix_ms")?
.map(|value| {
u64::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid provider_endpoints.created_at_unix_ms: {value}"
))
})
})
.transpose()?;
let updated_at_unix_secs = row_get::<Option<i64>>(row, "updated_at_unix_secs")?
.map(|value| {
u64::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid provider_endpoints.updated_at_unix_secs: {value}"
))
})
})
.transpose()?;
StoredProviderCatalogEndpoint::new(
row_get(row, "id")?,
row_get(row, "provider_id")?,
row_get(row, "api_format")?,
row_get(row, "api_family")?,
row_get(row, "endpoint_kind")?,
row_get(row, "is_active")?,
)?
.with_timestamps(created_at_unix_ms, updated_at_unix_secs)
.with_health_score(
row.try_get::<Option<f64>, _>("health_score")
.ok()
.flatten()
.unwrap_or(1.0),
)
.with_transport_fields(
row_get(row, "base_url")?,
row_get(row, "header_rules")?,
row_get(row, "body_rules")?,
row_get(row, "max_retries")?,
row_get(row, "custom_path")?,
row_get(row, "config")?,
row_get(row, "format_acceptance_config")?,
row_get(row, "proxy")?,
)
}
fn is_missing_endpoint_health_score_column(error: &sqlx::Error) -> bool {
error
.as_database_error()
.and_then(|db| db.code())
.is_some_and(|code| code == "42703")
&& error
.as_database_error()
.map(|db| db.message().contains("health_score"))
.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")?,
row_get(row, "total_keys")?,
row_get(row, "active_keys")?,
)
}
fn map_key_maintenance_summary_row(
row: &PgRow,
) -> Result<StoredProviderCatalogKeyMaintenanceSummary, DataLayerError> {
Ok(StoredProviderCatalogKeyMaintenanceSummary {
id: row_get(row, "id")?,
provider_id: row_get(row, "provider_id")?,
is_active: row_get(row, "is_active")?,
upstream_metadata: row_get(row, "upstream_metadata")?,
})
}
fn map_key_row(row: &PgRow) -> Result<StoredProviderCatalogKey, DataLayerError> {
let expires_at_unix_secs = optional_u64(
row_get(row, "expires_at_unix_secs")?,
"provider_api_keys.expires_at",
)?;
let rpm_limit = row_get::<Option<i32>>(row, "rpm_limit")?
.map(|value| {
u32::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid provider_api_keys.rpm_limit: {value}"
))
})
})
.transpose()?;
let concurrent_limit = row_get::<Option<i32>>(row, "concurrent_limit")?;
let learned_rpm_limit = row_get::<Option<i32>>(row, "learned_rpm_limit")?
.map(|value| {
u32::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid provider_api_keys.learned_rpm_limit: {value}"
))
})
})
.transpose()?;
let concurrent_429_count = row_get::<Option<i32>>(row, "concurrent_429_count")?
.map(|value| {
u32::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid provider_api_keys.concurrent_429_count: {value}"
))
})
})
.transpose()?;
let rpm_429_count = row_get::<Option<i32>>(row, "rpm_429_count")?
.map(|value| {
u32::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid provider_api_keys.rpm_429_count: {value}"
))
})
})
.transpose()?;
2026-05-19 00:58:09 +08:00
let request_count = row_get::<Option<i64>>(row, "request_count")?
.map(|value| {
u32::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid provider_api_keys.request_count: {value}"
))
})
})
.transpose()?;
let total_tokens = row_get::<Option<i64>>(row, "total_tokens")?
.unwrap_or(0)
.try_into()
.map_err(|_| {
DataLayerError::UnexpectedValue("invalid provider_api_keys.total_tokens".to_string())
})?;
let total_cost_usd = row_get::<Option<f64>>(row, "total_cost_usd")?.unwrap_or(0.0);
if !total_cost_usd.is_finite() {
return Err(DataLayerError::UnexpectedValue(
"invalid provider_api_keys.total_cost_usd".to_string(),
));
}
2026-05-19 00:58:09 +08:00
let success_count = row_get::<Option<i64>>(row, "success_count")?
.map(|value| {
u32::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid provider_api_keys.success_count: {value}"
))
})
})
.transpose()?;
2026-05-19 00:58:09 +08:00
let error_count = row_get::<Option<i64>>(row, "error_count")?
.map(|value| {
u32::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid provider_api_keys.error_count: {value}"
))
})
})
.transpose()?;
2026-05-19 00:58:09 +08:00
let total_response_time_ms = row_get::<Option<i64>>(row, "total_response_time_ms")?
.map(|value| {
2026-07-09 14:34:09 +08:00
u64::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid provider_api_keys.total_response_time_ms: {value}"
))
})
})
.transpose()?;
let last_probe_increase_at_unix_secs =
row_get::<Option<i64>>(row, "last_probe_increase_at_unix_secs")?
.map(|value| {
u64::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid provider_api_keys.last_probe_increase_at_unix_secs: {value}"
))
})
})
.transpose()?;
let last_rpm_peak = row_get::<Option<i32>>(row, "last_rpm_peak")?
.map(|value| {
u32::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid provider_api_keys.last_rpm_peak: {value}"
))
})
})
.transpose()?;
let last_models_fetch_at_unix_secs =
row_get::<Option<i64>>(row, "last_models_fetch_at_unix_secs")?
.map(|value| {
u64::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid provider_api_keys.last_models_fetch_at_unix_secs: {value}"
))
})
})
.transpose()?;
let oauth_invalid_at_unix_secs = row_get::<Option<i64>>(row, "oauth_invalid_at_unix_secs")?
.map(|value| {
u64::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid provider_api_keys.oauth_invalid_at_unix_secs: {value}"
))
})
})
.transpose()?;
let last_used_at_unix_secs = row_get::<Option<i64>>(row, "last_used_at_unix_secs")?
.map(|value| {
u64::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid provider_api_keys.last_used_at_unix_secs: {value}"
))
})
})
.transpose()?;
let created_at_unix_ms = row_get::<Option<i64>>(row, "created_at_unix_ms")?
.map(|value| {
u64::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid provider_api_keys.created_at_unix_ms: {value}"
))
})
})
.transpose()?;
let updated_at_unix_secs = row_get::<Option<i64>>(row, "updated_at_unix_secs")?
.map(|value| {
u64::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid provider_api_keys.updated_at_unix_secs: {value}"
))
})
})
.transpose()?;
let auth_type_by_format: Option<serde_json::Value> = row_get(row, "auth_type_by_format")?;
let allow_auth_channel_mismatch_formats: Option<serde_json::Value> =
row_get(row, "allow_auth_channel_mismatch_formats")?;
StoredProviderCatalogKey::new(
row_get(row, "id")?,
row_get(row, "provider_id")?,
row_get(row, "name")?,
row_get(row, "auth_type")?,
row_get(row, "capabilities")?,
row_get(row, "is_active")?,
)?
.with_transport_fields(
row_get(row, "api_formats")?,
2026-04-29 15:46:50 +08:00
row_get::<Option<String>>(row, "api_key")?,
row_get(row, "auth_config")?,
row_get(row, "rate_multipliers")?,
row_get(row, "global_priority_by_format")?,
row_get(row, "allowed_models")?,
expires_at_unix_secs,
row_get(row, "proxy")?,
row_get(row, "fingerprint")?,
)
.map(|key| {
let mut key = key
.with_rate_limit_fields(
rpm_limit,
concurrent_limit,
learned_rpm_limit,
concurrent_429_count,
rpm_429_count,
row.try_get::<Option<i64>, _>("last_429_at_unix_secs")
.ok()
.flatten()
.and_then(|value| u64::try_from(value).ok()),
row.try_get("adjustment_history").ok(),
request_count,
success_count,
)
.with_usage_fields(error_count, total_response_time_ms)
.with_usage_totals(total_tokens, total_cost_usd)
.with_health_fields(
row.try_get("health_by_format").ok(),
row.try_get("circuit_breaker_by_format").ok(),
);
key.note = row.try_get("note").ok();
key.internal_priority = row.try_get("internal_priority").unwrap_or(50);
key.cache_ttl_minutes = row.try_get("cache_ttl_minutes").unwrap_or(5);
key.max_probe_interval_minutes = row.try_get("max_probe_interval_minutes").unwrap_or(32);
key.last_429_type = row.try_get("last_429_type").ok();
key.utilization_samples = row.try_get("utilization_samples").ok();
key.last_probe_increase_at_unix_secs = last_probe_increase_at_unix_secs;
key.last_rpm_peak = last_rpm_peak;
key.last_used_at_unix_secs = last_used_at_unix_secs;
key.auto_fetch_models = row.try_get("auto_fetch_models").unwrap_or(false);
key.last_models_fetch_at_unix_secs = last_models_fetch_at_unix_secs;
key.last_models_fetch_error = row.try_get("last_models_fetch_error").ok();
key.locked_models = row.try_get("locked_models").ok();
key.model_include_patterns = row.try_get("model_include_patterns").ok();
key.model_exclude_patterns = row.try_get("model_exclude_patterns").ok();
key.upstream_metadata = row.try_get("upstream_metadata").ok();
key.oauth_invalid_at_unix_secs = oauth_invalid_at_unix_secs;
key.oauth_invalid_reason = row.try_get("oauth_invalid_reason").ok();
key.status_snapshot = row.try_get("status_snapshot").ok();
key.created_at_unix_ms = created_at_unix_ms;
key.updated_at_unix_secs = updated_at_unix_secs;
key
})
.and_then(|key| {
key.with_auth_channel_policy_fields(
auth_type_by_format,
allow_auth_channel_mismatch_formats,
)
})
}
#[cfg(test)]
mod tests {
use aether_data_contracts::repository::provider_catalog::ProviderCatalogKeyRuntimeMetadataUpdate;
use serde_json::json;
use super::{optional_u64, SqlxProviderCatalogReadRepository};
use crate::{PostgresPoolConfig, PostgresPoolFactory};
#[tokio::test]
async fn repository_constructs_from_lazy_pool() {
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
})
.expect("factory should build");
let pool = factory.connect_lazy().expect("pool should build");
let repository = SqlxProviderCatalogReadRepository::new(pool);
let _ = repository.pool();
}
#[test]
fn provider_catalog_negative_security_timestamps_fail_closed() {
assert!(optional_u64(Some(-1), "provider_api_keys.expires_at").is_err());
assert_eq!(
optional_u64(None, "provider_api_keys.expires_at")
.expect("SQL NULL should remain optional"),
None
);
}
#[test]
fn key_queries_include_usage_totals() {
for sql in [
super::LIST_KEYS_BY_IDS_PREFIX,
super::LIST_KEYS_BY_PROVIDER_IDS_PREFIX,
] {
assert!(sql.contains("total_tokens"));
assert!(sql.contains("total_cost_usd"));
}
}
#[test]
fn provider_api_keys_concurrent_limit_queries_include_field() {
for sql in [
super::LIST_KEYS_BY_IDS_PREFIX,
super::LIST_KEYS_BY_PROVIDER_IDS_PREFIX,
super::LIST_KEY_SUMMARIES_BY_PROVIDER_IDS_PREFIX,
] {
assert!(sql.contains("concurrent_limit"));
}
let source = include_str!("provider_catalog.rs").replace("\r\n", "\n");
assert!(source.contains("concurrent_limit,"));
assert!(source.contains("concurrent_limit = $13"));
assert!(source.contains(".bind(key.concurrent_limit)"));
assert!(source.contains("row_get::<Option<i32>>(row, \"concurrent_limit\")"));
}
#[test]
fn provider_api_keys_concurrent_limit_schema_is_nullable_without_default() {
let migration = include_str!(
"../migrations/20260502000000_add_provider_key_auth_channel_mismatch_formats.sql"
);
let concurrent_limit_line = migration
.lines()
.find(|line| line.contains("ADD COLUMN IF NOT EXISTS concurrent_limit"))
.expect("concurrent_limit migration line should exist")
.to_ascii_lowercase();
assert_eq!(
concurrent_limit_line.trim(),
"add column if not exists concurrent_limit integer;"
);
let baseline = include_str!("../migrations/20260403000000_baseline.sql");
assert!(baseline.contains("CREATE TABLE IF NOT EXISTS public.provider_api_keys"));
assert!(baseline.contains("concurrent_limit integer,"));
}
#[test]
fn provider_api_keys_auth_channel_mismatch_list_queries_include_field() {
for sql in [
super::LIST_KEYS_BY_IDS_PREFIX,
super::LIST_KEYS_BY_PROVIDER_IDS_PREFIX,
] {
assert!(sql.contains("allow_auth_channel_mismatch_formats"));
}
let source = include_str!("provider_catalog.rs").replace("\r\n", "\n");
assert!(
source
2026-05-06 00:20:06 +08:00
.matches(
"auth_type_by_format,\n allow_auth_channel_mismatch_formats,\n COALESCE(api_key, encrypted_key) AS api_key",
)
.count()
>= 2
);
assert!(source.contains("QueryBuilder::<Postgres>::new(select_prefix_for_in("));
assert!(source.contains(".bind(&key.allow_auth_channel_mismatch_formats)"));
assert!(source.contains("row_get(row, \"allow_auth_channel_mismatch_formats\")?"));
}
2026-05-04 12:21:54 +08:00
#[test]
fn provider_api_keys_create_key_insert_placeholders_match_bind_order() {
let source = include_str!("provider_catalog.rs").replace("\r\n", "\n");
2026-05-04 12:21:54 +08:00
assert!(source.contains(
" $24,\n $25,\n $26,\n CASE\n WHEN $27::double precision IS NULL THEN NULL"
));
assert!(source.contains(" $29,\n $30,\n COALESCE($31, 0),"));
assert!(source.contains(
" COALESCE($42, 0),\n COALESCE($43, 0),\n COALESCE($44, 0),\n CASE\n WHEN $45::double precision IS NULL THEN NULL"
));
assert!(source.contains(
" CASE\n WHEN $52::double precision IS NULL THEN NOW()\n ELSE TO_TIMESTAMP($52::double precision)\n END,\n $53"
));
}
#[test]
fn runtime_health_cas_never_updates_admin_activation() {
let sql = super::KEY_RUNTIME_HEALTH_CAS_SQL.to_ascii_lowercase();
assert!(sql.contains("health_by_format = $2"));
assert!(sql.contains("circuit_breaker_by_format = $3"));
assert!(sql.contains("health_by_format::jsonb is not distinct from $4::jsonb"));
assert!(!sql.contains("is_active"));
}
#[test]
fn runtime_metadata_cas_locks_only_the_requested_namespace() {
let lock_sql = super::KEY_RUNTIME_METADATA_NAMESPACE_LOCK_SQL.to_ascii_lowercase();
let update_sql = super::KEY_RUNTIME_METADATA_UPDATE_SQL.to_ascii_lowercase();
assert!(lock_sql.contains("upstream_metadata, '{}'::jsonb) -> $2"));
assert!(lock_sql.contains("upstream_metadata, '{}'::jsonb) ? $2"));
assert!(lock_sql.contains("for update"));
assert!(update_sql
.contains("jsonb_typeof(coalesce(upstream_metadata, '{}'::jsonb)) = 'object'"));
assert!(update_sql.contains("status_snapshot::jsonb"));
assert!(!update_sql.contains("is_active"));
}
#[test]
fn runtime_metadata_namespace_cas_distinguishes_missing_from_json_null() {
assert!(super::runtime_metadata_namespace_matches(
true, false, None, None,
));
assert!(!super::runtime_metadata_namespace_matches(
true,
true,
Some(&serde_json::Value::Null),
None,
));
assert!(super::runtime_metadata_namespace_matches(
true,
true,
Some(&serde_json::Value::Null),
Some(&serde_json::Value::Null),
));
assert!(!super::runtime_metadata_namespace_matches(
false, false, None, None,
));
}
#[tokio::test]
#[ignore = "requires AETHER_TEST_DATABASE_URL and PostgreSQL migrations"]
async fn live_runtime_metadata_cas_handles_high_precision_jsonb_numbers() {
let database_url = std::env::var("AETHER_TEST_DATABASE_URL")
.expect("AETHER_TEST_DATABASE_URL must point at the test database");
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
database_url,
min_connections: 1,
max_connections: 2,
acquire_timeout_ms: 10_000,
idle_timeout_ms: 30_000,
max_lifetime_ms: 60_000,
statement_cache_capacity: 64,
require_ssl: false,
})
.expect("factory should build");
let repository = SqlxProviderCatalogReadRepository::new(
factory.connect_lazy().expect("lazy pool should build"),
);
crate::run_migrations(repository.pool())
.await
.expect("test database migrations should succeed");
let suffix = uuid::Uuid::new_v4().simple().to_string();
let provider_id = uuid::Uuid::new_v4().to_string();
let key_id = uuid::Uuid::new_v4().to_string();
let provider_name = format!("provider-metadata-cas-{suffix}");
let key_name = format!("key-metadata-cas-{suffix}");
sqlx::query(
"INSERT INTO providers (id, name, provider_type) VALUES ($1, $2, 'antigravity')",
)
.bind(&provider_id)
.bind(&provider_name)
.execute(repository.pool())
.await
.expect("provider fixture should insert");
sqlx::query(
r#"
INSERT INTO provider_api_keys (
id, name, provider_id, total_tokens, total_cost_usd, upstream_metadata
)
VALUES ($1, $2, $3, 0, 0, $4::jsonb)
"#,
)
.bind(&key_id)
.bind(&key_name)
.bind(&provider_id)
.bind(r#"{"antigravity":{"used_percent":0.123456789012345678901234567890}}"#)
.execute(repository.pool())
.await
.expect("provider key fixture should insert");
let observed = sqlx::query_scalar::<_, serde_json::Value>(
"SELECT upstream_metadata -> 'antigravity' FROM provider_api_keys WHERE id = $1",
)
.bind(&key_id)
.fetch_one(repository.pool())
.await
.expect("metadata namespace should load");
assert_ne!(
serde_json::to_string(&observed).expect("metadata should serialize"),
r#"{"used_percent":0.123456789012345678901234567890}"#,
"the fixture must exercise precision loss in serde_json's default number representation",
);
let updated = repository
.update_key_runtime_metadata(&ProviderCatalogKeyRuntimeMetadataUpdate {
key_id: key_id.clone(),
namespace: "antigravity".to_string(),
expected_upstream_metadata_value: Some(observed),
upstream_metadata_value: json!({"used_percent": 12.5}),
status_snapshot_patch: json!({"quota": {"used_percent": 12.5}}),
updated_at_unix_secs: Some(1_700_000_000),
})
.await
.expect("runtime metadata CAS should execute");
assert!(
updated,
"matching metadata must not report a false CAS conflict"
);
let stored = sqlx::query_scalar::<_, serde_json::Value>(
"SELECT upstream_metadata -> 'antigravity' FROM provider_api_keys WHERE id = $1",
)
.bind(&key_id)
.fetch_one(repository.pool())
.await
.expect("updated metadata namespace should load");
assert_eq!(stored, json!({"used_percent": 12.5}));
sqlx::query("DELETE FROM provider_api_keys WHERE id = $1")
.bind(&key_id)
.execute(repository.pool())
.await
.expect("provider key fixture should delete");
sqlx::query("DELETE FROM providers WHERE id = $1")
.bind(&provider_id)
.execute(repository.pool())
.await
.expect("provider fixture should delete");
}
#[test]
fn ordinary_key_update_does_not_own_runtime_observation_fields() {
let sql = super::KEY_UPDATE_SQL.to_ascii_lowercase();
for runtime_assignment in [
"learned_rpm_limit =",
"rpm_429_count =",
"last_429_at =",
"last_429_type =",
"adjustment_history =",
"utilization_samples =",
"last_probe_increase_at =",
"last_rpm_peak =",
"upstream_metadata =",
"status_snapshot =",
"health_by_format =",
"circuit_breaker_by_format =",
] {
assert!(
!sql.contains(runtime_assignment),
"ordinary key update unexpectedly owns {runtime_assignment}"
);
}
assert!(sql.contains("is_active = $25"));
assert!(sql.contains("rpm_limit = $12"));
2026-08-14 09:28:07 +08:00
assert!(sql.contains("api_key is not distinct from $5"));
assert!(sql.contains("auth_config is not distinct from $6"));
}
#[test]
fn credential_cas_migrates_legacy_encrypted_key_with_null_safe_fence() {
let source = include_str!("provider_catalog.rs");
assert!(source.contains("SET api_key = $5, encrypted_key = NULL, auth_config = $6"));
assert!(source.contains("AND COALESCE(api_key, encrypted_key) IS NOT DISTINCT FROM $3"));
}
2026-08-14 09:28:07 +08:00
#[test]
fn admin_credential_cas_has_atomic_rotation_guards() {
let source = include_str!("provider_catalog.rs");
for predicate in [
"api_key IS NOT DISTINCT FROM ",
"auth_config IS NOT DISTINCT FROM ",
"jsonb_typeof(COALESCE(upstream_metadata, '{}'::jsonb)) = 'object'",
"jsonb_typeof(COALESCE(status_snapshot::jsonb, '{}'::jsonb)) = 'object'",
"jsonb_set(COALESCE(upstream_metadata, '{}'::jsonb), '{codex}'",
"jsonb_set(COALESCE(status_snapshot::jsonb, '{}'::jsonb), '{quota}', 'null'::jsonb, true)::json",
"oauth_invalid_at = NULL, oauth_invalid_reason = NULL, error_count = 0",
] {
assert!(
source.contains(predicate),
"missing admin CAS guard: {predicate}"
);
}
}
}