Aether/crates/aether-data/runtime/src/lifecycle/bootstrap/postgres.rs

use sqlx::{
    migrate::{Migrate, MigrateError, Migration, Migrator},
    query, query_scalar, Connection, PgConnection,
};
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"));
// Keep post-snapshot migrations executable on a fresh database so required
// schema changes still run. Compatibility migrations after this frontier are
// explicit no-ops and must not rewrite legacy rows.
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260821130000;

const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#"
SELECT COUNT(*)::BIGINT
FROM information_schema.tables
WHERE table_schema = 'public'
  AND table_type = 'BASE TABLE'
  AND table_name <> '_sqlx_migrations'
"#;
const AETHER_SCHEMA_FOOTPRINT_TABLE_COUNT_SQL: &str = r#"
SELECT COUNT(*)::BIGINT
FROM information_schema.tables
WHERE table_schema = 'public'
  AND table_type = 'BASE TABLE'
  AND table_name IN (
    'api_key_provider_mappings',
    'auth_modules',
    'gemini_file_mappings',
    'global_models',
    'oauth_providers',
    'provider_api_keys',
    'proxy_nodes',
    'routing_groups',
    'user_groups',
    'usage_routing_snapshots',
    'usage_settlement_snapshots'
  )
"#;
const INSERT_APPLIED_MIGRATION_SQL: &str = r#"
INSERT INTO _sqlx_migrations (
    version,
    description,
    success,
    checksum,
    execution_time
) VALUES (
    $1,
    $2,
    TRUE,
    $3,
    0
)
ON CONFLICT (version) DO NOTHING
"#;

pub(crate) async fn apply_snapshot_if_empty(
    conn: &mut PgConnection,
    migrator: &'static Migrator,
) -> Result<(), MigrateError> {
    if !should_apply_snapshot(conn).await? {
        return Ok(());
    }

    let migrations = snapshot_migrations(migrator)?;
    info!(
        cutoff_version = EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION,
        stamped_migrations = migrations.len(),
        "bootstrapping empty database from empty_database_snapshot"
    );

    let mut tx = conn.begin().await?;
    sqlx::raw_sql(EMPTY_DATABASE_SNAPSHOT_SQL)
        .execute(&mut *tx)
        .await?;
    for migration in migrations {
        query(INSERT_APPLIED_MIGRATION_SQL)
            .bind(migration.version)
            .bind(migration.description.as_ref())
            .bind(migration.checksum.as_ref())
            .execute(&mut *tx)
            .await?;
    }
    tx.commit().await?;

    Ok(())
}

pub(crate) async fn should_apply_snapshot(conn: &mut PgConnection) -> Result<bool, MigrateError> {
    let applied_migrations = conn.list_applied_migrations().await?;
    if !applied_migrations.is_empty() {
        return Ok(false);
    }

    let public_table_count: i64 = query_scalar(PUBLIC_BASE_TABLE_COUNT_SQL)
        .fetch_one(&mut *conn)
        .await?;
    if public_table_count == 0 {
        return Ok(true);
    }

    let aether_footprint_table_count: i64 = query_scalar(AETHER_SCHEMA_FOOTPRINT_TABLE_COUNT_SQL)
        .fetch_one(&mut *conn)
        .await?;
    if aether_footprint_table_count == 0 {
        info!(
            public_table_count,
            "no Aether schema footprint detected; allowing empty database snapshot bootstrap despite pre-existing public tables"
        );
        return Ok(true);
    }

    Ok(false)
}

pub(crate) fn snapshot_migrations(
    migrator: &'static Migrator,
) -> Result<Vec<&'static Migration>, MigrateError> {
    let migrations = migrator
        .iter()
        .filter(|migration| migration.migration_type.is_up_migration())
        .filter(|migration| migration.version <= EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION)
        .collect::<Vec<_>>();

    if migrations.is_empty() {
        return Err(MigrateError::Source(Box::new(std::io::Error::other(
            "empty database snapshot cutoff does not match any embedded migrations",
        ))));
    }

    Ok(migrations)
}