mirror of
https://github.com/clearml/clearml-server
synced 2025-01-31 10:56:48 +00:00
139 lines
4.6 KiB
Python
139 lines
4.6 KiB
Python
import importlib.util
|
|
from datetime import datetime
|
|
from inspect import signature
|
|
from logging import Logger
|
|
from pathlib import Path
|
|
|
|
import pymongo.database
|
|
from mongoengine.connection import get_db
|
|
from packaging.version import Version, parse
|
|
|
|
from apiserver.config_repo import config
|
|
from apiserver.database import utils
|
|
from apiserver.database import Database
|
|
from apiserver.database.model.version import Version as DatabaseVersion
|
|
from apiserver.utilities.dicts import nested_get
|
|
|
|
_migrations = "migrations"
|
|
_parent_dir = Path(__file__).resolve().parents[1]
|
|
_migration_dir = _parent_dir / _migrations
|
|
log = config.logger(__file__)
|
|
|
|
|
|
def check_mongo_empty() -> bool:
|
|
for alias in utils.get_options(Database):
|
|
collection_names = get_db(alias).list_collection_names()
|
|
if collection_names and any(
|
|
name in collection_names
|
|
for name in ["company", "user", "versions"]
|
|
):
|
|
return False
|
|
|
|
return True
|
|
|
|
|
|
def get_last_server_version() -> Version:
|
|
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}")
|
|
|
|
return previous_versions[0] if previous_versions else Version("0.0.0")
|
|
|
|
|
|
def _ensure_mongodb_version():
|
|
db: pymongo.database.Database = get_db(Database.backend)
|
|
db_version = db.client.server_info()["version"]
|
|
if not db_version.startswith("5.0"):
|
|
log.warning(f"Database version should be 5.0.x. Instead: {str(db_version)}")
|
|
return
|
|
|
|
res = db.client.admin.command({"getParameter": 1, "featureCompatibilityVersion": 1})
|
|
version = nested_get(res, ("featureCompatibilityVersion", "version"))
|
|
if version == "5.0":
|
|
return
|
|
if version != "4.4":
|
|
log.warning(f"Cannot upgrade DB version. Should be 4.4. {str(res)}")
|
|
return
|
|
|
|
log.info("Upgrading db version from 4.4 to 5.0")
|
|
res = db.client.admin.command({"setFeatureCompatibilityVersion": "5.0"})
|
|
log.info(res)
|
|
|
|
|
|
def _apply_migrations(log: Logger):
|
|
"""
|
|
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")
|
|
|
|
_ensure_mongodb_version()
|
|
|
|
if not _migration_dir.is_dir():
|
|
raise ValueError(f"Invalid migration dir {_migration_dir}")
|
|
|
|
empty_dbs = check_mongo_empty()
|
|
last_version = get_last_server_version()
|
|
|
|
try:
|
|
new_scripts = {
|
|
ver: path
|
|
for ver, path in (
|
|
(parse(f.stem.replace("_", ".")), 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(
|
|
".".join(("apiserver", _parent_dir.name, _migrations, 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:
|
|
sig = signature(func)
|
|
kwargs = {}
|
|
if len(sig.parameters) == 2:
|
|
name, param = list(sig.parameters.items())[-1]
|
|
key = name.replace("_", "-")
|
|
if issubclass(param.annotation, pymongo.database.Database) and key in dbs:
|
|
kwargs[name] = get_db(key)
|
|
log.info(f"Applying {script.stem}/{func_name}()")
|
|
func(get_db(alias), **kwargs)
|
|
except Exception:
|
|
log.exception(f"Failed applying {script}:{func_name}()")
|
|
raise ValueError(
|
|
"Migration failed, aborting. Please restore backup."
|
|
)
|
|
|
|
DatabaseVersion(
|
|
id=utils.id(),
|
|
num=str(script_version),
|
|
created=datetime.utcnow(),
|
|
desc="Applied on server startup",
|
|
).save()
|
|
|
|
log.info("Finished mongodb migrations")
|