def test_db_type():
for db_type in ["mysql", "mssql", "postgresql", "sqlite"]:
# should not raise an exception
_validate_db_type_string(db_type)
# error cases
for db_type in ["MySQL", "mongo", "cassandra", "sql", ""]:
with pytest.raises(MlflowException) as e:
_validate_db_type_string(db_type)
assert "Invalid database engine" in e.value.message
def extract_db_type_from_uri(db_uri):
"""
Parse the specified DB URI to extract the database type. Confirm the database type is
supported. If a driver is specified, confirm it passes a plausible regex.
"""
scheme = urllib.parse.urlparse(db_uri).scheme
scheme_plus_count = scheme.count('+')
if scheme_plus_count == 0:
db_type = scheme
elif scheme_plus_count == 1:
db_type, _ = scheme.split('+')
else:
error_msg = "Invalid database URI: '%s'. %s" % (db_uri, _INVALID_DB_URI_MSG)
raise MlflowException(error_msg, INVALID_PARAMETER_VALUE)
_validate_db_type_string(db_type)
return db_type
