def test_sqlalchemy_store_detects_schema_mismatch(
tmpdir, db_url
): # pylint: disable=unused-argument
def _assert_invalid_schema(engine):
with pytest.raises(MlflowException) as ex:
_verify_schema(engine)
assert ex.message.contains("Detected out-of-date database schema.")
# Initialize an empty database & verify that we detect a schema mismatch
engine = sqlalchemy.create_engine(db_url)
_assert_invalid_schema(engine)
# Create legacy tables, verify schema is still out of date
InitialBase.metadata.create_all(engine)
_assert_invalid_schema(engine)
# Run each migration. Until the last one, schema should be out of date
config = _get_alembic_config(db_url)
script = ScriptDirectory.from_config(config)
revisions = list(script.walk_revisions())
revisions.reverse()
for rev in revisions[:-1]:
command.upgrade(config, rev.revision)
_assert_invalid_schema(engine)
# Run migrations, schema verification should now pass
invoke_cli_runner(mlflow.db.commands, ["upgrade", db_url])
_verify_schema(engine)
def _get_latest_schema_revision():
"""Get latest schema revision as a string."""
# We aren't executing any commands against a DB, so we leave the DB URL unspecified
config = _get_alembic_config(db_url="")
script = ScriptDirectory.from_config(config)
heads = script.get_heads()
if len(heads) != 1:
raise MlflowException("Migration script directory was in unexpected state. Got %s head "
"database versions but expected only 1. Found versions: %s"
% (len(heads), heads))
return heads[0]
def test_alembic_escape_logic():
url = "fakesql://cooluser%40stillusername:apassword@localhost:3306/testingdb"
config = utils._get_alembic_config(url)
assert config.get_main_option("sqlalchemy.url") == url
