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

from pathlib import Path
from typing import Sequence, Union

from config import config
from config.info import get_default_company
from database.model.auth import Role
from service_repo.auth.fixed_user import FixedUser
from .migration import _apply_migrations
from .pre_populate import PrePopulate
from .user import ensure_fixed_user, _ensure_auth_user, _ensure_backend_user
from .util import _ensure_company, _ensure_default_queue, _ensure_uuid

log = config.logger(__package__)


def _pre_populate(company_id: str, zip_file: str):
    if not zip_file or not Path(zip_file).is_file():
        msg = f"Invalid pre-populate zip file: {zip_file}"
        if config.get("apiserver.pre_populate.fail_on_error", False):
            log.error(msg)
            raise ValueError(msg)
        else:
            log.warning(msg)
    else:
        log.info(f"Pre-populating using {zip_file}")

        user_id = _ensure_backend_user(
            "__allegroai__", company_id, "Allegro.ai"
        )

        PrePopulate.import_from_zip(
            zip_file,
            company_id="",
            user_id=user_id,
            artifacts_path=config.get(
                "apiserver.pre_populate.artifacts_path", None
            ),
        )


def _resolve_zip_files(zip_files: Union[Sequence[str], str]) -> Sequence[str]:
    if isinstance(zip_files, str):
        zip_files = [zip_files]
    for p in map(Path, zip_files):
        if p.is_file():
            yield p
        if p.is_dir():
            yield from p.glob("*.zip")
        log.warning(f"Invalid pre-populate entry {str(p)}, skipping")


def pre_populate_data():
    for zip_file in _resolve_zip_files(config.get("apiserver.pre_populate.zip_files")):
        _pre_populate(company_id=get_default_company(), zip_file=zip_file)


def init_mongo_data() -> bool:
    try:
        empty_dbs = _apply_migrations(log)

        _ensure_uuid()

        company_id = _ensure_company(log)

        _ensure_default_queue(company_id)

        fixed_mode = FixedUser.enabled()

        for user, credentials in config.get("secure.credentials", {}).items():
            user_data = {
                "name": user,
                "role": credentials.role,
                "email": f"{user}@example.com",
                "key": credentials.user_key,
                "secret": credentials.user_secret,
            }
            revoke = fixed_mode and credentials.get("revoke_in_fixed_mode", False)
            user_id = _ensure_auth_user(user_data, company_id, log=log, revoke=revoke)
            if credentials.role == Role.user:
                _ensure_backend_user(user_id, company_id, credentials.display_name)

        if fixed_mode:
            log.info("Fixed users mode is enabled")
            FixedUser.validate()
            for user in FixedUser.from_config():
                try:
                    ensure_fixed_user(user, company_id, log=log)
                except Exception as ex:
                    log.error(f"Failed creating fixed user {user.name}: {ex}")

        return empty_dbs
    except Exception as ex:
        log.exception("Failed initializing mongodb")