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): 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(SqlAlchemyStore, self).__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 = sqlalchemy.create_engine(db_uri, pool_pre_ping=True) insp = sqlalchemy.inspect(self.engine) # 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 = set([ SqlExperiment.__tablename__, SqlRun.__tablename__, SqlMetric.__tablename__, SqlParam.__tablename__, SqlTag.__tablename__, SqlExperimentTag.__tablename__, SqlLatestMetric.__tablename__, ]) if len(expected_tables & set(insp.get_table_names())) == 0: SqlAlchemyStore._initialize_tables(self.engine) Base.metadata.bind = self.engine SessionMaker = sqlalchemy.orm.sessionmaker(bind=self.engine) self.ManagedSessionMaker = self._get_managed_session_maker( SessionMaker) SqlAlchemyStore._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, 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) # Indicate to MLflow that this is an artifact repository plugin self.is_plugin = True