from __future__ import annotations import os from logging.config import fileConfig from alembic import context from sqlalchemy import engine_from_config, pool, text from iaai_sync_service.models import Base config = context.config if config.config_file_name is not None: fileConfig(config.config_file_name) target_metadata = Base.metadata SYNC_VERSION_TABLE = "iaai_sync_alembic_version" LEGACY_VERSION_TABLE = "alembic_version" SYNC_REVISIONS = { "0001_ensure_sync_schema", "0002_extend_origin_enum_for_iaai", "0003_create_iaai_source_tables", } def _normalize_database_url(url: str) -> str: normalized = url.strip() if normalized.startswith("jdbc:"): normalized = normalized[len("jdbc:") :] if normalized.startswith("postgres://"): normalized = "postgresql+psycopg://" + normalized[len("postgres://") :] elif normalized.startswith("postgresql://"): normalized = "postgresql+psycopg://" + normalized[len("postgresql://") :] return normalized def _database_url() -> str: raw_url = os.getenv("DATABASE_URL", config.get_main_option("sqlalchemy.url")) return _normalize_database_url(raw_url) def _select_version_table(connection) -> str: # noqa: ANN001 inspector = connection.dialect if inspector.has_table(connection, SYNC_VERSION_TABLE): return SYNC_VERSION_TABLE if inspector.has_table(connection, LEGACY_VERSION_TABLE): version_num = connection.execute(text(f"SELECT version_num FROM {LEGACY_VERSION_TABLE} LIMIT 1")).scalar() if version_num in SYNC_REVISIONS: return LEGACY_VERSION_TABLE return SYNC_VERSION_TABLE def run_migrations_offline() -> None: url = _database_url() context.configure( url=url, target_metadata=target_metadata, version_table=SYNC_VERSION_TABLE, literal_binds=True, dialect_opts={"paramstyle": "named"}, ) with context.begin_transaction(): context.run_migrations() def run_migrations_online() -> None: configuration = config.get_section(config.config_ini_section) or {} configuration["sqlalchemy.url"] = _database_url() connectable = engine_from_config(configuration, prefix="sqlalchemy.", poolclass=pool.NullPool) with connectable.connect() as connection: version_table = _select_version_table(connection) context.configure( connection=connection, target_metadata=target_metadata, version_table=version_table, ) with context.begin_transaction(): context.run_migrations() if context.is_offline_mode(): run_migrations_offline() else: run_migrations_online()