use std::collections::{HashMap, HashSet}; use std::future::Future; use std::pin::Pin; use sqlx::{ migrate::{AppliedMigration, Migrate, MigrateError, Migrator}, query, query_scalar, PgConnection, PgPool, }; use tracing::{error, info, warn}; use aether_data_contracts::PendingMigrationInfo; pub static POSTGRES_MIGRATOR: Migrator = sqlx::migrate!("./migrations"); const MIGRATIONS_TABLE_EXISTS_SQL: &str = "SELECT to_regclass('public._sqlx_migrations') IS NOT NULL"; const USAGE_LEGACY_BODY_REF_CLEANUP_INDEX_MIGRATION_VERSION: i64 = 20260715000000; const USAGE_SETTLEMENT_DASHBOARD_INDEX_MIGRATION_VERSION: i64 = 20260715130000; const USAGE_STALE_PENDING_CLEANUP_INDEX_MIGRATION_VERSION: i64 = 20260720000000; const INVALID_USAGE_LEGACY_BODY_REF_CLEANUP_INDEX_EXISTS_SQL: &str = r#" SELECT EXISTS ( SELECT 1 FROM pg_catalog.pg_class AS index_relation JOIN pg_catalog.pg_namespace AS index_namespace ON index_namespace.oid = index_relation.relnamespace JOIN pg_catalog.pg_index AS index_state ON index_state.indexrelid = index_relation.oid WHERE index_namespace.nspname = 'public' AND index_relation.relname = 'idx_usage_legacy_body_ref_cleanup_created_at' AND NOT index_state.indisvalid ) "#; const DROP_USAGE_LEGACY_BODY_REF_CLEANUP_INDEX_SQL: &str = "DROP INDEX CONCURRENTLY IF EXISTS public.idx_usage_legacy_body_ref_cleanup_created_at"; const INVALID_USAGE_SETTLEMENT_DASHBOARD_INDEX_EXISTS_SQL: &str = r#" SELECT EXISTS ( SELECT 1 FROM pg_catalog.pg_class AS index_relation JOIN pg_catalog.pg_namespace AS index_namespace ON index_namespace.oid = index_relation.relnamespace JOIN pg_catalog.pg_index AS index_state ON index_state.indexrelid = index_relation.oid WHERE index_namespace.nspname = 'public' AND index_relation.relname = 'idx_usage_settlement_dashboard_cover' AND NOT index_state.indisvalid ) "#; const DROP_USAGE_SETTLEMENT_DASHBOARD_INDEX_SQL: &str = "DROP INDEX CONCURRENTLY IF EXISTS public.idx_usage_settlement_dashboard_cover"; const INVALID_USAGE_STALE_PENDING_CLEANUP_INDEX_EXISTS_SQL: &str = r#" SELECT EXISTS ( SELECT 1 FROM pg_catalog.pg_class AS index_relation JOIN pg_catalog.pg_namespace AS index_namespace ON index_namespace.oid = index_relation.relnamespace JOIN pg_catalog.pg_index AS index_state ON index_state.indexrelid = index_relation.oid WHERE index_namespace.nspname = 'public' AND index_relation.relname = 'idx_usage_stale_pending_created_request' AND NOT index_state.indisvalid ) "#; const DROP_USAGE_STALE_PENDING_CLEANUP_INDEX_SQL: &str = "DROP INDEX CONCURRENTLY IF EXISTS public.idx_usage_stale_pending_created_request"; pub type BootstrapFuture<'a> = Pin> + 'a>>; pub trait PostgresMigrationBootstrap: Send + Sync { fn apply_snapshot<'a>( &self, conn: &'a mut PgConnection, migrator: &'static Migrator, ) -> BootstrapFuture<'a>; } #[derive(Debug, Clone, Copy, Default)] struct NoopBootstrap; impl PostgresMigrationBootstrap for NoopBootstrap { fn apply_snapshot<'a>( &self, _conn: &'a mut PgConnection, _migrator: &'static Migrator, ) -> BootstrapFuture<'a> { Box::pin(async { Ok(()) }) } } pub async fn run_migrations(pool: &PgPool) -> Result<(), MigrateError> { run_migrations_with_bootstrap(pool, &NoopBootstrap).await } pub async fn run_migrations_with_bootstrap( pool: &PgPool, bootstrap: &dyn PostgresMigrationBootstrap, ) -> Result<(), MigrateError> { let mut conn = crate::pool::acquire_postgres_migration_connection(pool).await?; if POSTGRES_MIGRATOR.locking { conn.lock().await?; } let result = run_migrations_locked(&mut conn, bootstrap).await; if POSTGRES_MIGRATOR.locking { match conn.unlock().await { Ok(()) => {} Err(unlock_error) if result.is_ok() => return Err(unlock_error), Err(unlock_error) => { warn!( error = %unlock_error, "database migration lock release failed after migration error" ); } } } result } pub async fn pending_migrations(pool: &PgPool) -> Result, MigrateError> { let mut conn = pool.acquire().await?; pending_migrations_locked(&mut conn).await } pub async fn prepare_database_for_startup( pool: &PgPool, ) -> Result, MigrateError> { prepare_database_for_startup_with_bootstrap(pool, &NoopBootstrap).await } pub async fn prepare_database_for_startup_with_bootstrap( pool: &PgPool, bootstrap: &dyn PostgresMigrationBootstrap, ) -> Result, MigrateError> { let mut conn = crate::pool::acquire_postgres_migration_connection(pool).await?; if POSTGRES_MIGRATOR.locking { conn.lock().await?; } let result = prepare_database_for_startup_locked(&mut conn, bootstrap).await; if POSTGRES_MIGRATOR.locking { match conn.unlock().await { Ok(()) => {} Err(unlock_error) if result.is_ok() => return Err(unlock_error), Err(unlock_error) => { warn!( error = %unlock_error, "database migration lock release failed after startup preparation error" ); } } } result } async fn run_migrations_locked( conn: &mut PgConnection, bootstrap: &dyn PostgresMigrationBootstrap, ) -> Result<(), MigrateError> { conn.ensure_migrations_table().await?; bootstrap.apply_snapshot(conn, &POSTGRES_MIGRATOR).await?; if let Some(version) = conn.dirty_version().await? { error!(version, "database migration state is dirty"); return Err(MigrateError::Dirty(version)); } let applied_migrations = conn.list_applied_migrations().await?; validate_applied_migrations(&applied_migrations)?; let applied_migrations_by_version = applied_migrations .into_iter() .map(|migration| (migration.version, migration)) .collect::>(); let pending_migrations = POSTGRES_MIGRATOR .iter() .filter(|migration| migration.migration_type.is_up_migration()) .filter(|migration| !applied_migrations_by_version.contains_key(&migration.version)) .collect::>(); let total_migrations = POSTGRES_MIGRATOR .iter() .filter(|migration| migration.migration_type.is_up_migration()) .count(); let applied_count = total_migrations.saturating_sub(pending_migrations.len()); if pending_migrations.is_empty() { info!( total_migrations, applied_migrations = applied_count, pending_migrations = 0, "database migrations already up to date" ); return Ok(()); } info!( total_migrations, applied_migrations = applied_count, pending_migrations = pending_migrations.len(), "database migrations pending" ); for (index, migration) in pending_migrations.iter().enumerate() { let current = index + 1; info!( current, total = pending_migrations.len(), version = migration.version, description = %migration.description, "applying database migration" ); repair_invalid_concurrent_index(conn, migration.version).await?; let elapsed = conn.apply(migration).await?; info!( current, total = pending_migrations.len(), version = migration.version, description = %migration.description, elapsed_ms = elapsed.as_millis() as u64, "applied database migration" ); } info!( total_migrations, applied_migrations = total_migrations, pending_migrations = 0, "database migrations complete" ); Ok(()) } async fn repair_invalid_concurrent_index( conn: &mut PgConnection, migration_version: i64, ) -> Result<(), MigrateError> { let (index_name, invalid_index_exists_sql, drop_index_sql) = match migration_version { USAGE_LEGACY_BODY_REF_CLEANUP_INDEX_MIGRATION_VERSION => ( "idx_usage_legacy_body_ref_cleanup_created_at", INVALID_USAGE_LEGACY_BODY_REF_CLEANUP_INDEX_EXISTS_SQL, DROP_USAGE_LEGACY_BODY_REF_CLEANUP_INDEX_SQL, ), USAGE_SETTLEMENT_DASHBOARD_INDEX_MIGRATION_VERSION => ( "idx_usage_settlement_dashboard_cover", INVALID_USAGE_SETTLEMENT_DASHBOARD_INDEX_EXISTS_SQL, DROP_USAGE_SETTLEMENT_DASHBOARD_INDEX_SQL, ), USAGE_STALE_PENDING_CLEANUP_INDEX_MIGRATION_VERSION => ( "idx_usage_stale_pending_created_request", INVALID_USAGE_STALE_PENDING_CLEANUP_INDEX_EXISTS_SQL, DROP_USAGE_STALE_PENDING_CLEANUP_INDEX_SQL, ), _ => return Ok(()), }; let invalid_index_exists: bool = query_scalar(invalid_index_exists_sql) .fetch_one(&mut *conn) .await?; if !invalid_index_exists { return Ok(()); } warn!( migration_version, index = index_name, "dropping invalid index left by an interrupted concurrent migration" ); query(drop_index_sql).execute(&mut *conn).await?; Ok(()) } async fn prepare_database_for_startup_locked( conn: &mut PgConnection, bootstrap: &dyn PostgresMigrationBootstrap, ) -> Result, MigrateError> { conn.ensure_migrations_table().await?; bootstrap.apply_snapshot(conn, &POSTGRES_MIGRATOR).await?; pending_migrations_locked(conn).await } async fn pending_migrations_locked( conn: &mut PgConnection, ) -> Result, MigrateError> { if !migrations_table_exists(conn).await? { return Ok(all_up_migrations()); } if let Some(version) = conn.dirty_version().await? { error!(version, "database migration state is dirty"); return Err(MigrateError::Dirty(version)); } let applied_migrations = conn.list_applied_migrations().await?; validate_applied_migrations(&applied_migrations)?; Ok(pending_migrations_from_applied(&applied_migrations)) } async fn migrations_table_exists(conn: &mut PgConnection) -> Result { Ok(query_scalar(MIGRATIONS_TABLE_EXISTS_SQL) .fetch_one(&mut *conn) .await?) } pub fn all_up_migrations() -> Vec { POSTGRES_MIGRATOR .iter() .filter(|migration| migration.migration_type.is_up_migration()) .map(|migration| PendingMigrationInfo { version: migration.version, description: migration.description.to_string(), }) .collect() } pub fn pending_migrations_from_applied( applied_migrations: &[AppliedMigration], ) -> Vec { let applied_versions = applied_migrations .iter() .map(|migration| migration.version) .collect::>(); POSTGRES_MIGRATOR .iter() .filter(|migration| migration.migration_type.is_up_migration()) .filter(|migration| !applied_versions.contains(&migration.version)) .map(|migration| PendingMigrationInfo { version: migration.version, description: migration.description.to_string(), }) .collect() } fn validate_applied_migrations( applied_migrations: &[AppliedMigration], ) -> Result<(), MigrateError> { if POSTGRES_MIGRATOR.ignore_missing { return Ok(()); } let known_versions = POSTGRES_MIGRATOR .iter() .map(|migration| migration.version) .collect::>(); for applied_migration in applied_migrations { if !known_versions.contains(&applied_migration.version) { error!( version = applied_migration.version, "applied database migration is missing from embedded migrations" ); return Err(MigrateError::VersionMissing(applied_migration.version)); } } for migration in POSTGRES_MIGRATOR .iter() .filter(|migration| migration.migration_type.is_up_migration()) { if let Some(applied_migration) = applied_migrations .iter() .find(|applied_migration| applied_migration.version == migration.version) { if migration.checksum != applied_migration.checksum { warn!( version = migration.version, description = %migration.description, "database migration checksum mismatch (ignored: version-only validation)" ); } } } Ok(()) } #[cfg(test)] mod tests { use super::{all_up_migrations, pending_migrations_from_applied, POSTGRES_MIGRATOR}; #[test] fn embeds_ordered_postgres_migration_sources() { let versions = POSTGRES_MIGRATOR .iter() .map(|migration| migration.version) .collect::>(); assert!(!versions.is_empty()); assert!(versions.windows(2).all(|pair| pair[0] < pair[1])); assert_eq!(pending_migrations_from_applied(&[]), all_up_migrations()); } #[test] fn embeds_scoped_codex_live_permission_migration() { let migration = POSTGRES_MIGRATOR .iter() .find(|migration| migration.version == 20260821000000) .expect("Codex Live permission migration should be embedded"); let sql = migration.sql.as_ref(); for required_fragment in [ "UPDATE public.users", "UPDATE public.user_groups", "UPDATE public.api_keys", "UPDATE public.provider_api_keys", "provider.provider_type", "openai:responses", "codex:live", ] { assert!( sql.contains(required_fragment), "Codex Live permission migration is missing {required_fragment}" ); } } #[test] fn concurrent_index_migrations_opt_out_of_transactions() { for version in [20260715000000, 20260715130000, 20260720000000] { let migration = POSTGRES_MIGRATOR .iter() .find(|migration| migration.version == version) .expect("concurrent index migration should be embedded"); assert!( migration.no_tx, "migration {version} must run without a transaction" ); } let analyze_tuning = POSTGRES_MIGRATOR .iter() .find(|migration| migration.version == 20260715130100) .expect("analyze tuning migration should be embedded"); assert!(!analyze_tuning.no_tx); } #[test] fn tunnel_generation_backfill_does_not_require_pgcrypto() { let migration = POSTGRES_MIGRATOR .iter() .find(|migration| migration.version == 20260831010000) .expect("tunnel generation migration should be embedded"); let sql = migration.sql.as_ref(); // Existing installations may not have the optional pgcrypto extension. // Legacy rows only need an opaque epoch marker; new registrations use a // CSPRNG in the application layer. assert!(sql.contains("SET tunnel_generation = md5(")); assert!(sql.contains("ctid::text")); assert!(!sql.contains("gen_random_uuid")); } #[test] fn gemini_file_mapping_metadata_migration_expands_legacy_columns() { let migration = POSTGRES_MIGRATOR .iter() .find(|migration| migration.version == 20260901000000) .expect("Gemini file mapping metadata migration should be embedded"); let sql = migration.sql.as_ref(); for expected in [ "file_name TYPE character varying(512)", "display_name TYPE character varying(512)", "mime_type TYPE character varying(255)", ] { assert!(sql.contains(expected), "migration is missing {expected}"); } } }