clearml-server/server/mongo/initialize/migration.py

import importlib.util
from datetime import datetime
from logging import Logger
from pathlib import Path

from mongoengine.connection import get_db
from semantic_version import Version

import database.utils
from database import Database
from database.model.version import Version as DatabaseVersion

migration_dir = Path(__file__).resolve().parent.with_name("migrations")


def _apply_migrations(log: Logger) -> bool:
    """
    Apply migrations as found in the migration dir.
    Returns a boolean indicating whether the database was empty prior to migration.
    """
    log = log.getChild(Path(__file__).stem)

    log.info(f"Started mongodb migrations")

    if not migration_dir.is_dir():
        raise ValueError(f"Invalid migration dir {migration_dir}")

    empty_dbs = not any(
        get_db(alias).collection_names()
        for alias in database.utils.get_options(Database)
    )

    try:
        previous_versions = sorted(
            (Version(ver.num) for ver in DatabaseVersion.objects().only("num")),
            reverse=True,
        )
    except ValueError as ex:
        raise ValueError(f"Invalid database version number encountered: {ex}")

    last_version = previous_versions[0] if previous_versions else Version("0.0.0")

    try:
        new_scripts = {
            ver: path
            for ver, path in ((Version(f.stem), f) for f in migration_dir.glob("*.py"))
            if ver > last_version
        }
    except ValueError as ex:
        raise ValueError(f"Failed parsing migration version from file: {ex}")

    dbs = {Database.auth: "migrate_auth", Database.backend: "migrate_backend"}

    for script_version in sorted(new_scripts):
        script = new_scripts[script_version]

        if empty_dbs:
            log.info(f"Skipping migration {script.name} (empty databases)")
        else:
            spec = importlib.util.spec_from_file_location(script.stem, str(script))
            module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(module)

            for alias, func_name in dbs.items():
                func = getattr(module, func_name, None)
                if not func:
                    continue
                try:
                    log.info(f"Applying {script.stem}/{func_name}()")
                    func(get_db(alias))
                except Exception:
                    log.exception(f"Failed applying {script}:{func_name}()")
                    raise ValueError(
                        "Migration failed, aborting. Please restore backup."
                    )

        DatabaseVersion(
            id=database.utils.id(),
            num=script.stem,
            created=datetime.utcnow(),
            desc="Applied on server startup",
        ).save()

    log.info("Finished mongodb migrations")

    return empty_dbs