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

Optimize provider scheduler database hotspots

Browse Source
This commit is contained in:
fawney19
2026-05-24 15:47:56 +08:00
parent 78523f122d
commit 576918daa5
25 changed files with 836 additions and 109 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-data/migrations/mysql/20260524000000_add_scheduler_hotspot_indexes.sql
+31
View File
@@ -0,0 +1,31 @@
SET @aether_provider_key_active_priority_index_sql := IF(
(
SELECT COUNT(*)
FROM information_schema.statistics
WHERE table_schema = DATABASE()
AND table_name = 'provider_api_keys'
AND index_name = 'idx_provider_api_keys_provider_active_priority_id'
) = 0,
'CREATE INDEX idx_provider_api_keys_provider_active_priority_id ON provider_api_keys (provider_id, is_active, internal_priority, id)',
'DO 0'
);
PREPARE aether_provider_key_active_priority_index_stmt FROM @aether_provider_key_active_priority_index_sql;
EXECUTE aether_provider_key_active_priority_index_stmt;
DEALLOCATE PREPARE aether_provider_key_active_priority_index_stmt;
SET @aether_pool_score_scheduler_rank_index_sql := IF(
(
SELECT COUNT(*)
FROM information_schema.statistics
WHERE table_schema = DATABASE()
AND table_name = 'pool_member_scores'
AND index_name = 'pool_member_scores_scheduler_account_rank_idx'
) = 0,
'CREATE INDEX pool_member_scores_scheduler_account_rank_idx ON pool_member_scores (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC, last_ranked_at DESC, member_id, id)',
'DO 0'
);
PREPARE aether_pool_score_scheduler_rank_index_stmt FROM @aether_pool_score_scheduler_rank_index_sql;
EXECUTE aether_pool_score_scheduler_rank_index_stmt;
DEALLOCATE PREPARE aether_pool_score_scheduler_rank_index_stmt;
crates/aether-data/migrations/postgres/20260524000000_add_scheduler_hotspot_indexes.sql
+16
View File
@@ -0,0 +1,16 @@
CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_active_priority_id
ON public.provider_api_keys USING btree (provider_id, internal_priority, id)
WHERE is_active IS TRUE;
CREATE INDEX IF NOT EXISTS pool_member_scores_scheduler_account_rank_idx
ON public.pool_member_scores USING btree (
pool_kind,
pool_id,
capability,
scope_kind,
score DESC,
last_ranked_at DESC NULLS LAST,
member_id,
id
)
WHERE scope_id IS NULL AND hard_state IN ('available', 'unknown');
crates/aether-data/migrations/sqlite/20260524000000_add_scheduler_hotspot_indexes.sql
+16
View File
@@ -0,0 +1,16 @@
CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_active_priority_id
ON provider_api_keys (provider_id, is_active, internal_priority, id);
CREATE INDEX IF NOT EXISTS pool_member_scores_scheduler_account_rank_idx
ON pool_member_scores (
pool_kind,
pool_id,
capability,
scope_kind,
scope_id,
hard_state,
score DESC,
last_ranked_at DESC,
member_id,
id
);
crates/aether-data/schema/bootstrap/postgres/004_indexes.sql
+16
View File
@@ -197,6 +197,14 @@ CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_default_sort ON public
--
-- Name: idx_provider_api_keys_provider_active_priority_id; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_active_priority_id ON public.provider_api_keys USING btree (provider_id, internal_priority, id) WHERE is_active IS TRUE;
--
-- Name: idx_provider_api_keys_provider_created_at_desc; Type: INDEX; Schema: public; Owner: -
--
@@ -1037,6 +1045,14 @@ CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON public.pool_member_sco
--
-- Name: pool_member_scores_scheduler_account_rank_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_scheduler_account_rank_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, capability, scope_kind, score DESC, last_ranked_at DESC NULLS LAST, member_id, id) WHERE scope_id IS NULL AND hard_state IN ('available', 'unknown');
--
-- Name: pool_member_scores_member_idx; Type: INDEX; Schema: public; Owner: -
--
crates/aether-data/src/backend/stats/postgres_daily/sql.rs
+24 -18
View File
@@ -1,8 +1,14 @@
pub(super) const SELECT_STATS_DAILY_AGGREGATE_SQL: &str = r#"
WITH bucket_usage AS MATERIALIZED (
SELECT *
FROM usage_billing_facts
WHERE created_at >= $1
AND created_at < $2
)
SELECT
(
SELECT CAST(COUNT(cache_hit_usage.id) AS BIGINT)
FROM usage_billing_facts AS cache_hit_usage
FROM bucket_usage AS cache_hit_usage
WHERE cache_hit_usage.created_at >= $1
AND cache_hit_usage.created_at < $2
) AS cache_hit_total_requests,
@@ -12,13 +18,13 @@ SELECT
WHERE GREATEST(COALESCE(cache_hit_usage.cache_read_input_tokens, 0), 0) > 0
) AS BIGINT
)
FROM usage_billing_facts AS cache_hit_usage
FROM bucket_usage AS cache_hit_usage
WHERE cache_hit_usage.created_at >= $1
AND cache_hit_usage.created_at < $2
) AS cache_hit_requests,
(
SELECT CAST(COUNT(completed_usage.id) AS BIGINT)
FROM usage_billing_facts AS completed_usage
FROM bucket_usage AS completed_usage
WHERE completed_usage.created_at >= $1
AND completed_usage.created_at < $2
AND completed_usage.status = 'completed'
@@ -29,7 +35,7 @@ SELECT
WHERE GREATEST(COALESCE(completed_usage.cache_read_input_tokens, 0), 0) > 0
) AS BIGINT
)
FROM usage_billing_facts AS completed_usage
FROM bucket_usage AS completed_usage
WHERE completed_usage.created_at >= $1
AND completed_usage.created_at < $2
AND completed_usage.status = 'completed'
@@ -38,7 +44,7 @@ SELECT
SELECT CAST(
COALESCE(SUM(GREATEST(COALESCE(completed_usage.input_tokens, 0), 0)), 0) AS BIGINT
)
FROM usage_billing_facts AS completed_usage
FROM bucket_usage AS completed_usage
WHERE completed_usage.created_at >= $1
AND completed_usage.created_at < $2
AND completed_usage.status = 'completed'
@@ -61,7 +67,7 @@ SELECT
0
) AS BIGINT
)
FROM usage_billing_facts AS completed_usage
FROM bucket_usage AS completed_usage
WHERE completed_usage.created_at >= $1
AND completed_usage.created_at < $2
AND completed_usage.status = 'completed'
@@ -73,7 +79,7 @@ SELECT
0
) AS BIGINT
)
FROM usage_billing_facts AS completed_usage
FROM bucket_usage AS completed_usage
WHERE completed_usage.created_at >= $1
AND completed_usage.created_at < $2
AND completed_usage.status = 'completed'
@@ -211,7 +217,7 @@ SELECT
0
) AS BIGINT
)
FROM usage_billing_facts AS completed_usage
FROM bucket_usage AS completed_usage
WHERE completed_usage.created_at >= $1
AND completed_usage.created_at < $2
AND completed_usage.status = 'completed'
@@ -228,7 +234,7 @@ SELECT
0
) AS DOUBLE PRECISION
)
FROM usage_billing_facts AS completed_usage
FROM bucket_usage AS completed_usage
WHERE completed_usage.created_at >= $1
AND completed_usage.created_at < $2
AND completed_usage.status = 'completed'
@@ -242,7 +248,7 @@ SELECT
0
) AS DOUBLE PRECISION
)
FROM usage_billing_facts AS completed_usage
FROM bucket_usage AS completed_usage
WHERE completed_usage.created_at >= $1
AND completed_usage.created_at < $2
AND completed_usage.status = 'completed'
@@ -252,7 +258,7 @@ SELECT
COALESCE(SUM(COALESCE(CAST(settled_usage.total_cost_usd AS DOUBLE PRECISION), 0)), 0)
AS DOUBLE PRECISION
)
FROM usage_billing_facts AS settled_usage
FROM bucket_usage AS settled_usage
WHERE settled_usage.created_at >= $1
AND settled_usage.created_at < $2
AND settled_usage.billing_status = 'settled'
@@ -260,7 +266,7 @@ SELECT
) AS settled_total_cost,
(
SELECT CAST(COUNT(settled_usage.id) AS BIGINT)
FROM usage_billing_facts AS settled_usage
FROM bucket_usage AS settled_usage
WHERE settled_usage.created_at >= $1
AND settled_usage.created_at < $2
AND settled_usage.billing_status = 'settled'
@@ -270,7 +276,7 @@ SELECT
SELECT CAST(
COALESCE(SUM(GREATEST(COALESCE(settled_usage.input_tokens, 0), 0)), 0) AS BIGINT
)
FROM usage_billing_facts AS settled_usage
FROM bucket_usage AS settled_usage
WHERE settled_usage.created_at >= $1
AND settled_usage.created_at < $2
AND settled_usage.billing_status = 'settled'
@@ -280,7 +286,7 @@ SELECT
SELECT CAST(
COALESCE(SUM(GREATEST(COALESCE(settled_usage.output_tokens, 0), 0)), 0) AS BIGINT
)
FROM usage_billing_facts AS settled_usage
FROM bucket_usage AS settled_usage
WHERE settled_usage.created_at >= $1
AND settled_usage.created_at < $2
AND settled_usage.billing_status = 'settled'
@@ -293,7 +299,7 @@ SELECT
0
) AS BIGINT
)
FROM usage_billing_facts AS settled_usage
FROM bucket_usage AS settled_usage
WHERE settled_usage.created_at >= $1
AND settled_usage.created_at < $2
AND settled_usage.billing_status = 'settled'
@@ -306,7 +312,7 @@ SELECT
0
) AS BIGINT
)
FROM usage_billing_facts AS settled_usage
FROM bucket_usage AS settled_usage
WHERE settled_usage.created_at >= $1
AND settled_usage.created_at < $2
AND settled_usage.billing_status = 'settled'
@@ -314,7 +320,7 @@ SELECT
) AS settled_cache_read_tokens,
(
SELECT MIN(CAST(EXTRACT(EPOCH FROM settled_usage.finalized_at) AS BIGINT))
FROM usage_billing_facts AS settled_usage
FROM bucket_usage AS settled_usage
WHERE settled_usage.created_at >= $1
AND settled_usage.created_at < $2
AND settled_usage.billing_status = 'settled'
@@ -322,7 +328,7 @@ SELECT
) AS settled_first_finalized_at_unix_secs,
(
SELECT MAX(CAST(EXTRACT(EPOCH FROM settled_usage.finalized_at) AS BIGINT))
FROM usage_billing_facts AS settled_usage
FROM bucket_usage AS settled_usage
WHERE settled_usage.created_at >= $1
AND settled_usage.created_at < $2
AND settled_usage.billing_status = 'settled'
crates/aether-data/src/backend/stats/postgres_hourly/sql.rs
+24 -18
View File
@@ -12,10 +12,16 @@ WHERE created_at >= $1
AND provider_name NOT IN ('unknown', 'pending')
"#;
pub(super) const SELECT_STATS_HOURLY_AGGREGATE_SQL: &str = r#"
WITH bucket_usage AS MATERIALIZED (
SELECT *
FROM usage_billing_facts
WHERE created_at >= $1
AND created_at < $2
)
SELECT
(
SELECT CAST(COUNT(cache_hit_usage.id) AS BIGINT)
FROM usage_billing_facts AS cache_hit_usage
FROM bucket_usage AS cache_hit_usage
WHERE cache_hit_usage.created_at >= $1
AND cache_hit_usage.created_at < $2
) AS cache_hit_total_requests,
@@ -25,13 +31,13 @@ SELECT
WHERE GREATEST(COALESCE(cache_hit_usage.cache_read_input_tokens, 0), 0) > 0
) AS BIGINT
)
FROM usage_billing_facts AS cache_hit_usage
FROM bucket_usage AS cache_hit_usage
WHERE cache_hit_usage.created_at >= $1
AND cache_hit_usage.created_at < $2
) AS cache_hit_requests,
(
SELECT CAST(COUNT(completed_usage.id) AS BIGINT)
FROM usage_billing_facts AS completed_usage
FROM bucket_usage AS completed_usage
WHERE completed_usage.created_at >= $1
AND completed_usage.created_at < $2
AND completed_usage.status = 'completed'
@@ -42,7 +48,7 @@ SELECT
WHERE GREATEST(COALESCE(completed_usage.cache_read_input_tokens, 0), 0) > 0
) AS BIGINT
)
FROM usage_billing_facts AS completed_usage
FROM bucket_usage AS completed_usage
WHERE completed_usage.created_at >= $1
AND completed_usage.created_at < $2
AND completed_usage.status = 'completed'
@@ -51,7 +57,7 @@ SELECT
SELECT CAST(
COALESCE(SUM(GREATEST(COALESCE(completed_usage.input_tokens, 0), 0)), 0) AS BIGINT
)
FROM usage_billing_facts AS completed_usage
FROM bucket_usage AS completed_usage
WHERE completed_usage.created_at >= $1
AND completed_usage.created_at < $2
AND completed_usage.status = 'completed'
@@ -74,7 +80,7 @@ SELECT
0
) AS BIGINT
)
FROM usage_billing_facts AS completed_usage
FROM bucket_usage AS completed_usage
WHERE completed_usage.created_at >= $1
AND completed_usage.created_at < $2
AND completed_usage.status = 'completed'
@@ -86,7 +92,7 @@ SELECT
0
) AS BIGINT
)
FROM usage_billing_facts AS completed_usage
FROM bucket_usage AS completed_usage
WHERE completed_usage.created_at >= $1
AND completed_usage.created_at < $2
AND completed_usage.status = 'completed'
@@ -224,7 +230,7 @@ SELECT
0
) AS BIGINT
)
FROM usage_billing_facts AS completed_usage
FROM bucket_usage AS completed_usage
WHERE completed_usage.created_at >= $1
AND completed_usage.created_at < $2
AND completed_usage.status = 'completed'
@@ -241,7 +247,7 @@ SELECT
0
) AS DOUBLE PRECISION
)
FROM usage_billing_facts AS completed_usage
FROM bucket_usage AS completed_usage
WHERE completed_usage.created_at >= $1
AND completed_usage.created_at < $2
AND completed_usage.status = 'completed'
@@ -255,7 +261,7 @@ SELECT
0
) AS DOUBLE PRECISION
)
FROM usage_billing_facts AS completed_usage
FROM bucket_usage AS completed_usage
WHERE completed_usage.created_at >= $1
AND completed_usage.created_at < $2
AND completed_usage.status = 'completed'
@@ -265,7 +271,7 @@ SELECT
COALESCE(SUM(COALESCE(CAST(settled_usage.total_cost_usd AS DOUBLE PRECISION), 0)), 0)
AS DOUBLE PRECISION
)
FROM usage_billing_facts AS settled_usage
FROM bucket_usage AS settled_usage
WHERE settled_usage.created_at >= $1
AND settled_usage.created_at < $2
AND settled_usage.billing_status = 'settled'
@@ -273,7 +279,7 @@ SELECT
) AS settled_total_cost,
(
SELECT CAST(COUNT(settled_usage.id) AS BIGINT)
FROM usage_billing_facts AS settled_usage
FROM bucket_usage AS settled_usage
WHERE settled_usage.created_at >= $1
AND settled_usage.created_at < $2
AND settled_usage.billing_status = 'settled'
@@ -283,7 +289,7 @@ SELECT
SELECT CAST(
COALESCE(SUM(GREATEST(COALESCE(settled_usage.input_tokens, 0), 0)), 0) AS BIGINT
)
FROM usage_billing_facts AS settled_usage
FROM bucket_usage AS settled_usage
WHERE settled_usage.created_at >= $1
AND settled_usage.created_at < $2
AND settled_usage.billing_status = 'settled'
@@ -293,7 +299,7 @@ SELECT
SELECT CAST(
COALESCE(SUM(GREATEST(COALESCE(settled_usage.output_tokens, 0), 0)), 0) AS BIGINT
)
FROM usage_billing_facts AS settled_usage
FROM bucket_usage AS settled_usage
WHERE settled_usage.created_at >= $1
AND settled_usage.created_at < $2
AND settled_usage.billing_status = 'settled'
@@ -306,7 +312,7 @@ SELECT
0
) AS BIGINT
)
FROM usage_billing_facts AS settled_usage
FROM bucket_usage AS settled_usage
WHERE settled_usage.created_at >= $1
AND settled_usage.created_at < $2
AND settled_usage.billing_status = 'settled'
@@ -319,7 +325,7 @@ SELECT
0
) AS BIGINT
)
FROM usage_billing_facts AS settled_usage
FROM bucket_usage AS settled_usage
WHERE settled_usage.created_at >= $1
AND settled_usage.created_at < $2
AND settled_usage.billing_status = 'settled'
@@ -327,7 +333,7 @@ SELECT
) AS settled_cache_read_tokens,
(
SELECT MIN(CAST(EXTRACT(EPOCH FROM settled_usage.finalized_at) AS BIGINT))
FROM usage_billing_facts AS settled_usage
FROM bucket_usage AS settled_usage
WHERE settled_usage.created_at >= $1
AND settled_usage.created_at < $2
AND settled_usage.billing_status = 'settled'
@@ -335,7 +341,7 @@ SELECT
) AS settled_first_finalized_at_unix_secs,
(
SELECT MAX(CAST(EXTRACT(EPOCH FROM settled_usage.finalized_at) AS BIGINT))
FROM usage_billing_facts AS settled_usage
FROM bucket_usage AS settled_usage
WHERE settled_usage.created_at >= $1
AND settled_usage.created_at < $2
AND settled_usage.billing_status = 'settled'
crates/aether-data/src/lifecycle/bootstrap/postgres.rs
+1 -1
View File
@@ -7,7 +7,7 @@ use tracing::info;
// Generated by build.rs from schema/bootstrap/postgres.
pub(crate) static EMPTY_DATABASE_SNAPSHOT_SQL: &str =
include_str!(concat!(env!("OUT_DIR"), "/empty_database_snapshot.sql"));
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260522000000;
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260524000000;
const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#"
SELECT COUNT(*)::BIGINT
crates/aether-data/src/lifecycle/migrate/tests.rs
+8
View File
@@ -313,6 +313,7 @@ fn empty_database_snapshot_covers_current_cutoff_versions() {
20260520000000,
20260520010000,
20260522000000,
20260524000000,
]
);
}
@@ -389,6 +390,10 @@ fn empty_database_snapshot_sql_includes_usage_body_blobs_and_audit_admin_role()
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("ix_usage_counter_deltas_unprocessed"));
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("idx_entitlement_usage_entitlement_date"));
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("idx_provider_api_keys_provider_default_sort"));
assert!(
EMPTY_DATABASE_SNAPSHOT_SQL.contains("idx_provider_api_keys_provider_active_priority_id")
);
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("pool_member_scores_scheduler_account_rank_idx"));
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("idx_video_tasks_due_poll"));
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("request_count bigint DEFAULT 0"));
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("usage_count bigint DEFAULT 0 NOT NULL"));
@@ -662,6 +667,7 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
20260519130000,
20260520000000,
20260520010000,
20260524000000,
]
);
assert_eq!(
@@ -685,6 +691,7 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
20260519130000,
20260520000000,
20260520010000,
20260524000000,
]
);
}
@@ -1209,6 +1216,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
20260520000000,
20260520010000,
20260522000000,
20260524000000,
]
);
}
crates/aether-data/src/repository/candidate_selection/postgres.rs
+196 -7
View File
@@ -54,15 +54,100 @@ SELECT
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 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: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)) = 'api_key'
AND LOWER($3) IN ('gemini:generate_content', 'gemini:embedding')
)
OR (
LOWER(BTRIM(pak.auth_type)) IN ('service_account', 'vertex_ai')
AND LOWER($3) IN ('claude:messages', '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 = TRUE
AND pak.is_active IS TRUE
AND m.is_active = TRUE
AND m.is_available = TRUE
AND gm.is_active = TRUE
@@ -248,15 +333,100 @@ SELECT
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 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: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)) = 'api_key'
AND LOWER($4) IN ('gemini:generate_content', 'gemini:embedding')
)
OR (
LOWER(BTRIM(pak.auth_type)) IN ('service_account', 'vertex_ai')
AND LOWER($4) IN ('claude:messages', '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 = TRUE
AND pak.is_active IS TRUE
AND m.is_active = TRUE
AND m.is_available = TRUE
AND gm.is_active = TRUE
@@ -448,7 +618,7 @@ 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 = TRUE
AND pak.is_active IS TRUE
AND m.is_active = TRUE
AND m.is_available = TRUE
AND gm.is_active = TRUE
@@ -1308,6 +1478,25 @@ mod tests {
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();
crates/aether-data/src/repository/provider_catalog/memory.rs
+30 -1
View File
@@ -8,7 +8,8 @@ use serde_json::{json, Map, Value};
use super::{
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogReadRepository,
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
};
use crate::repository::usage::{ProviderApiKeyUsageContribution, ProviderApiKeyUsageDelta};
use crate::DataLayerError;
@@ -431,6 +432,34 @@ impl ProviderCatalogReadRepository for InMemoryProviderCatalogReadRepository {
Self::list_keys_by_provider_ids(self, provider_ids).await
}
async fn list_key_maintenance_summaries_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKeyMaintenanceSummary>, DataLayerError> {
let index = self.index.read().expect("provider catalog repository lock");
let mut keys = index
.keys
.values()
.filter(|key| {
provider_ids
.iter()
.any(|provider_id| provider_id == &key.provider_id)
})
.map(|key| StoredProviderCatalogKeyMaintenanceSummary {
id: key.id.clone(),
provider_id: key.provider_id.clone(),
is_active: key.is_active,
upstream_metadata: key.upstream_metadata.clone(),
})
.collect::<Vec<_>>();
keys.sort_by(|left, right| {
left.provider_id
.cmp(&right.provider_id)
.then(left.id.cmp(&right.id))
});
Ok(keys)
}
async fn list_keys_page(
&self,
query: &ProviderCatalogKeyListQuery,
crates/aether-data/src/repository/provider_catalog/mod.rs
+2 -1
View File
@@ -7,7 +7,8 @@ mod sqlite;
pub(crate) use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogReadRepository,
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
};
pub use memory::InMemoryProviderCatalogReadRepository;
pub use mysql::MysqlProviderCatalogReadRepository;
crates/aether-data/src/repository/provider_catalog/mysql.rs
+12 -2
View File
@@ -4,8 +4,8 @@ use sqlx::{mysql::MySqlRow, Row};
use super::{
InMemoryProviderCatalogReadRepository, ProviderCatalogKeyListQuery,
ProviderCatalogReadRepository, ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint,
StoredProviderCatalogKey, StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats,
StoredProviderCatalogProvider,
StoredProviderCatalogKey, StoredProviderCatalogKeyMaintenanceSummary,
StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
};
use crate::driver::mysql::MysqlPool;
use crate::error::SqlResultExt;
@@ -928,6 +928,16 @@ impl ProviderCatalogReadRepository for MysqlProviderCatalogReadRepository {
.await
}
async fn list_key_maintenance_summaries_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKeyMaintenanceSummary>, DataLayerError> {
self.load_memory()
.await?
.list_key_maintenance_summaries_by_provider_ids(provider_ids)
.await
}
async fn list_keys_page(
&self,
query: &ProviderCatalogKeyListQuery,
crates/aether-data/src/repository/provider_catalog/postgres.rs
+51 -1
View File
@@ -5,7 +5,8 @@ use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row};
use super::{
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogReadRepository,
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
};
use crate::{
error::{postgres_error, SqlxResultExt},
@@ -323,6 +324,16 @@ 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 (
"#;
#[derive(Debug, Clone)]
pub struct SqlxProviderCatalogReadRepository {
pool: PgPool,
@@ -513,6 +524,27 @@ impl SqlxProviderCatalogReadRepository {
.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,
@@ -1917,6 +1949,13 @@ impl ProviderCatalogReadRepository for SqlxProviderCatalogReadRepository {
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,
@@ -2255,6 +2294,17 @@ fn map_key_stats_row(row: &PgRow) -> Result<StoredProviderCatalogKeyStats, DataL
)
}
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 rpm_limit = row_get::<Option<i32>>(row, "rpm_limit")?
.map(|value| {
crates/aether-data/src/repository/provider_catalog/sqlite.rs
+53 -1
View File
@@ -4,7 +4,8 @@ use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite};
use super::{
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogReadRepository,
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
};
use crate::driver::sqlite::{sqlite_optional_real, SqlitePool};
use crate::error::SqlResultExt;
@@ -278,6 +279,16 @@ 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 (
"#;
#[derive(Debug, Clone)]
pub struct SqliteProviderCatalogReadRepository {
pool: SqlitePool,
@@ -423,6 +434,26 @@ impl SqliteProviderCatalogReadRepository {
rows.iter().map(map_key_row).collect()
}
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());
}
let rows = build_list_query(
LIST_KEY_MAINTENANCE_SUMMARIES_BY_PROVIDER_IDS_PREFIX,
provider_ids,
" ORDER BY provider_id ASC, id ASC",
)
.build()
.fetch_all(&self.pool)
.await
.map_sql_err()?;
rows.iter().map(map_key_maintenance_summary_row).collect()
}
pub async fn list_keys_page(
&self,
query: &ProviderCatalogKeyListQuery,
@@ -1322,6 +1353,13 @@ impl ProviderCatalogReadRepository for SqliteProviderCatalogReadRepository {
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,
@@ -1805,6 +1843,20 @@ fn map_key_stats_row(row: &SqliteRow) -> Result<StoredProviderCatalogKeyStats, D
)
}
fn map_key_maintenance_summary_row(
row: &SqliteRow,
) -> Result<StoredProviderCatalogKeyMaintenanceSummary, DataLayerError> {
Ok(StoredProviderCatalogKeyMaintenanceSummary {
id: row.try_get("id").map_sql_err()?,
provider_id: row.try_get("provider_id").map_sql_err()?,
is_active: row.try_get("is_active").map_sql_err()?,
upstream_metadata: optional_json_from_string(
row.try_get("upstream_metadata").map_sql_err()?,
"provider_api_keys.upstream_metadata",
)?,
})
}
fn map_key_row(row: &SqliteRow) -> Result<StoredProviderCatalogKey, DataLayerError> {
let total_cost_usd = sqlite_optional_real(row, "total_cost_usd")?.unwrap_or(0.0);
if !total_cost_usd.is_finite() {