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

Cache provider catalog lookups

Browse Source
This commit is contained in:
fawney19
2026-05-27 13:56:39 +08:00
parent 674cc85005
commit ccfc4cbddc
16 changed files with 582 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-data/migrations/mysql/20260527000000_add_provider_key_name_hotspot_index.sql
+15
View File
@@ -0,0 +1,15 @@
SET @aether_provider_key_name_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_name_id'
) = 0,
'CREATE INDEX idx_provider_api_keys_provider_name_id ON provider_api_keys (provider_id, name, id)',
'DO 0'
);
PREPARE aether_provider_key_name_index_stmt FROM @aether_provider_key_name_index_sql;
EXECUTE aether_provider_key_name_index_stmt;
DEALLOCATE PREPARE aether_provider_key_name_index_stmt;
crates/aether-data/migrations/postgres/20260527000000_add_provider_key_name_hotspot_index.sql
+2
View File
@@ -0,0 +1,2 @@
CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_name_id
ON public.provider_api_keys USING btree (provider_id, name, id);
crates/aether-data/migrations/sqlite/20260527000000_add_provider_key_name_hotspot_index.sql
+2
View File
@@ -0,0 +1,2 @@
CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_name_id
ON provider_api_keys (provider_id, name, id);
crates/aether-data/schema/bootstrap/postgres/004_indexes.sql
+8
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_name_id; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_name_id ON public.provider_api_keys USING btree (provider_id, name, id);
--
-- Name: idx_provider_api_keys_provider_active_priority_id; Type: INDEX; Schema: public; Owner: -
--
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 = 20260524000000;
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260527000000;
const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#"
SELECT COUNT(*)::BIGINT
crates/aether-data/src/lifecycle/migrate/tests.rs
+5
View File
@@ -314,6 +314,7 @@ fn empty_database_snapshot_covers_current_cutoff_versions() {
20260520010000,
20260522000000,
20260524000000,
20260527000000,
]
);
}
@@ -390,6 +391,7 @@ 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_name_id"));
assert!(
EMPTY_DATABASE_SNAPSHOT_SQL.contains("idx_provider_api_keys_provider_active_priority_id")
);
@@ -668,6 +670,7 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
20260520000000,
20260520010000,
20260524000000,
20260527000000,
]
);
assert_eq!(
@@ -692,6 +695,7 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
20260520000000,
20260520010000,
20260524000000,
20260527000000,
]
);
}
@@ -1217,6 +1221,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
20260520010000,
20260522000000,
20260524000000,
20260527000000,
]
);
}
crates/aether-data/src/repository/pool_scores/postgres.rs
+7 -8
View File
@@ -341,14 +341,13 @@ impl PoolScoreReadRepository for PostgresPoolMemberScoreRepository {
if query.ids.is_empty() {
return Ok(Vec::new());
}
let mut builder = QueryBuilder::<Postgres>::new(SCORE_COLUMNS);
let mut where_clause = WhereClause::new();
push_in(&mut builder, &mut where_clause, "id", &query.ids);
let rows = builder
.build()
.fetch_all(&self.pool)
.await
.map_postgres_err()?;
let rows = sqlx::query(&format!(
"{SCORE_COLUMNS} WHERE id = ANY($1) ORDER BY id ASC"
))
.bind(&query.ids)
.fetch_all(&self.pool)
.await
.map_postgres_err()?;
rows.iter().map(map_score_row).collect()
}
}
crates/aether-data/src/repository/provider_catalog/mysql.rs
+75 -40
View File
@@ -1,5 +1,5 @@
use async_trait::async_trait;
use sqlx::{mysql::MySqlRow, Row};
use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row};
use super::{
InMemoryProviderCatalogReadRepository, ProviderCatalogKeyListQuery,
@@ -16,6 +16,33 @@ pub struct MysqlProviderCatalogReadRepository {
pool: MysqlPool,
}
const KEY_SELECT_SQL: &str = r#"
SELECT
id, provider_id, name, auth_type, capabilities, is_active, api_formats,
auth_type_by_format, allow_auth_channel_mismatch_formats,
COALESCE(api_key, encrypted_key) AS api_key,
auth_config, note, internal_priority, rate_multipliers,
global_priority_by_format, allowed_models,
expires_at 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, last_429_at AS last_429_at_unix_secs, last_429_type,
adjustment_history, utilization_samples,
last_probe_increase_at AS last_probe_increase_at_unix_secs,
last_rpm_peak, request_count, total_tokens, total_cost_usd,
success_count, error_count, total_response_time_ms,
last_used_at AS last_used_at_unix_secs, auto_fetch_models,
last_models_fetch_at AS last_models_fetch_at_unix_secs,
last_models_fetch_error, locked_models, model_include_patterns,
model_exclude_patterns, upstream_metadata,
oauth_invalid_at AS oauth_invalid_at_unix_secs,
oauth_invalid_reason, status_snapshot,
created_at AS created_at_unix_ms,
updated_at AS updated_at_unix_secs,
health_by_format, circuit_breaker_by_format
FROM provider_api_keys
"#;
impl MysqlProviderCatalogReadRepository {
pub fn new(pool: MysqlPool) -> Self {
Self { pool }
@@ -71,37 +98,26 @@ WHERE api_format IS NOT NULL
}
async fn load_keys(&self) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
let rows = sqlx::query(
r#"
SELECT
id, provider_id, name, auth_type, capabilities, is_active, api_formats,
auth_type_by_format, allow_auth_channel_mismatch_formats,
COALESCE(api_key, encrypted_key) AS api_key,
auth_config, note, internal_priority, rate_multipliers,
global_priority_by_format, allowed_models,
expires_at 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, last_429_at AS last_429_at_unix_secs, last_429_type,
adjustment_history, utilization_samples,
last_probe_increase_at AS last_probe_increase_at_unix_secs,
last_rpm_peak, request_count, total_tokens, total_cost_usd,
success_count, error_count, total_response_time_ms,
last_used_at AS last_used_at_unix_secs, auto_fetch_models,
last_models_fetch_at AS last_models_fetch_at_unix_secs,
last_models_fetch_error, locked_models, model_include_patterns,
model_exclude_patterns, upstream_metadata,
oauth_invalid_at AS oauth_invalid_at_unix_secs,
oauth_invalid_reason, status_snapshot,
created_at AS created_at_unix_ms,
updated_at AS updated_at_unix_secs,
health_by_format, circuit_breaker_by_format
FROM provider_api_keys
"#,
)
.fetch_all(&self.pool)
.await
.map_sql_err()?;
let rows = sqlx::query(KEY_SELECT_SQL)
.fetch_all(&self.pool)
.await
.map_sql_err()?;
rows.iter().map(map_key_row).collect()
}
async fn list_keys_by_provider_ids_direct(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
if provider_ids.is_empty() {
return Ok(Vec::new());
}
let rows = build_list_keys_by_provider_ids_query(provider_ids)
.build()
.fetch_all(&self.pool)
.await
.map_sql_err()?;
rows.iter().map(map_key_row).collect()
}
@@ -912,20 +928,14 @@ impl ProviderCatalogReadRepository for MysqlProviderCatalogReadRepository {
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
self.load_memory()
.await?
.list_keys_by_provider_ids(provider_ids)
.await
self.list_keys_by_provider_ids_direct(provider_ids).await
}
async fn list_key_summaries_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
self.load_memory()
.await?
.list_key_summaries_by_provider_ids(provider_ids)
.await
self.list_keys_by_provider_ids_direct(provider_ids).await
}
async fn list_key_maintenance_summaries_by_provider_ids(
@@ -1162,6 +1172,19 @@ fn optional_json_to_string(
optional_json_ref_to_string(value.as_ref(), field_name)
}
fn build_list_keys_by_provider_ids_query(provider_ids: &[String]) -> QueryBuilder<'_, MySql> {
let mut builder = QueryBuilder::<MySql>::new(KEY_SELECT_SQL);
builder.push("WHERE provider_id IN (");
{
let mut separated = builder.separated(", ");
for provider_id in provider_ids {
separated.push_bind(provider_id.clone());
}
}
builder.push(") ORDER BY provider_id ASC, name ASC, id ASC");
builder
}
fn key_insert_sql() -> &'static str {
r#"
INSERT INTO provider_api_keys (
@@ -1568,6 +1591,7 @@ mod tests {
StoredProviderCatalogProvider,
};
use serde_json::json;
use sqlx::Execute;
#[tokio::test]
async fn repository_builds_from_lazy_pool() {
@@ -1580,6 +1604,17 @@ mod tests {
let _repository = MysqlProviderCatalogReadRepository::new(pool);
}
#[test]
fn list_keys_by_provider_ids_query_targets_index_aligned_ordering() {
let provider_ids = vec!["provider-a".to_string(), "provider-b".to_string()];
let mut builder = super::build_list_keys_by_provider_ids_query(&provider_ids);
let query = builder.build();
let sql = query.sql();
assert!(sql.contains("WHERE provider_id IN ("));
assert!(sql.contains("ORDER BY provider_id ASC, name ASC, id ASC"));
}
#[tokio::test]
async fn mysql_provider_catalog_repository_round_trips_when_url_is_set() {
let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL")