wehub-resource-sync
79 lines
2.9 KiB
Python
79 lines
2.9 KiB
Python
from pathlib import Path
|
|
|
|
from alembic.command import stamp, upgrade
|
|
from alembic.config import Config
|
|
from alembic.migration import MigrationContext
|
|
from alembic.script import ScriptDirectory
|
|
from sqlalchemy import inspect
|
|
from sqlalchemy.engine.base import Engine
|
|
|
|
INITIAL_REVISION = "8606fa83a998"
|
|
|
|
|
|
def _get_alembic_dir() -> str:
|
|
return Path(__file__).parent / "migrations"
|
|
|
|
|
|
def _get_alembic_config(url: str) -> Config:
|
|
alembic_dir = _get_alembic_dir()
|
|
alembic_ini_path = alembic_dir / "alembic.ini"
|
|
alembic_cfg = Config(alembic_ini_path)
|
|
alembic_cfg.set_main_option("script_location", str(alembic_dir))
|
|
url = url.replace("%", "%%") # Same as here: https://github.com/mlflow/mlflow/issues/1487
|
|
alembic_cfg.set_main_option("sqlalchemy.url", url)
|
|
return alembic_cfg
|
|
|
|
|
|
def _is_legacy_database(engine: Engine) -> bool:
|
|
"""Check if this is a pre-3.6.0 auth database that needs version table migration.
|
|
|
|
Returns True if:
|
|
- Auth tables (users, experiment_permissions, etc.) exist
|
|
- alembic_version_auth table does NOT exist
|
|
|
|
This indicates a database from MLflow < 3.6.0 that was using the old
|
|
initialization method without Alembic version tracking.
|
|
"""
|
|
inspector = inspect(engine)
|
|
existing_tables = inspector.get_table_names()
|
|
|
|
has_auth_tables = "users" in existing_tables
|
|
has_version_table = "alembic_version_auth" in existing_tables
|
|
|
|
return has_auth_tables and not has_version_table
|
|
|
|
|
|
def _stamp_legacy_database(engine: Engine, revision: str) -> None:
|
|
"""Stamp a legacy database with the initial Alembic revision.
|
|
|
|
This creates the alembic_version_auth table and marks the database
|
|
as being at the specified revision, avoiding the need to re-run
|
|
the initial migration on databases that already have the tables.
|
|
"""
|
|
alembic_cfg = _get_alembic_config(engine.url.render_as_string(hide_password=False))
|
|
with engine.begin() as conn:
|
|
alembic_cfg.attributes["connection"] = conn
|
|
stamp(alembic_cfg, revision)
|
|
|
|
|
|
def migrate(engine: Engine, revision: str) -> None:
|
|
if _is_legacy_database(engine):
|
|
_stamp_legacy_database(engine, INITIAL_REVISION)
|
|
|
|
alembic_cfg = _get_alembic_config(engine.url.render_as_string(hide_password=False))
|
|
with engine.begin() as conn:
|
|
alembic_cfg.attributes["connection"] = conn
|
|
upgrade(alembic_cfg, revision)
|
|
|
|
|
|
def migrate_if_needed(engine: Engine, revision: str) -> None:
|
|
if _is_legacy_database(engine):
|
|
_stamp_legacy_database(engine, INITIAL_REVISION)
|
|
|
|
alembic_cfg = _get_alembic_config(engine.url.render_as_string(hide_password=False))
|
|
script_dir = ScriptDirectory.from_config(alembic_cfg)
|
|
with engine.begin() as conn:
|
|
context = MigrationContext.configure(conn, opts={"version_table": "alembic_version_auth"})
|
|
if context.get_current_revision() != script_dir.get_current_head():
|
|
upgrade(alembic_cfg, revision)
|