Files
Aether/crates/aether-data/adapters/postgres/src/candidate_selection.rs
T
elky 531cf11025 feat(gateway): harden provider request execution
Preserve exact request payloads and model client surface and API operation explicitly.

Add Anthropic compatibility profiles, bounded stream commitment, and scoped OAuth retry behavior across provider transports.
2026-07-27 09:36:31 +08:00

1692 lines
57 KiB
Rust

use async_trait::async_trait;
use futures_util::{stream::TryStream, TryStreamExt};
use sqlx::{PgPool, Row};
use std::collections::BTreeSet;
use aether_data_contracts::repository::candidate_selection::{
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsByKeyIdsQuery,
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredRequestedModelCandidateRowsQuery,
};
use aether_data_contracts::DataLayerError;
use crate::error::SqlxResultExt;
const LIST_FOR_EXACT_API_FORMAT_SQL: &str = r#"
WITH candidate_rows AS (
SELECT
p.id AS provider_id,
p.name AS provider_name,
p.provider_type AS provider_type,
p.provider_priority AS provider_priority,
p.is_active AS provider_is_active,
pe.id AS endpoint_id,
pe.api_format AS endpoint_api_format,
pe.api_family AS endpoint_api_family,
pe.endpoint_kind AS endpoint_kind,
pe.is_active AS endpoint_is_active,
pak.id AS key_id,
pak.name AS key_name,
pak.auth_type AS key_auth_type,
pak.is_active AS key_is_active,
pak.api_formats AS key_api_formats,
pak.allowed_models AS key_allowed_models,
pak.capabilities AS key_capabilities,
pak.internal_priority AS key_internal_priority,
pak.global_priority_by_format AS key_global_priority_by_format,
m.id AS model_id,
m.global_model_id AS global_model_id,
gm.name AS global_model_name,
CASE
WHEN gm.config IS NOT NULL THEN gm.config -> 'model_mappings'
ELSE NULL
END AS global_model_mappings,
CASE
WHEN gm.config IS NOT NULL AND gm.config ? 'streaming'
THEN (gm.config ->> 'streaming')::BOOLEAN
ELSE NULL
END AS global_model_supports_streaming,
m.provider_model_name AS model_provider_model_name,
m.provider_model_mappings AS model_provider_model_mappings,
m.supports_streaming AS model_supports_streaming,
m.is_active AS model_is_active,
m.is_available AS model_is_available,
(p.config -> 'pool_advanced') IS NOT NULL AS provider_pool_enabled
FROM providers p
INNER JOIN provider_endpoints pe
ON pe.provider_id = p.id
INNER JOIN LATERAL (
SELECT pak.*
FROM provider_api_keys pak
WHERE pak.provider_id = p.id
AND pak.is_active IS TRUE
AND (
pak.api_formats IS NULL
OR EXISTS (
SELECT 1
FROM json_array_elements_text(pak.api_formats) AS fmt(value)
WHERE LOWER(BTRIM(fmt.value)) = ANY($2::text[])
)
)
AND (
(
LOWER(BTRIM(p.provider_type)) = 'codex'
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($3) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image')
)
OR (
LOWER(BTRIM(p.provider_type)) = 'chatgpt_web'
AND LOWER(BTRIM(pak.auth_type)) IN ('oauth', 'bearer')
AND LOWER($3) = 'openai:image'
)
OR (
LOWER(BTRIM(p.provider_type)) = 'claude_code'
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($3) = 'claude:messages'
)
OR (
LOWER(BTRIM(p.provider_type)) = 'kiro'
AND LOWER($3) = 'claude:messages'
AND (
LOWER(BTRIM(pak.auth_type)) = 'oauth'
OR (
LOWER(BTRIM(pak.auth_type)) = 'bearer'
AND pak.auth_config IS NOT NULL
AND BTRIM(pak.auth_config) <> ''
)
)
)
OR (
LOWER(BTRIM(p.provider_type)) = 'grok'
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($3) IN ('openai:chat', 'openai:responses', 'claude:messages', 'openai:image')
)
OR (
LOWER(BTRIM(p.provider_type)) IN ('gemini_cli', 'antigravity')
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($3) = 'gemini:generate_content'
)
OR (
LOWER(BTRIM(p.provider_type)) = 'windsurf'
AND LOWER(BTRIM(pak.auth_type)) IN ('oauth', 'api_key', 'bearer')
AND LOWER($3) = 'openai:chat'
)
OR (
LOWER(BTRIM(p.provider_type)) = 'vertex_ai'
AND LOWER(BTRIM(pak.auth_type)) IN ('api_key', 'service_account', 'vertex_ai')
AND LOWER($3) IN ('gemini:generate_content', 'gemini:embedding')
)
OR (
LOWER(BTRIM(p.provider_type)) NOT IN (
'chatgpt_web',
'claude_code',
'codex',
'gemini_cli',
'grok',
'vertex_ai',
'antigravity',
'kiro',
'windsurf'
)
AND LOWER(BTRIM(pak.auth_type)) <> 'oauth'
)
)
ORDER BY pak.internal_priority ASC, pak.id ASC
LIMIT CASE WHEN (p.config -> 'pool_advanced') IS NOT NULL THEN 1 ELSE 2147483647 END
) pak ON TRUE
INNER JOIN models m
ON m.provider_id = p.id
INNER JOIN global_models gm
ON gm.id = m.global_model_id
WHERE p.is_active = TRUE
AND pe.is_active = TRUE
AND pak.is_active IS TRUE
AND m.is_active = TRUE
AND m.is_available = TRUE
AND gm.is_active = TRUE
AND LOWER(pe.api_format) = LOWER($1)
AND (
pak.api_formats IS NULL
OR EXISTS (
SELECT 1
FROM json_array_elements_text(pak.api_formats) AS fmt(value)
WHERE LOWER(BTRIM(fmt.value)) = ANY($2::text[])
)
)
AND (
(
LOWER(BTRIM(p.provider_type)) = 'codex'
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($3) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image')
)
OR (
LOWER(BTRIM(p.provider_type)) = 'chatgpt_web'
AND LOWER(BTRIM(pak.auth_type)) IN ('oauth', 'bearer')
AND LOWER($3) = 'openai:image'
)
OR (
LOWER(BTRIM(p.provider_type)) = 'claude_code'
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($3) = 'claude:messages'
)
OR (
LOWER(BTRIM(p.provider_type)) = 'kiro'
AND LOWER($3) = 'claude:messages'
AND (
LOWER(BTRIM(pak.auth_type)) = 'oauth'
OR (
LOWER(BTRIM(pak.auth_type)) = 'bearer'
AND pak.auth_config IS NOT NULL
AND BTRIM(pak.auth_config) <> ''
)
)
)
OR (
LOWER(BTRIM(p.provider_type)) = 'grok'
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($3) IN ('openai:chat', 'openai:responses', 'claude:messages', 'openai:image')
)
OR (
LOWER(BTRIM(p.provider_type)) IN ('gemini_cli', 'antigravity')
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($3) = 'gemini:generate_content'
)
OR (
LOWER(BTRIM(p.provider_type)) = 'windsurf'
AND LOWER(BTRIM(pak.auth_type)) IN ('oauth', 'api_key', 'bearer')
AND LOWER($3) = 'openai:chat'
)
OR (
LOWER(BTRIM(p.provider_type)) = 'vertex_ai'
AND LOWER(BTRIM(pak.auth_type)) IN ('api_key', 'service_account', 'vertex_ai')
AND LOWER($3) IN ('gemini:generate_content', 'gemini:embedding')
)
OR (
LOWER(BTRIM(p.provider_type)) NOT IN (
'chatgpt_web',
'claude_code',
'codex',
'gemini_cli',
'grok',
'vertex_ai',
'antigravity',
'kiro',
'windsurf'
)
AND LOWER(BTRIM(pak.auth_type)) <> 'oauth'
)
)
),
pool_rows AS (
SELECT DISTINCT ON (provider_id, endpoint_id, model_id)
*
FROM candidate_rows
WHERE provider_pool_enabled
ORDER BY
provider_id ASC,
endpoint_id ASC,
model_id ASC,
key_internal_priority ASC,
key_id ASC
),
selected_rows AS (
SELECT * FROM candidate_rows WHERE NOT provider_pool_enabled
UNION ALL
SELECT * FROM pool_rows
)
SELECT
provider_id,
provider_name,
provider_type,
provider_priority,
provider_is_active,
endpoint_id,
endpoint_api_format,
endpoint_api_family,
endpoint_kind,
endpoint_is_active,
key_id,
key_name,
key_auth_type,
key_is_active,
key_api_formats,
key_allowed_models,
key_capabilities,
key_internal_priority,
key_global_priority_by_format,
model_id,
global_model_id,
global_model_name,
global_model_mappings,
global_model_supports_streaming,
model_provider_model_name,
model_provider_model_mappings,
model_supports_streaming,
model_is_active,
model_is_available
FROM selected_rows
ORDER BY
global_model_name ASC,
provider_priority ASC,
key_internal_priority ASC,
provider_id ASC,
endpoint_id ASC,
key_id ASC,
model_id ASC
"#;
const LIST_FOR_EXACT_API_FORMAT_AND_GLOBAL_MODEL_SQL: &str = r#"
WITH candidate_rows AS (
SELECT
p.id AS provider_id,
p.name AS provider_name,
p.provider_type AS provider_type,
p.provider_priority AS provider_priority,
p.is_active AS provider_is_active,
pe.id AS endpoint_id,
pe.api_format AS endpoint_api_format,
pe.api_family AS endpoint_api_family,
pe.endpoint_kind AS endpoint_kind,
pe.is_active AS endpoint_is_active,
pak.id AS key_id,
pak.name AS key_name,
pak.auth_type AS key_auth_type,
pak.is_active AS key_is_active,
pak.api_formats AS key_api_formats,
pak.allowed_models AS key_allowed_models,
pak.capabilities AS key_capabilities,
pak.internal_priority AS key_internal_priority,
pak.global_priority_by_format AS key_global_priority_by_format,
m.id AS model_id,
m.global_model_id AS global_model_id,
gm.name AS global_model_name,
CASE
WHEN gm.config IS NOT NULL THEN gm.config -> 'model_mappings'
ELSE NULL
END AS global_model_mappings,
CASE
WHEN gm.config IS NOT NULL AND gm.config ? 'streaming'
THEN (gm.config ->> 'streaming')::BOOLEAN
ELSE NULL
END AS global_model_supports_streaming,
m.provider_model_name AS model_provider_model_name,
m.provider_model_mappings AS model_provider_model_mappings,
m.supports_streaming AS model_supports_streaming,
m.is_active AS model_is_active,
m.is_available AS model_is_available,
(p.config -> 'pool_advanced') IS NOT NULL AS provider_pool_enabled
FROM providers p
INNER JOIN provider_endpoints pe
ON pe.provider_id = p.id
INNER JOIN LATERAL (
SELECT pak.*
FROM provider_api_keys pak
WHERE pak.provider_id = p.id
AND pak.is_active IS TRUE
AND (
pak.api_formats IS NULL
OR EXISTS (
SELECT 1
FROM json_array_elements_text(pak.api_formats) AS fmt(value)
WHERE LOWER(BTRIM(fmt.value)) = ANY($3::text[])
)
)
AND (
(
LOWER(BTRIM(p.provider_type)) = 'codex'
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($4) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image')
)
OR (
LOWER(BTRIM(p.provider_type)) = 'chatgpt_web'
AND LOWER(BTRIM(pak.auth_type)) IN ('oauth', 'bearer')
AND LOWER($4) = 'openai:image'
)
OR (
LOWER(BTRIM(p.provider_type)) = 'claude_code'
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($4) = 'claude:messages'
)
OR (
LOWER(BTRIM(p.provider_type)) = 'kiro'
AND LOWER($4) = 'claude:messages'
AND (
LOWER(BTRIM(pak.auth_type)) = 'oauth'
OR (
LOWER(BTRIM(pak.auth_type)) = 'bearer'
AND pak.auth_config IS NOT NULL
AND BTRIM(pak.auth_config) <> ''
)
)
)
OR (
LOWER(BTRIM(p.provider_type)) = 'grok'
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($4) IN ('openai:chat', 'openai:responses', 'claude:messages', 'openai:image')
)
OR (
LOWER(BTRIM(p.provider_type)) IN ('gemini_cli', 'antigravity')
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($4) = 'gemini:generate_content'
)
OR (
LOWER(BTRIM(p.provider_type)) = 'windsurf'
AND LOWER(BTRIM(pak.auth_type)) IN ('oauth', 'api_key', 'bearer')
AND LOWER($4) = 'openai:chat'
)
OR (
LOWER(BTRIM(p.provider_type)) = 'vertex_ai'
AND LOWER(BTRIM(pak.auth_type)) IN ('api_key', 'service_account', 'vertex_ai')
AND LOWER($4) IN ('gemini:generate_content', 'gemini:embedding')
)
OR (
LOWER(BTRIM(p.provider_type)) NOT IN (
'chatgpt_web',
'claude_code',
'codex',
'gemini_cli',
'grok',
'vertex_ai',
'antigravity',
'kiro',
'windsurf'
)
AND LOWER(BTRIM(pak.auth_type)) <> 'oauth'
)
)
ORDER BY pak.internal_priority ASC, pak.id ASC
LIMIT CASE WHEN (p.config -> 'pool_advanced') IS NOT NULL THEN 1 ELSE 2147483647 END
) pak ON TRUE
INNER JOIN models m
ON m.provider_id = p.id
INNER JOIN global_models gm
ON gm.id = m.global_model_id
WHERE p.is_active = TRUE
AND pe.is_active = TRUE
AND pak.is_active IS TRUE
AND m.is_active = TRUE
AND m.is_available = TRUE
AND gm.is_active = TRUE
AND LOWER(pe.api_format) = LOWER($1)
AND gm.name = $2
AND (
pak.api_formats IS NULL
OR EXISTS (
SELECT 1
FROM json_array_elements_text(pak.api_formats) AS fmt(value)
WHERE LOWER(BTRIM(fmt.value)) = ANY($3::text[])
)
)
AND (
(
LOWER(BTRIM(p.provider_type)) = 'codex'
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($4) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image')
)
OR (
LOWER(BTRIM(p.provider_type)) = 'chatgpt_web'
AND LOWER(BTRIM(pak.auth_type)) IN ('oauth', 'bearer')
AND LOWER($4) = 'openai:image'
)
OR (
LOWER(BTRIM(p.provider_type)) = 'claude_code'
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($4) = 'claude:messages'
)
OR (
LOWER(BTRIM(p.provider_type)) = 'kiro'
AND LOWER($4) = 'claude:messages'
AND (
LOWER(BTRIM(pak.auth_type)) = 'oauth'
OR (
LOWER(BTRIM(pak.auth_type)) = 'bearer'
AND pak.auth_config IS NOT NULL
AND BTRIM(pak.auth_config) <> ''
)
)
)
OR (
LOWER(BTRIM(p.provider_type)) = 'grok'
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($4) IN ('openai:chat', 'openai:responses', 'claude:messages', 'openai:image')
)
OR (
LOWER(BTRIM(p.provider_type)) IN ('gemini_cli', 'antigravity')
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($4) = 'gemini:generate_content'
)
OR (
LOWER(BTRIM(p.provider_type)) = 'windsurf'
AND LOWER(BTRIM(pak.auth_type)) IN ('oauth', 'api_key', 'bearer')
AND LOWER($4) = 'openai:chat'
)
OR (
LOWER(BTRIM(p.provider_type)) = 'vertex_ai'
AND LOWER(BTRIM(pak.auth_type)) IN ('api_key', 'service_account', 'vertex_ai')
AND LOWER($4) IN ('gemini:generate_content', 'gemini:embedding')
)
OR (
LOWER(BTRIM(p.provider_type)) NOT IN (
'chatgpt_web',
'claude_code',
'codex',
'gemini_cli',
'grok',
'vertex_ai',
'antigravity',
'kiro',
'windsurf'
)
AND LOWER(BTRIM(pak.auth_type)) <> 'oauth'
)
)
),
pool_rows AS (
SELECT DISTINCT ON (provider_id, endpoint_id, model_id)
*
FROM candidate_rows
WHERE provider_pool_enabled
ORDER BY
provider_id ASC,
endpoint_id ASC,
model_id ASC,
key_internal_priority ASC,
key_id ASC
),
selected_rows AS (
SELECT * FROM candidate_rows WHERE NOT provider_pool_enabled
UNION ALL
SELECT * FROM pool_rows
)
SELECT
provider_id,
provider_name,
provider_type,
provider_priority,
provider_is_active,
endpoint_id,
endpoint_api_format,
endpoint_api_family,
endpoint_kind,
endpoint_is_active,
key_id,
key_name,
key_auth_type,
key_is_active,
key_api_formats,
key_allowed_models,
key_capabilities,
key_internal_priority,
key_global_priority_by_format,
model_id,
global_model_id,
global_model_name,
global_model_mappings,
global_model_supports_streaming,
model_provider_model_name,
model_provider_model_mappings,
model_supports_streaming,
model_is_active,
model_is_available
FROM selected_rows
ORDER BY
provider_priority ASC,
key_internal_priority ASC,
provider_id ASC,
endpoint_id ASC,
key_id ASC,
model_id ASC
"#;
const LIST_POOL_KEYS_FOR_GROUP_SQL: &str = r#"
SELECT
p.id AS provider_id,
p.name AS provider_name,
p.provider_type AS provider_type,
p.provider_priority AS provider_priority,
p.is_active AS provider_is_active,
pe.id AS endpoint_id,
pe.api_format AS endpoint_api_format,
pe.api_family AS endpoint_api_family,
pe.endpoint_kind AS endpoint_kind,
pe.is_active AS endpoint_is_active,
pak.id AS key_id,
pak.name AS key_name,
pak.auth_type AS key_auth_type,
pak.is_active AS key_is_active,
pak.api_formats AS key_api_formats,
pak.allowed_models AS key_allowed_models,
pak.capabilities AS key_capabilities,
pak.internal_priority AS key_internal_priority,
pak.global_priority_by_format AS key_global_priority_by_format,
m.id AS model_id,
m.global_model_id AS global_model_id,
gm.name AS global_model_name,
CASE
WHEN gm.config IS NOT NULL THEN gm.config -> 'model_mappings'
ELSE NULL
END AS global_model_mappings,
CASE
WHEN gm.config IS NOT NULL AND gm.config ? 'streaming'
THEN (gm.config ->> 'streaming')::BOOLEAN
ELSE NULL
END AS global_model_supports_streaming,
m.provider_model_name AS model_provider_model_name,
m.provider_model_mappings AS model_provider_model_mappings,
m.supports_streaming AS model_supports_streaming,
m.is_active AS model_is_active,
m.is_available AS model_is_available
FROM providers p
INNER JOIN provider_endpoints pe
ON pe.provider_id = p.id
INNER JOIN provider_api_keys pak
ON pak.provider_id = p.id
INNER JOIN models m
ON m.provider_id = p.id
INNER JOIN global_models gm
ON gm.id = m.global_model_id
WHERE p.is_active = TRUE
AND pe.is_active = TRUE
AND pak.is_active IS TRUE
AND m.is_active = TRUE
AND m.is_available = TRUE
AND gm.is_active = TRUE
AND LOWER(pe.api_format) = LOWER($1)
AND p.id = $2
AND pe.id = $3
AND m.id = $4
AND (
pak.api_formats IS NULL
OR EXISTS (
SELECT 1
FROM json_array_elements_text(pak.api_formats) AS fmt(value)
WHERE LOWER(BTRIM(fmt.value)) = ANY($5::text[])
)
)
AND (
(
LOWER(BTRIM(p.provider_type)) = 'codex'
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($6) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image')
)
OR (
LOWER(BTRIM(p.provider_type)) = 'chatgpt_web'
AND LOWER(BTRIM(pak.auth_type)) IN ('oauth', 'bearer')
AND LOWER($6) = 'openai:image'
)
OR (
LOWER(BTRIM(p.provider_type)) = 'claude_code'
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($6) = 'claude:messages'
)
OR (
LOWER(BTRIM(p.provider_type)) = 'kiro'
AND LOWER($6) = 'claude:messages'
AND (
LOWER(BTRIM(pak.auth_type)) = 'oauth'
OR (
LOWER(BTRIM(pak.auth_type)) = 'bearer'
AND pak.auth_config IS NOT NULL
AND BTRIM(pak.auth_config) <> ''
)
)
)
OR (
LOWER(BTRIM(p.provider_type)) = 'grok'
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($6) IN ('openai:chat', 'openai:responses', 'claude:messages', 'openai:image')
)
OR (
LOWER(BTRIM(p.provider_type)) IN ('gemini_cli', 'antigravity')
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($6) = 'gemini:generate_content'
)
OR (
LOWER(BTRIM(p.provider_type)) = 'windsurf'
AND LOWER(BTRIM(pak.auth_type)) IN ('oauth', 'api_key', 'bearer')
AND LOWER($6) = 'openai:chat'
)
OR (
LOWER(BTRIM(p.provider_type)) = 'vertex_ai'
AND LOWER(BTRIM(pak.auth_type)) IN ('api_key', 'service_account', 'vertex_ai')
AND LOWER($6) IN ('gemini:generate_content', 'gemini:embedding')
)
OR (
LOWER(BTRIM(p.provider_type)) NOT IN (
'chatgpt_web',
'claude_code',
'codex',
'gemini_cli',
'grok',
'vertex_ai',
'antigravity',
'kiro',
'windsurf'
)
AND LOWER(BTRIM(pak.auth_type)) <> 'oauth'
)
)
ORDER BY
pak.internal_priority ASC,
pak.id ASC
LIMIT $7
OFFSET $8
"#;
fn pool_key_candidate_order_by_sql(order: &StoredPoolKeyCandidateOrder) -> &'static str {
match order {
StoredPoolKeyCandidateOrder::InternalPriority => {
"ORDER BY\n pak.internal_priority ASC,\n pak.id ASC"
}
StoredPoolKeyCandidateOrder::Lru => {
"ORDER BY\n pak.last_used_at ASC NULLS FIRST,\n pak.internal_priority ASC,\n pak.id ASC"
}
StoredPoolKeyCandidateOrder::CacheAffinity => {
"ORDER BY\n pak.last_used_at DESC NULLS LAST,\n pak.internal_priority ASC,\n pak.id ASC"
}
StoredPoolKeyCandidateOrder::SingleAccount => {
"ORDER BY\n pak.internal_priority ASC,\n pak.last_used_at DESC NULLS LAST,\n pak.id ASC"
}
StoredPoolKeyCandidateOrder::LoadBalance { .. } => {
"ORDER BY\n md5($9 || ':' || pak.id) ASC,\n pak.id ASC"
}
}
}
fn pool_key_candidate_selection_sql(order: &StoredPoolKeyCandidateOrder) -> String {
let default_order =
"ORDER BY\n pak.internal_priority ASC,\n pak.id ASC\nLIMIT $7\nOFFSET $8\n";
let replacement = format!(
"{}\nLIMIT $7\nOFFSET $8\n",
pool_key_candidate_order_by_sql(order)
);
LIST_POOL_KEYS_FOR_GROUP_SQL.replace(default_order, &replacement)
}
fn pool_key_candidate_selection_by_key_ids_sql() -> String {
let default_order =
"ORDER BY\n pak.internal_priority ASC,\n pak.id ASC\nLIMIT $7\nOFFSET $8\n";
let replacement = "AND pak.id = ANY($7::text[])\nORDER BY\n array_position($7::text[], pak.id) ASC,\n pak.id ASC\n";
LIST_POOL_KEYS_FOR_GROUP_SQL.replace(default_order, replacement)
}
#[derive(Debug, Clone)]
pub struct SqlxMinimalCandidateSelectionReadRepository {
pool: PgPool,
}
impl SqlxMinimalCandidateSelectionReadRepository {
pub fn new(pool: PgPool) -> Self {
Self { pool }
}
pub fn pool(&self) -> &PgPool {
&self.pool
}
async fn collect_query_rows<T, S>(
mut rows: S,
map_row: fn(&sqlx::postgres::PgRow) -> Result<T, DataLayerError>,
) -> Result<Vec<T>, DataLayerError>
where
S: TryStream<Ok = sqlx::postgres::PgRow, Error = sqlx::Error> + Unpin,
{
let mut items = Vec::new();
while let Some(row) = rows.try_next().await.map_postgres_err()? {
items.push(map_row(&row)?);
}
Ok(items)
}
pub async fn list_for_exact_api_format(
&self,
api_format: &str,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
let mut rows = Vec::new();
let canonical_api_format = normalize_api_format(api_format);
let storage_aliases = api_format_aliases(&canonical_api_format);
let sql_match_aliases =
sql_match_aliases(&api_format_permission_aliases(&canonical_api_format));
for api_format in storage_aliases {
rows.extend(
Self::collect_query_rows(
sqlx::query(LIST_FOR_EXACT_API_FORMAT_SQL)
.bind(api_format)
.bind(sql_match_aliases.clone())
.bind(canonical_api_format.clone())
.fetch(&self.pool),
map_candidate_selection_row,
)
.await?,
);
}
Ok(dedupe_candidate_selection_rows(rows))
}
pub async fn list_for_exact_api_format_and_global_model(
&self,
api_format: &str,
global_model_name: &str,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
let mut rows = Vec::new();
let canonical_api_format = normalize_api_format(api_format);
let storage_aliases = api_format_aliases(&canonical_api_format);
let sql_match_aliases =
sql_match_aliases(&api_format_permission_aliases(&canonical_api_format));
for api_format in storage_aliases {
rows.extend(
Self::collect_query_rows(
sqlx::query(LIST_FOR_EXACT_API_FORMAT_AND_GLOBAL_MODEL_SQL)
.bind(api_format)
.bind(global_model_name)
.bind(sql_match_aliases.clone())
.bind(canonical_api_format.clone())
.fetch(&self.pool),
map_candidate_selection_row,
)
.await?,
);
}
Ok(dedupe_candidate_selection_rows(rows))
}
pub async fn list_for_exact_api_format_and_requested_model(
&self,
api_format: &str,
requested_model_name: &str,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
let mut rows = Vec::new();
let canonical_api_format = normalize_api_format(api_format);
let storage_aliases = api_format_aliases(&canonical_api_format);
let sql_match_aliases =
sql_match_aliases(&api_format_permission_aliases(&canonical_api_format));
let sql = requested_model_selection_sql();
for api_format in storage_aliases {
rows.extend(
Self::collect_query_rows(
sqlx::query(sql.as_str())
.bind(api_format)
.bind(requested_model_name)
.bind(sql_match_aliases.clone())
.bind(canonical_api_format.clone())
.fetch(&self.pool),
map_candidate_selection_row,
)
.await?,
);
}
Ok(dedupe_candidate_selection_rows(rows))
}
pub async fn list_for_exact_api_format_and_requested_model_page(
&self,
query: &StoredRequestedModelCandidateRowsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
let mut rows = Vec::new();
let canonical_api_format = normalize_api_format(&query.api_format);
let storage_aliases = api_format_aliases(&canonical_api_format);
let sql_match_aliases =
sql_match_aliases(&api_format_permission_aliases(&canonical_api_format));
let limit = i64::from(query.limit.max(1));
let offset = i64::from(query.offset);
let sql = requested_model_selection_page_sql();
for api_format in storage_aliases {
rows.extend(
Self::collect_query_rows(
sqlx::query(sql.as_str())
.bind(api_format)
.bind(query.requested_model_name.as_str())
.bind(sql_match_aliases.clone())
.bind(canonical_api_format.clone())
.bind(limit)
.bind(offset)
.fetch(&self.pool),
map_candidate_selection_row,
)
.await?,
);
}
Ok(dedupe_candidate_selection_rows(rows))
}
pub async fn list_pool_key_rows_for_group(
&self,
query: &StoredPoolKeyCandidateRowsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
let mut rows = Vec::new();
let canonical_api_format = normalize_api_format(&query.api_format);
let storage_aliases = api_format_aliases(&canonical_api_format);
let sql_match_aliases =
sql_match_aliases(&api_format_permission_aliases(&canonical_api_format));
let limit = i64::from(query.limit.max(1));
let offset = i64::from(query.offset);
let sql = pool_key_candidate_selection_sql(&query.order);
for api_format in storage_aliases {
let mut query_builder = sqlx::query(sql.as_str())
.bind(api_format)
.bind(query.provider_id.as_str())
.bind(query.endpoint_id.as_str())
.bind(query.model_id.as_str())
.bind(sql_match_aliases.clone())
.bind(canonical_api_format.clone())
.bind(limit)
.bind(offset);
if let StoredPoolKeyCandidateOrder::LoadBalance { seed } = &query.order {
query_builder = query_builder.bind(seed.as_str());
}
rows.extend(
Self::collect_query_rows(
query_builder.fetch(&self.pool),
map_candidate_selection_row,
)
.await?,
);
}
Ok(dedupe_candidate_selection_rows(rows))
}
pub async fn list_pool_key_rows_for_group_key_ids(
&self,
query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
if query.key_ids.is_empty() {
return Ok(Vec::new());
}
let mut rows = Vec::new();
let canonical_api_format = normalize_api_format(&query.api_format);
let storage_aliases = api_format_aliases(&canonical_api_format);
let sql_match_aliases =
sql_match_aliases(&api_format_permission_aliases(&canonical_api_format));
let sql = pool_key_candidate_selection_by_key_ids_sql();
for api_format in storage_aliases {
rows.extend(
Self::collect_query_rows(
sqlx::query(sql.as_str())
.bind(api_format)
.bind(query.provider_id.as_str())
.bind(query.endpoint_id.as_str())
.bind(query.model_id.as_str())
.bind(sql_match_aliases.clone())
.bind(canonical_api_format.clone())
.bind(query.key_ids.clone())
.fetch(&self.pool),
map_candidate_selection_row,
)
.await?,
);
}
let key_order = query
.key_ids
.iter()
.enumerate()
.map(|(index, key_id)| (key_id.as_str(), index))
.collect::<std::collections::BTreeMap<_, _>>();
let mut rows = dedupe_candidate_selection_rows(rows);
rows.sort_by(|left, right| {
key_order
.get(left.key_id.as_str())
.cmp(&key_order.get(right.key_id.as_str()))
.then(left.key_id.cmp(&right.key_id))
});
Ok(rows)
}
}
fn requested_model_selection_sql() -> String {
LIST_FOR_EXACT_API_FORMAT_AND_GLOBAL_MODEL_SQL
.replace(
"AND gm.name = $2",
r#"AND (
(
gm.name = $2
AND (
m.provider_model_mappings IS NULL
OR jsonb_typeof(m.provider_model_mappings) <> 'array'
OR EXISTS (
SELECT 1
FROM jsonb_array_elements(
CASE
WHEN jsonb_typeof(m.provider_model_mappings) = 'array'
THEN m.provider_model_mappings
ELSE '[]'::jsonb
END
) AS mapping(value)
WHERE (
mapping.value -> 'api_formats' IS NULL
OR jsonb_typeof(mapping.value -> 'api_formats') <> 'array'
OR EXISTS (
SELECT 1
FROM jsonb_array_elements_text(mapping.value -> 'api_formats') AS fmt(value)
WHERE LOWER(BTRIM(fmt.value)) = ANY($3::text[])
)
)
AND (
mapping.value -> 'endpoint_ids' IS NULL
OR jsonb_typeof(mapping.value -> 'endpoint_ids') <> 'array'
OR EXISTS (
SELECT 1
FROM jsonb_array_elements_text(mapping.value -> 'endpoint_ids') AS endpoint(value)
WHERE endpoint.value = pe.id
)
)
)
OR NOT EXISTS (
SELECT 1
FROM jsonb_array_elements(
CASE
WHEN jsonb_typeof(m.provider_model_mappings) = 'array'
THEN m.provider_model_mappings
ELSE '[]'::jsonb
END
) AS mapping(value)
WHERE mapping.value ->> 'name' = m.provider_model_name
)
)
)
OR (
m.provider_model_name = $2
AND (
m.provider_model_mappings IS NULL
OR jsonb_typeof(m.provider_model_mappings) <> 'array'
OR NOT EXISTS (
SELECT 1
FROM jsonb_array_elements(
CASE
WHEN jsonb_typeof(m.provider_model_mappings) = 'array'
THEN m.provider_model_mappings
ELSE '[]'::jsonb
END
) AS mapping(value)
WHERE mapping.value ->> 'name' = m.provider_model_name
)
OR EXISTS (
SELECT 1
FROM jsonb_array_elements(
CASE
WHEN jsonb_typeof(m.provider_model_mappings) = 'array'
THEN m.provider_model_mappings
ELSE '[]'::jsonb
END
) AS mapping(value)
WHERE mapping.value ->> 'name' = m.provider_model_name
AND (
mapping.value -> 'api_formats' IS NULL
OR jsonb_typeof(mapping.value -> 'api_formats') <> 'array'
OR EXISTS (
SELECT 1
FROM jsonb_array_elements_text(mapping.value -> 'api_formats') AS fmt(value)
WHERE LOWER(BTRIM(fmt.value)) = ANY($3::text[])
)
)
AND (
mapping.value -> 'endpoint_ids' IS NULL
OR jsonb_typeof(mapping.value -> 'endpoint_ids') <> 'array'
OR EXISTS (
SELECT 1
FROM jsonb_array_elements_text(mapping.value -> 'endpoint_ids') AS endpoint(value)
WHERE endpoint.value = pe.id
)
)
)
)
)
OR (
jsonb_typeof(m.provider_model_mappings) = 'array'
AND EXISTS (
SELECT 1
FROM jsonb_array_elements(
CASE
WHEN jsonb_typeof(m.provider_model_mappings) = 'array'
THEN m.provider_model_mappings
ELSE '[]'::jsonb
END
) AS mapping(value)
WHERE mapping.value ->> 'name' = $2
AND (
mapping.value -> 'api_formats' IS NULL
OR jsonb_typeof(mapping.value -> 'api_formats') <> 'array'
OR EXISTS (
SELECT 1
FROM jsonb_array_elements_text(mapping.value -> 'api_formats') AS fmt(value)
WHERE LOWER(BTRIM(fmt.value)) = ANY($3::text[])
)
)
AND (
mapping.value -> 'endpoint_ids' IS NULL
OR jsonb_typeof(mapping.value -> 'endpoint_ids') <> 'array'
OR EXISTS (
SELECT 1
FROM jsonb_array_elements_text(mapping.value -> 'endpoint_ids') AS endpoint(value)
WHERE endpoint.value = pe.id
)
)
)
)
)"#,
)
.replace(
"ORDER BY\n provider_priority ASC,",
"ORDER BY\n global_model_name ASC,\n provider_priority ASC,",
)
}
fn requested_model_selection_page_sql() -> String {
format!("{}\nLIMIT $5\nOFFSET $6", requested_model_selection_sql())
}
fn api_format_aliases(api_format: &str) -> Vec<String> {
aether_ai_formats::api_format_storage_aliases(api_format)
}
fn api_format_permission_aliases(api_format: &str) -> Vec<String> {
aether_ai_formats::api_format_permission_storage_aliases(api_format)
}
fn normalize_api_format(api_format: &str) -> String {
aether_ai_formats::normalize_api_format_alias(api_format)
}
fn sql_match_aliases(api_formats: &[String]) -> Vec<String> {
api_formats
.iter()
.map(|value| value.trim().to_ascii_lowercase())
.collect()
}
fn dedupe_candidate_selection_rows(
rows: Vec<StoredMinimalCandidateSelectionRow>,
) -> Vec<StoredMinimalCandidateSelectionRow> {
let mut seen = BTreeSet::new();
rows.into_iter()
.filter(|row| {
seen.insert((
row.endpoint_id.clone(),
row.key_id.clone(),
row.model_id.clone(),
))
})
.collect()
}
#[async_trait]
impl MinimalCandidateSelectionReadRepository for SqlxMinimalCandidateSelectionReadRepository {
async fn list_for_exact_api_format(
&self,
api_format: &str,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
Self::list_for_exact_api_format(self, api_format).await
}
async fn list_for_exact_api_format_and_global_model(
&self,
api_format: &str,
global_model_name: &str,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
Self::list_for_exact_api_format_and_global_model(self, api_format, global_model_name).await
}
async fn list_for_exact_api_format_and_requested_model(
&self,
api_format: &str,
requested_model_name: &str,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
Self::list_for_exact_api_format_and_requested_model(self, api_format, requested_model_name)
.await
}
async fn list_for_exact_api_format_and_requested_model_page(
&self,
query: &StoredRequestedModelCandidateRowsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
Self::list_for_exact_api_format_and_requested_model_page(self, query).await
}
async fn list_pool_key_rows_for_group(
&self,
query: &StoredPoolKeyCandidateRowsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
Self::list_pool_key_rows_for_group(self, query).await
}
async fn list_pool_key_rows_for_group_key_ids(
&self,
query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
Self::list_pool_key_rows_for_group_key_ids(self, query).await
}
}
fn map_candidate_selection_row(
row: &sqlx::postgres::PgRow,
) -> Result<StoredMinimalCandidateSelectionRow, DataLayerError> {
Ok(StoredMinimalCandidateSelectionRow {
provider_id: row.try_get("provider_id").map_postgres_err()?,
provider_name: row.try_get("provider_name").map_postgres_err()?,
provider_type: row.try_get("provider_type").map_postgres_err()?,
provider_priority: row.try_get("provider_priority").map_postgres_err()?,
provider_is_active: row.try_get("provider_is_active").map_postgres_err()?,
endpoint_id: row.try_get("endpoint_id").map_postgres_err()?,
endpoint_api_format: row.try_get("endpoint_api_format").map_postgres_err()?,
endpoint_api_family: row.try_get("endpoint_api_family").map_postgres_err()?,
endpoint_kind: row.try_get("endpoint_kind").map_postgres_err()?,
endpoint_is_active: row.try_get("endpoint_is_active").map_postgres_err()?,
key_id: row.try_get("key_id").map_postgres_err()?,
key_name: row.try_get("key_name").map_postgres_err()?,
key_auth_type: row.try_get("key_auth_type").map_postgres_err()?,
key_is_active: row.try_get("key_is_active").map_postgres_err()?,
key_api_formats: parse_string_list(
row.try_get("key_api_formats").map_postgres_err()?,
"provider_api_keys.api_formats",
)?,
key_allowed_models: parse_string_list(
row.try_get("key_allowed_models").map_postgres_err()?,
"provider_api_keys.allowed_models",
)?,
key_capabilities: row.try_get("key_capabilities").map_postgres_err()?,
key_internal_priority: row.try_get("key_internal_priority").map_postgres_err()?,
key_global_priority_by_format: row
.try_get("key_global_priority_by_format")
.map_postgres_err()?,
model_id: row.try_get("model_id").map_postgres_err()?,
global_model_id: row.try_get("global_model_id").map_postgres_err()?,
global_model_name: row.try_get("global_model_name").map_postgres_err()?,
global_model_mappings: parse_string_list(
row.try_get("global_model_mappings").map_postgres_err()?,
"global_models.config.model_mappings",
)?,
global_model_supports_streaming: row
.try_get("global_model_supports_streaming")
.map_postgres_err()?,
model_provider_model_name: row
.try_get("model_provider_model_name")
.map_postgres_err()?,
model_provider_model_mappings: parse_provider_model_mappings(
row.try_get("model_provider_model_mappings")
.map_postgres_err()?,
)?,
model_supports_streaming: row.try_get("model_supports_streaming").map_postgres_err()?,
model_is_active: row.try_get("model_is_active").map_postgres_err()?,
model_is_available: row.try_get("model_is_available").map_postgres_err()?,
})
}
fn parse_string_list(
value: Option<serde_json::Value>,
field_name: &str,
) -> Result<Option<Vec<String>>, DataLayerError> {
let Some(value) = value else {
return Ok(None);
};
parse_string_list_value(&value, field_name)
}
fn parse_string_list_value(
value: &serde_json::Value,
field_name: &str,
) -> Result<Option<Vec<String>>, DataLayerError> {
match value {
serde_json::Value::Null => Ok(None),
serde_json::Value::Array(array) => parse_string_list_array(array, field_name).map(Some),
serde_json::Value::String(raw) => parse_embedded_string_list(raw, field_name),
_ => Err(DataLayerError::UnexpectedValue(format!(
"{field_name} is not a JSON array"
))),
}
}
fn parse_embedded_string_list(
raw: &str,
field_name: &str,
) -> Result<Option<Vec<String>>, DataLayerError> {
let raw = raw.trim();
if raw.is_empty() || raw.eq_ignore_ascii_case("null") {
return Ok(None);
}
if let Ok(decoded) = serde_json::from_str::<serde_json::Value>(raw) {
return parse_string_list_value(&decoded, field_name);
}
Ok(Some(vec![raw.to_string()]))
}
fn parse_string_list_array(
array: &[serde_json::Value],
field_name: &str,
) -> Result<Vec<String>, DataLayerError> {
let mut items = Vec::with_capacity(array.len());
for item in array {
let Some(item) = item.as_str() else {
return Err(DataLayerError::UnexpectedValue(format!(
"{field_name} contains a non-string item"
)));
};
let item = item.trim();
if !item.is_empty() {
items.push(item.to_string());
}
}
Ok(items)
}
fn parse_provider_model_mappings(
value: Option<serde_json::Value>,
) -> Result<Option<Vec<StoredProviderModelMapping>>, DataLayerError> {
let Some(value) = value else {
return Ok(None);
};
parse_provider_model_mappings_value(&value)
}
fn parse_provider_model_mappings_value(
value: &serde_json::Value,
) -> Result<Option<Vec<StoredProviderModelMapping>>, DataLayerError> {
match value {
serde_json::Value::Null => Ok(None),
serde_json::Value::Array(array) => parse_provider_model_mappings_array(array),
serde_json::Value::Object(object) => {
parse_provider_model_mapping_object(object).map(|mapping| Some(vec![mapping]))
}
serde_json::Value::String(raw) => parse_embedded_provider_model_mappings(raw),
_ => Err(DataLayerError::UnexpectedValue(
"models.provider_model_mappings is not a JSON array".to_string(),
)),
}
}
fn parse_embedded_provider_model_mappings(
raw: &str,
) -> Result<Option<Vec<StoredProviderModelMapping>>, DataLayerError> {
let raw = raw.trim();
if raw.is_empty() || raw.eq_ignore_ascii_case("null") {
return Ok(None);
}
if let Ok(decoded) = serde_json::from_str::<serde_json::Value>(raw) {
return parse_provider_model_mappings_value(&decoded);
}
Ok(Some(vec![StoredProviderModelMapping {
name: raw.to_string(),
priority: 1,
api_formats: None,
endpoint_ids: None,
operations: None,
}]))
}
fn parse_provider_model_mappings_array(
array: &[serde_json::Value],
) -> Result<Option<Vec<StoredProviderModelMapping>>, DataLayerError> {
let mut mappings = Vec::with_capacity(array.len());
for raw in array {
match raw {
serde_json::Value::Object(object) => {
if let Some(mapping) = parse_provider_model_mapping_object_lenient(object)? {
mappings.push(mapping);
}
}
serde_json::Value::String(raw) => {
let raw = raw.trim();
if !raw.is_empty() {
mappings.push(StoredProviderModelMapping {
name: raw.to_string(),
priority: 1,
api_formats: None,
endpoint_ids: None,
operations: None,
});
}
}
serde_json::Value::Null => {}
_ => {}
}
}
if mappings.is_empty() {
Ok(None)
} else {
Ok(Some(mappings))
}
}
fn parse_provider_model_mapping_object(
object: &serde_json::Map<String, serde_json::Value>,
) -> Result<StoredProviderModelMapping, DataLayerError> {
parse_provider_model_mapping_object_lenient(object)?.ok_or_else(|| {
DataLayerError::UnexpectedValue(
"models.provider_model_mappings item is missing a valid name".to_string(),
)
})
}
fn parse_provider_model_mapping_object_lenient(
object: &serde_json::Map<String, serde_json::Value>,
) -> Result<Option<StoredProviderModelMapping>, DataLayerError> {
let Some(name) = object
.get("name")
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
else {
return Ok(None);
};
let priority = object
.get("priority")
.and_then(serde_json::Value::as_i64)
.unwrap_or(1)
.max(1);
let api_formats = parse_string_list(
object.get("api_formats").cloned(),
"models.provider_model_mappings.api_formats",
)?
.map(|formats| {
formats
.into_iter()
.map(|value| aether_ai_formats::normalize_api_format_alias(&value))
.collect()
});
let endpoint_ids = parse_string_list(
object.get("endpoint_ids").cloned(),
"models.provider_model_mappings.endpoint_ids",
)?;
let operations = parse_string_list(
object.get("operations").cloned(),
"models.provider_model_mappings.operations",
)?
.and_then(normalize_request_operations);
Ok(Some(StoredProviderModelMapping {
name: name.to_string(),
priority: i32::try_from(priority).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid models.provider_model_mappings.priority: {priority}"
))
})?,
api_formats,
endpoint_ids,
operations,
}))
}
fn normalize_request_operations(values: Vec<String>) -> Option<Vec<String>> {
let operations = values
.into_iter()
.map(|value| value.trim().to_ascii_lowercase())
.filter(|value| !value.is_empty())
.collect::<Vec<_>>();
(!operations.is_empty()).then_some(operations)
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::{
parse_provider_model_mappings, parse_string_list, pool_key_candidate_selection_sql,
requested_model_selection_page_sql, requested_model_selection_sql,
SqlxMinimalCandidateSelectionReadRepository,
LIST_FOR_EXACT_API_FORMAT_AND_GLOBAL_MODEL_SQL, LIST_FOR_EXACT_API_FORMAT_SQL,
LIST_POOL_KEYS_FOR_GROUP_SQL,
};
use crate::{PostgresPoolConfig, PostgresPoolFactory};
use aether_data_contracts::repository::candidate_selection::{
StoredPoolKeyCandidateOrder, StoredProviderModelMapping,
};
#[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 = SqlxMinimalCandidateSelectionReadRepository::new(pool);
let _ = repository.pool();
}
#[test]
fn requested_model_selection_sql_filters_before_row_materialization() {
let sql = requested_model_selection_sql();
assert!(sql.contains("m.provider_model_name = $2"));
assert!(sql.contains("jsonb_array_elements("));
assert!(!sql.contains("json_typeof(m.provider_model_mappings)"));
assert!(!sql.contains("json_array_elements_text(gm.config -> 'model_mappings')"));
assert!(sql.contains("ORDER BY\n global_model_name ASC,"));
assert!(!sql.contains("AND gm.name = $2\n AND"));
}
#[test]
fn exact_candidate_selection_sql_limits_pool_key_expansion() {
for sql in [
LIST_FOR_EXACT_API_FORMAT_SQL,
LIST_FOR_EXACT_API_FORMAT_AND_GLOBAL_MODEL_SQL,
] {
assert!(sql.contains("INNER JOIN LATERAL"));
assert!(sql.contains("FROM provider_api_keys pak"));
assert!(sql.contains("AND pak.is_active IS TRUE"));
assert!(!sql.contains("AND pak.is_active = TRUE"));
assert!(sql.contains(
"LIMIT CASE WHEN (p.config -> 'pool_advanced') IS NOT NULL THEN 1 ELSE 2147483647 END"
));
}
assert!(!LIST_POOL_KEYS_FOR_GROUP_SQL.contains("INNER JOIN LATERAL"));
assert!(LIST_POOL_KEYS_FOR_GROUP_SQL.contains("AND pak.is_active IS TRUE"));
assert!(!LIST_POOL_KEYS_FOR_GROUP_SQL.contains("AND pak.is_active = TRUE"));
}
#[test]
fn candidate_selection_sql_allows_chatgpt_web_image_auth() {
let requested_model_sql = requested_model_selection_sql();
for sql in [
LIST_FOR_EXACT_API_FORMAT_SQL,
LIST_FOR_EXACT_API_FORMAT_AND_GLOBAL_MODEL_SQL,
LIST_POOL_KEYS_FOR_GROUP_SQL,
requested_model_sql.as_str(),
] {
assert!(sql.contains("LOWER(BTRIM(p.provider_type)) = 'chatgpt_web'"));
assert!(sql.contains("LOWER(BTRIM(pak.auth_type)) IN ('oauth', 'bearer')"));
assert!(sql.contains("'chatgpt_web',"));
}
}
#[test]
fn candidate_selection_sql_allows_grok_oauth_chat_auth() {
let requested_model_sql = requested_model_selection_sql();
for sql in [
LIST_FOR_EXACT_API_FORMAT_SQL,
LIST_FOR_EXACT_API_FORMAT_AND_GLOBAL_MODEL_SQL,
LIST_POOL_KEYS_FOR_GROUP_SQL,
requested_model_sql.as_str(),
] {
assert!(sql.contains("LOWER(BTRIM(p.provider_type)) = 'grok'"));
assert!(sql.contains("LOWER(BTRIM(pak.auth_type)) = 'oauth'"));
assert!(sql
.contains("'openai:chat', 'openai:responses', 'claude:messages', 'openai:image'"));
assert!(sql.contains("'grok',"));
}
}
#[test]
fn candidate_selection_sql_allows_windsurf_openai_chat_managed_keys() {
let requested_model_sql = requested_model_selection_sql();
for sql in [
LIST_FOR_EXACT_API_FORMAT_SQL,
LIST_FOR_EXACT_API_FORMAT_AND_GLOBAL_MODEL_SQL,
LIST_POOL_KEYS_FOR_GROUP_SQL,
requested_model_sql.as_str(),
] {
assert!(sql.contains("LOWER(BTRIM(p.provider_type)) = 'windsurf'"));
assert!(
sql.contains("LOWER($3) = 'openai:chat'")
|| sql.contains("LOWER($4) = 'openai:chat'")
|| sql.contains("LOWER($6) = 'openai:chat'")
);
assert!(sql.contains("LOWER(BTRIM(pak.auth_type)) IN ('oauth', 'api_key', 'bearer')"));
assert!(sql.contains("'windsurf'"));
}
}
#[test]
fn candidate_selection_sql_rejects_retired_vertex_claude_auth() {
let requested_model_sql = requested_model_selection_sql();
for (name, sql, expected_occurrences) in [
("exact", LIST_FOR_EXACT_API_FORMAT_SQL, 2),
(
"global_model",
LIST_FOR_EXACT_API_FORMAT_AND_GLOBAL_MODEL_SQL,
2,
),
("pool_keys", LIST_POOL_KEYS_FOR_GROUP_SQL, 1),
("requested_model", requested_model_sql.as_str(), 2),
] {
let mut remaining = sql;
let mut occurrences = 0;
while let Some((_, suffix)) =
remaining.split_once("LOWER(BTRIM(p.provider_type)) = 'vertex_ai'")
{
let (vertex_clause, rest) = suffix
.split_once("LOWER(BTRIM(p.provider_type)) NOT IN")
.expect("each Vertex auth clause should have a following fallback clause");
assert!(vertex_clause.contains(
"LOWER(BTRIM(pak.auth_type)) IN ('api_key', 'service_account', 'vertex_ai')"
));
assert!(vertex_clause.contains("gemini:embedding"));
assert!(vertex_clause.contains("gemini:generate_content"));
assert!(!vertex_clause.contains("claude:messages"));
occurrences += 1;
remaining = rest;
}
assert_eq!(
occurrences, expected_occurrences,
"unexpected Vertex auth clause count in {name} SQL"
);
}
}
#[test]
fn requested_model_selection_page_sql_adds_limit_and_offset() {
let sql = requested_model_selection_page_sql();
assert!(sql.ends_with("LIMIT $5\nOFFSET $6"));
}
#[test]
fn pool_key_selection_sql_applies_query_order() {
let load_balance_sql =
pool_key_candidate_selection_sql(&StoredPoolKeyCandidateOrder::LoadBalance {
seed: "seed".to_string(),
});
assert!(load_balance_sql.contains("md5($9 || ':' || pak.id) ASC"));
assert!(load_balance_sql.ends_with("LIMIT $7\nOFFSET $8\n"));
let lru_sql = pool_key_candidate_selection_sql(&StoredPoolKeyCandidateOrder::Lru);
assert!(lru_sql.contains("pak.last_used_at ASC NULLS FIRST"));
let cache_affinity_sql =
pool_key_candidate_selection_sql(&StoredPoolKeyCandidateOrder::CacheAffinity);
assert!(cache_affinity_sql.contains("pak.last_used_at DESC NULLS LAST"));
}
#[test]
fn parse_string_list_accepts_stringified_array() {
let parsed = parse_string_list(
Some(json!("[\"gpt-5.2\", \"gpt-5\"]")),
"provider_api_keys.allowed_models",
)
.expect("stringified array should parse");
assert_eq!(
parsed,
Some(vec!["gpt-5.2".to_string(), "gpt-5".to_string()])
);
}
#[test]
fn parse_string_list_accepts_single_string() {
let parsed = parse_string_list(Some(json!("gpt-5.2")), "provider_api_keys.allowed_models")
.expect("single string should parse");
assert_eq!(parsed, Some(vec!["gpt-5.2".to_string()]));
}
#[test]
fn parse_provider_model_mappings_accepts_stringified_array() {
let parsed = parse_provider_model_mappings(Some(json!(
"[{\"name\":\"gpt-5.2\",\"priority\":2,\"api_formats\":[\"openai:chat\"],\"operations\":[\"COMPACT\"]}]"
)))
.expect("stringified provider_model_mappings should parse");
assert_eq!(
parsed,
Some(vec![StoredProviderModelMapping {
name: "gpt-5.2".to_string(),
priority: 2,
api_formats: Some(vec!["openai:chat".to_string()]),
endpoint_ids: None,
operations: Some(vec!["compact".to_string()]),
}])
);
}
#[test]
fn parse_provider_model_mappings_accepts_single_string_alias() {
let parsed = parse_provider_model_mappings(Some(json!("gpt-5.2")))
.expect("single-string provider_model_mappings should parse");
assert_eq!(
parsed,
Some(vec![StoredProviderModelMapping {
name: "gpt-5.2".to_string(),
priority: 1,
api_formats: None,
endpoint_ids: None,
operations: None,
}])
);
}
#[test]
fn parse_provider_model_mappings_skips_invalid_array_items() {
let parsed = parse_provider_model_mappings(Some(json!([
{"name": "gpt-5.2", "priority": 1},
{"priority": 2},
3,
null,
"gpt-5.2-mini"
])))
.expect("mixed provider_model_mappings should parse");
assert_eq!(
parsed,
Some(vec![
StoredProviderModelMapping {
name: "gpt-5.2".to_string(),
priority: 1,
api_formats: None,
endpoint_ids: None,
operations: None,
},
StoredProviderModelMapping {
name: "gpt-5.2-mini".to_string(),
priority: 1,
api_formats: None,
endpoint_ids: None,
operations: None,
}
])
);
}
}