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

Fix provider deletion cleanup

Browse Source
This commit is contained in:
elky
2026-06-12 16:25:11 +08:00
parent 30b545785f
commit 308cc88ef7
17 changed files with 385 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-data-contracts/src/repository/provider_catalog/types.rs
+1
View File
@@ -651,6 +651,7 @@ pub trait ProviderCatalogWriteRepository: Send + Sync {
async fn cleanup_deleted_provider_refs(
&self,
provider_id: &str,
provider_deleted: bool,
endpoint_ids: &[String],
key_ids: &[String],
) -> Result<(), crate::DataLayerError>;
crates/aether-data/src/repository/provider_catalog/memory.rs
+1
View File
@@ -652,6 +652,7 @@ impl ProviderCatalogWriteRepository for InMemoryProviderCatalogReadRepository {
async fn cleanup_deleted_provider_refs(
&self,
_provider_id: &str,
_provider_deleted: bool,
_endpoint_ids: &[String],
_key_ids: &[String],
) -> Result<(), DataLayerError> {
crates/aether-data/src/repository/provider_catalog/mysql.rs
+61 -3
View File
@@ -310,10 +310,60 @@ WHERE id = ?
pub async fn cleanup_deleted_provider_refs(
&self,
provider_id: &str,
_endpoint_ids: &[String],
_key_ids: &[String],
provider_deleted: bool,
endpoint_ids: &[String],
key_ids: &[String],
) -> Result<(), DataLayerError> {
validate_non_empty(provider_id, "provider catalog provider_id")?;
let mut tx = self.pool.begin().await.map_sql_err()?;
if provider_deleted {
sqlx::query(
"UPDATE user_preferences SET default_provider_id = NULL WHERE default_provider_id = ?",
)
.bind(provider_id)
.execute(&mut *tx)
.await
.map_sql_err()?;
sqlx::query("UPDATE video_tasks SET provider_id = NULL WHERE provider_id = ?")
.bind(provider_id)
.execute(&mut *tx)
.await
.map_sql_err()?;
sqlx::query("DELETE FROM request_candidates WHERE provider_id = ?")
.bind(provider_id)
.execute(&mut *tx)
.await
.map_sql_err()?;
}
for endpoint_id in endpoint_ids {
sqlx::query("UPDATE video_tasks SET endpoint_id = NULL WHERE endpoint_id = ?")
.bind(endpoint_id)
.execute(&mut *tx)
.await
.map_sql_err()?;
sqlx::query("DELETE FROM request_candidates WHERE endpoint_id = ?")
.bind(endpoint_id)
.execute(&mut *tx)
.await
.map_sql_err()?;
}
for key_id in key_ids {
sqlx::query("DELETE FROM gemini_file_mappings WHERE key_id = ?")
.bind(key_id)
.execute(&mut *tx)
.await
.map_sql_err()?;
sqlx::query("UPDATE video_tasks SET key_id = NULL WHERE key_id = ?")
.bind(key_id)
.execute(&mut *tx)
.await
.map_sql_err()?;
}
tx.commit().await.map_sql_err()?;
Ok(())
}
@@ -990,10 +1040,18 @@ impl ProviderCatalogWriteRepository for MysqlProviderCatalogReadRepository {
async fn cleanup_deleted_provider_refs(
&self,
provider_id: &str,
provider_deleted: bool,
endpoint_ids: &[String],
key_ids: &[String],
) -> Result<(), DataLayerError> {
Self::cleanup_deleted_provider_refs(self, provider_id, endpoint_ids, key_ids).await
Self::cleanup_deleted_provider_refs(
self,
provider_id,
provider_deleted,
endpoint_ids,
key_ids,
)
.await
}
async fn create_endpoint(
crates/aether-data/src/repository/provider_catalog/postgres.rs
+37 -43
View File
@@ -994,6 +994,7 @@ WHERE id = $1
pub async fn cleanup_deleted_provider_refs(
&self,
provider_id: &str,
provider_deleted: bool,
endpoint_ids: &[String],
key_ids: &[String],
) -> Result<(), DataLayerError> {
@@ -1005,37 +1006,27 @@ WHERE id = $1
let mut tx = self.pool.begin().await.map_postgres_err()?;
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()?;
sqlx::query("UPDATE usage SET provider_id = NULL WHERE provider_id = $1")
.bind(provider_id)
.execute(&mut *tx)
.await
.map_postgres_err()?;
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")
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()?;
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 usage SET provider_endpoint_id = NULL WHERE provider_endpoint_id = $1",
)
.bind(endpoint_id)
.execute(&mut *tx)
.await
.map_postgres_err()?;
sqlx::query("UPDATE video_tasks SET endpoint_id = NULL WHERE endpoint_id = $1")
.bind(endpoint_id)
.execute(&mut *tx)
@@ -1054,13 +1045,6 @@ WHERE id = $1
.execute(&mut *tx)
.await
.map_postgres_err()?;
sqlx::query(
"UPDATE usage SET provider_api_key_id = NULL WHERE provider_api_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)
@@ -1068,16 +1052,18 @@ WHERE id = $1
.map_postgres_err()?;
}
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()?;
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(())
@@ -1995,10 +1981,18 @@ impl ProviderCatalogWriteRepository for SqlxProviderCatalogReadRepository {
async fn cleanup_deleted_provider_refs(
&self,
provider_id: &str,
provider_deleted: bool,
endpoint_ids: &[String],
key_ids: &[String],
) -> Result<(), DataLayerError> {
Self::cleanup_deleted_provider_refs(self, provider_id, endpoint_ids, key_ids).await
Self::cleanup_deleted_provider_refs(
self,
provider_id,
provider_deleted,
endpoint_ids,
key_ids,
)
.await
}
async fn create_endpoint(
crates/aether-data/src/repository/provider_catalog/sqlite.rs
+61 -3
View File
@@ -734,10 +734,60 @@ WHERE id = ?
pub async fn cleanup_deleted_provider_refs(
&self,
provider_id: &str,
_endpoint_ids: &[String],
_key_ids: &[String],
provider_deleted: bool,
endpoint_ids: &[String],
key_ids: &[String],
) -> Result<(), DataLayerError> {
validate_non_empty(provider_id, "provider catalog provider_id")?;
let mut tx = self.pool.begin().await.map_sql_err()?;
if provider_deleted {
sqlx::query(
"UPDATE user_preferences SET default_provider_id = NULL WHERE default_provider_id = ?",
)
.bind(provider_id)
.execute(&mut *tx)
.await
.map_sql_err()?;
sqlx::query("UPDATE video_tasks SET provider_id = NULL WHERE provider_id = ?")
.bind(provider_id)
.execute(&mut *tx)
.await
.map_sql_err()?;
sqlx::query("DELETE FROM request_candidates WHERE provider_id = ?")
.bind(provider_id)
.execute(&mut *tx)
.await
.map_sql_err()?;
}
for endpoint_id in endpoint_ids {
sqlx::query("UPDATE video_tasks SET endpoint_id = NULL WHERE endpoint_id = ?")
.bind(endpoint_id)
.execute(&mut *tx)
.await
.map_sql_err()?;
sqlx::query("DELETE FROM request_candidates WHERE endpoint_id = ?")
.bind(endpoint_id)
.execute(&mut *tx)
.await
.map_sql_err()?;
}
for key_id in key_ids {
sqlx::query("DELETE FROM gemini_file_mappings WHERE key_id = ?")
.bind(key_id)
.execute(&mut *tx)
.await
.map_sql_err()?;
sqlx::query("UPDATE video_tasks SET key_id = NULL WHERE key_id = ?")
.bind(key_id)
.execute(&mut *tx)
.await
.map_sql_err()?;
}
tx.commit().await.map_sql_err()?;
Ok(())
}
@@ -1399,10 +1449,18 @@ impl ProviderCatalogWriteRepository for SqliteProviderCatalogReadRepository {
async fn cleanup_deleted_provider_refs(
&self,
provider_id: &str,
provider_deleted: bool,
endpoint_ids: &[String],
key_ids: &[String],
) -> Result<(), DataLayerError> {
Self::cleanup_deleted_provider_refs(self, provider_id, endpoint_ids, key_ids).await
Self::cleanup_deleted_provider_refs(
self,
provider_id,
provider_deleted,
endpoint_ids,
key_ids,
)
.await
}
async fn create_endpoint(
crates/aether-data/src/repository/usage/postgres/tests.rs
+53 -15
View File
@@ -832,6 +832,37 @@ fn usage_sql_writes_usage_settlement_pricing_snapshots() {
assert!(super::UPSERT_USAGE_SETTLEMENT_PRICING_SNAPSHOT_SQL.contains("price_per_request"));
}
#[test]
fn usage_sql_settlement_pricing_snapshot_billing_values_use_authoritative_incoming_values() {
let sql = super::UPSERT_USAGE_SETTLEMENT_PRICING_SNAPSHOT_SQL;
for field in [
"billing_input_tokens",
"billing_effective_input_tokens",
"billing_output_tokens",
"billing_cache_creation_tokens",
"billing_cache_creation_5m_tokens",
"billing_cache_creation_1h_tokens",
"billing_cache_read_tokens",
"billing_total_input_context",
"billing_cache_creation_cost_usd",
"billing_cache_read_cost_usd",
"billing_total_cost_usd",
"billing_actual_total_cost_usd",
] {
let assignment = format!(
"{field} = COALESCE(\n EXCLUDED.{field},\n usage_settlement_snapshots.{field}\n )"
);
assert!(
sql.contains(assignment.as_str()),
"missing authoritative billing snapshot assignment: {assignment}"
);
assert!(
!sql.contains(format!("{field} = GREATEST(").as_str()),
"billing snapshot field should not use max-only conflict resolution: {field}"
);
}
}
#[test]
fn usage_sql_upsert_recovers_missing_provider_links_after_billing_finalizes() {
for assignment in [
@@ -848,25 +879,32 @@ fn usage_sql_upsert_recovers_missing_provider_links_after_billing_finalizes() {
#[test]
fn usage_sql_updates_usage_mirror_columns_from_terminal_events_only() {
for assignment in [
"input_tokens = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".input_tokens, EXCLUDED.input_tokens) ELSE \"usage\".input_tokens END",
"output_tokens = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".output_tokens, EXCLUDED.output_tokens) ELSE \"usage\".output_tokens END",
"total_tokens = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".total_tokens, EXCLUDED.total_tokens) ELSE \"usage\".total_tokens END",
"input_output_total_tokens = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".input_output_total_tokens, EXCLUDED.input_output_total_tokens) ELSE \"usage\".input_output_total_tokens END",
"input_context_tokens = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".input_context_tokens, EXCLUDED.input_context_tokens) ELSE \"usage\".input_context_tokens END",
"cache_creation_input_tokens = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".cache_creation_input_tokens, EXCLUDED.cache_creation_input_tokens) ELSE \"usage\".cache_creation_input_tokens END",
"cache_creation_input_tokens_5m = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".cache_creation_input_tokens_5m, EXCLUDED.cache_creation_input_tokens_5m) ELSE \"usage\".cache_creation_input_tokens_5m END",
"cache_creation_input_tokens_1h = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".cache_creation_input_tokens_1h, EXCLUDED.cache_creation_input_tokens_1h) ELSE \"usage\".cache_creation_input_tokens_1h END",
"cache_read_input_tokens = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".cache_read_input_tokens, EXCLUDED.cache_read_input_tokens) ELSE \"usage\".cache_read_input_tokens END",
"cache_creation_cost_usd = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".cache_creation_cost_usd, EXCLUDED.cache_creation_cost_usd) ELSE \"usage\".cache_creation_cost_usd END",
"cache_read_cost_usd = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".cache_read_cost_usd, EXCLUDED.cache_read_cost_usd) ELSE \"usage\".cache_read_cost_usd END",
"total_cost_usd = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".total_cost_usd, EXCLUDED.total_cost_usd) ELSE \"usage\".total_cost_usd END",
"actual_total_cost_usd = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".actual_total_cost_usd, EXCLUDED.actual_total_cost_usd) ELSE \"usage\".actual_total_cost_usd END",
for field in [
"input_tokens",
"output_tokens",
"total_tokens",
"input_output_total_tokens",
"input_context_tokens",
"cache_creation_input_tokens",
"cache_creation_input_tokens_5m",
"cache_creation_input_tokens_1h",
"cache_read_input_tokens",
"cache_creation_cost_usd",
"cache_read_cost_usd",
"total_cost_usd",
"actual_total_cost_usd",
] {
let assignment = format!(
"{field} = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".{field}, EXCLUDED.{field}) ELSE \"usage\".{field} END"
);
assert!(
super::UPSERT_SQL.contains(assignment),
super::UPSERT_SQL.contains(assignment.as_str()),
"missing terminal mirror assignment: {assignment}"
);
assert!(
!super::UPSERT_SQL.contains(format!("{field} = CASE WHEN EXCLUDED.status IN").as_str()),
"terminal mirror assignment must keep pending billing guard: {field}"
);
}
}