def test_extract_db_type_from_uri():
uri = "{}://username:password@host:port/database"
for legit_db in DATABASE_ENGINES:
assert legit_db == extract_db_type_from_uri(uri.format(legit_db))
assert legit_db == get_uri_scheme(uri.format(legit_db))
with_driver = legit_db + "+driver-string"
assert legit_db == extract_db_type_from_uri(uri.format(with_driver))
assert legit_db == get_uri_scheme(uri.format(with_driver))
for unsupported_db in ["a", "aa", "sql"]:
with pytest.raises(MlflowException, match="Invalid database engine"):
extract_db_type_from_uri(unsupported_db)
def __init__(self, db_uri, default_artifact_root):
"""
Create a database backed store.
:param db_uri: The SQLAlchemy database URI string to connect to the database. See
the `SQLAlchemy docs
<https://docs.sqlalchemy.org/en/latest/core/engines.html#database-urls>`_
for format specifications. Mlflow supports the dialects ``mysql``,
``mssql``, ``sqlite``, and ``postgresql``.
:param default_artifact_root: Path/URI to location suitable for large data (such as a blob
store object, DBFS path, or shared NFS file system).
"""
super().__init__()
self.db_uri = db_uri
self.db_type = extract_db_type_from_uri(db_uri)
self.artifact_root_uri = default_artifact_root
# Quick check to see if the respective SQLAlchemy database engine has already been created.
if db_uri not in SqlAlchemyStore._db_uri_sql_alchemy_engine_map:
with SqlAlchemyStore._db_uri_sql_alchemy_engine_map_lock:
# Repeat check to prevent race conditions where one thread checks for an existing
# engine while another is creating the respective one, resulting in multiple
# engines being created. It isn't combined with the above check to prevent
# inefficiency from multiple threads waiting for the lock to check for engine
# existence if it has already been created.
if db_uri not in SqlAlchemyStore._db_uri_sql_alchemy_engine_map:
SqlAlchemyStore._db_uri_sql_alchemy_engine_map[
db_uri] = mlflow.store.db.utils.create_sqlalchemy_engine_with_retry(
db_uri)
self.engine = SqlAlchemyStore._db_uri_sql_alchemy_engine_map[db_uri]
# On a completely fresh MLflow installation against an empty database (verify database
# emptiness by checking that 'experiments' etc aren't in the list of table names), run all
# DB migrations
expected_tables = [
SqlExperiment.__tablename__,
SqlRun.__tablename__,
SqlMetric.__tablename__,
SqlParam.__tablename__,
SqlTag.__tablename__,
SqlExperimentTag.__tablename__,
SqlLatestMetric.__tablename__,
]
inspected_tables = set(
sqlalchemy.inspect(self.engine).get_table_names())
if any([table not in inspected_tables for table in expected_tables]):
mlflow.store.db.utils._initialize_tables(self.engine)
Base.metadata.bind = self.engine
SessionMaker = sqlalchemy.orm.sessionmaker(bind=self.engine)
self.ManagedSessionMaker = mlflow.store.db.utils._get_managed_session_maker(
SessionMaker, self.db_type)
mlflow.store.db.utils._verify_schema(self.engine)
if is_local_uri(default_artifact_root):
mkdir(local_file_uri_to_path(default_artifact_root))
if len(self.list_experiments(view_type=ViewType.ALL)) == 0:
with self.ManagedSessionMaker() as session:
self._create_default_experiment(session)
def __init__(self, artifact_uri):
self.db_uri, self.root = extract_db_uri_and_root_path(artifact_uri)
self.db_type = extract_db_type_from_uri(self.db_uri)
self.engine = sqlalchemy.create_engine(self.db_uri)
super(DBArtifactRepository, self).__init__(self.db_uri)
insp = sqlalchemy.inspect(self.engine)
self.expected_tables = set([
SqlArtifact.__tablename__,
])
if len(self.expected_tables & set(insp.get_table_names())) == 0:
DBArtifactRepository._initialize_tables(self.engine)
InitialBase.metadata.bind = self.engine
SessionMaker = sqlalchemy.orm.sessionmaker(bind=self.engine)
self.ManagedSessionMaker = self._get_managed_session_maker(
SessionMaker)
def __init__(self, db_uri, default_artifact_root):
"""
Create a database backed store.
:param db_uri: The SQLAlchemy database URI string to connect to the database. See
the `SQLAlchemy docs
<https://docs.sqlalchemy.org/en/latest/core/engines.html#database-urls>`_
for format specifications. Mlflow supports the dialects ``mysql``,
``mssql``, ``sqlite``, and ``postgresql``.
:param default_artifact_root: Path/URI to location suitable for large data (such as a blob
store object, DBFS path, or shared NFS file system).
"""
super().__init__()
self.db_uri = db_uri
self.db_type = extract_db_type_from_uri(db_uri)
self.artifact_root_uri = default_artifact_root
self.engine = mlflow.store.db.utils.create_sqlalchemy_engine_with_retry(
db_uri)
# On a completely fresh MLflow installation against an empty database (verify database
# emptiness by checking that 'experiments' etc aren't in the list of table names), run all
# DB migrations
expected_tables = [
SqlExperiment.__tablename__,
SqlRun.__tablename__,
SqlMetric.__tablename__,
SqlParam.__tablename__,
SqlTag.__tablename__,
SqlExperimentTag.__tablename__,
SqlLatestMetric.__tablename__,
]
inspected_tables = set(
sqlalchemy.inspect(self.engine).get_table_names())
if any([table not in inspected_tables for table in expected_tables]):
mlflow.store.db.utils._initialize_tables(self.engine)
Base.metadata.bind = self.engine
SessionMaker = sqlalchemy.orm.sessionmaker(bind=self.engine)
self.ManagedSessionMaker = mlflow.store.db.utils._get_managed_session_maker(
SessionMaker, self.db_type)
mlflow.store.db.utils._verify_schema(self.engine)
if is_local_uri(default_artifact_root):
mkdir(local_file_uri_to_path(default_artifact_root))
if len(self.list_experiments()) == 0:
with self.ManagedSessionMaker() as session:
self._create_default_experiment(session)
def __init__(self, db_uri):
"""
Create a database backed store.
:param db_uri: The SQLAlchemy database URI string to connect to the database. See
the `SQLAlchemy docs
<https://docs.sqlalchemy.org/en/latest/core/engines.html#database-urls>`_
for format specifications. Mlflow supports the dialects ``mysql``,
``mssql``, ``sqlite``, and ``postgresql``.
:param default_artifact_root: Path/URI to location suitable for large data (such as a blob
store object, DBFS path, or shared NFS file system).
"""
super(SqlAlchemyStore, self).__init__()
self.db_uri = db_uri
self.db_type = extract_db_type_from_uri(db_uri)
self.engine = mlflow.store.db.utils.create_sqlalchemy_engine(db_uri)
Base.metadata.create_all(self.engine)
# Verify that all model registry tables exist.
SqlAlchemyStore._verify_registry_tables_exist(self.engine)
Base.metadata.bind = self.engine
SessionMaker = sqlalchemy.orm.sessionmaker(bind=self.engine)
self.ManagedSessionMaker = mlflow.store.db.utils._get_managed_session_maker(SessionMaker)
