def test_add_data_version_empty_db():
db, db_session = _initialize_db_without_migrations()
# currently the latest is 1, which is also the value we'll put there if the db is not empty so change it to 3 to
# differentiate between the two
original_latest_data_version = mlrun.api.initial_data.latest_data_version
mlrun.api.initial_data.latest_data_version = "3"
assert db.get_current_data_version(db_session, raise_on_not_found=False) is None
mlrun.api.initial_data._add_initial_data(db_session)
assert (
db.get_current_data_version(db_session, raise_on_not_found=True)
== mlrun.api.initial_data.latest_data_version
)
mlrun.api.initial_data.latest_data_version = original_latest_data_version
def test_perform_data_migrations_from_zero_version():
db, db_session = _initialize_db_without_migrations()
# set version to 0
db.create_data_version(db_session, "0")
original_perform_version_1_data_migrations = (
mlrun.api.initial_data._perform_version_1_data_migrations
)
mlrun.api.initial_data._perform_version_1_data_migrations = unittest.mock.Mock()
mlrun.api.initial_data._perform_data_migrations(db_session)
mlrun.api.initial_data._perform_version_1_data_migrations.assert_called_once()
# calling again should trigger migrations again
mlrun.api.initial_data._perform_data_migrations(db_session)
mlrun.api.initial_data._perform_version_1_data_migrations.assert_called_once()
mlrun.api.initial_data._perform_version_1_data_migrations = (
original_perform_version_1_data_migrations
)
assert db.get_current_data_version(db_session, raise_on_not_found=True) == str(
mlrun.api.initial_data.latest_data_version
)
def test_add_data_version_non_empty_db():
db, db_session = _initialize_db_without_migrations()
# currently the latest is 1, which is also the value we'll put there if the db is not empty so change it to 3 to
# differentiate between the two
original_latest_data_version = mlrun.api.initial_data.latest_data_version
mlrun.api.initial_data.latest_data_version = "3"
assert db.get_current_data_version(db_session, raise_on_not_found=False) is None
# fill db
db.create_project(
db_session,
mlrun.api.schemas.Project(
metadata=mlrun.api.schemas.ProjectMetadata(name="project-name"),
),
)
mlrun.api.initial_data._add_initial_data(db_session)
assert db.get_current_data_version(db_session, raise_on_not_found=True) == "1"
mlrun.api.initial_data.latest_data_version = original_latest_data_version
def _add_data_version(db: mlrun.api.db.sqldb.db.SQLDB,
db_session: sqlalchemy.orm.Session):
if db.get_current_data_version(db_session,
raise_on_not_found=False) is None:
data_version = _resolve_current_data_version(db, db_session)
logger.info(
"No data version, setting data version",
data_version=data_version,
)
db.create_data_version(db_session, data_version)
def _perform_data_migrations(db_session: sqlalchemy.orm.Session):
if config.httpdb.db.data_migrations_mode == "enabled":
# FileDB is not really a thing anymore, so using SQLDB directly
db = mlrun.api.db.sqldb.db.SQLDB("")
current_data_version = int(db.get_current_data_version(db_session))
if current_data_version != latest_data_version:
logger.info(
"Performing data migrations",
current_data_version=current_data_version,
latest_data_version=latest_data_version,
)
if current_data_version < 1:
_perform_version_1_data_migrations(db, db_session)
db.create_data_version(db_session, str(latest_data_version))
def _resolve_current_data_version(db: mlrun.api.db.sqldb.db.SQLDB,
db_session: sqlalchemy.orm.Session):
try:
return int(db.get_current_data_version(db_session))
except (
sqlalchemy.exc.ProgrammingError,
sqlalchemy.exc.OperationalError,
pymysql.err.ProgrammingError,
pymysql.err.OperationalError,
mlrun.errors.MLRunNotFoundError,
) as exc:
try:
projects = db.list_projects(db_session)
except (
sqlalchemy.exc.ProgrammingError,
sqlalchemy.exc.OperationalError,
pymysql.err.ProgrammingError,
pymysql.err.OperationalError,
):
projects = None
# heuristic - if there are no projects it's a new DB - data version is latest
if not projects or not projects.projects:
logger.info(
"No projects in DB, assuming latest data version",
exc=exc,
latest_data_version=latest_data_version,
)
return latest_data_version
elif "no such table" in str(exc) or ("Table" in str(exc)
and "doesn't exist" in str(exc)):
logger.info(
"Data version table does not exist, assuming prior version",
exc=exc,
data_version_prior_to_table_addition=
data_version_prior_to_table_addition,
)
return data_version_prior_to_table_addition
elif isinstance(exc, mlrun.errors.MLRunNotFoundError):
logger.info(
"Data version table exist without version, assuming prior version",
exc=exc,
data_version_prior_to_table_addition=
data_version_prior_to_table_addition,
)
return data_version_prior_to_table_addition
raise exc