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 get_store(self, store_uri=None, artifact_uri=None):
"""Get a store from the registry based on the scheme of store_uri
:param store_uri: The store URI. If None, it will be inferred from the environment. This URI
is used to select which tracking store implementation to instantiate and
is passed to the constructor of the implementation.
:param artifact_uri: Artifact repository URI. Passed through to the tracking store
implementation.
:return: An instance of `mlflow.store.AbstractStore` that fulfills the store URI
requirements.
"""
from mlflow.tracking import utils
store_uri = store_uri if store_uri is not None else utils.get_tracking_uri(
)
scheme = store_uri if store_uri == "databricks" else get_uri_scheme(
store_uri)
try:
store_builder = self._registry[scheme]
except KeyError:
raise MlflowException(
"Unexpected URI scheme '{}' for tracking store. "
"Valid schemes are: {}".format(store_uri,
list(self._registry.keys())))
return store_builder(store_uri=store_uri, artifact_uri=artifact_uri)
def get_artifact_repository(self, artifact_uri):
"""Get an artifact repository from the registry based on the scheme of artifact_uri
:param store_uri: The store URI. This URI is used to select which artifact repository
implementation to instantiate and is passed to the
constructor of the implementation.
:return: An instance of `mlflow.store.ArtifactRepository` that fulfills the artifact URI
requirements.
"""
scheme = get_uri_scheme(artifact_uri)
repository = self._registry.get(scheme)
if repository is None:
raise MlflowException(
"Could not find a registered artifact repository for: {}. "
"Currently registered schemes are: {}".format(
artifact_uri, list(self._registry.keys())))
return repository(artifact_uri)
def get_artifact_repository_type(artifact_uri):
scheme = get_uri_scheme(artifact_uri)
if scheme == 'sqlite' or scheme == 'mssql':
return ArtifactRepositoryType.DB
else:
return ArtifactRepositoryType.FileSystem
