wehub-resource-sync
193 lines
7.5 KiB
Python
193 lines
7.5 KiB
Python
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
from sqlalchemy.pool import NullPool
|
|
from sqlalchemy.pool.impl import QueuePool
|
|
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.store.db import utils
|
|
|
|
|
|
def test_create_sqlalchemy_engine_inject_pool_options(monkeypatch):
|
|
monkeypatch.setenv("MLFLOW_SQLALCHEMYSTORE_POOL_SIZE", "2")
|
|
monkeypatch.setenv("MLFLOW_SQLALCHEMYSTORE_POOL_RECYCLE", "3600")
|
|
monkeypatch.setenv("MLFLOW_SQLALCHEMYSTORE_MAX_OVERFLOW", "4")
|
|
monkeypatch.setenv("MLFLOW_SQLALCHEMYSTORE_ECHO", "TRUE")
|
|
monkeypatch.setenv("MLFLOW_SQLALCHEMYSTORE_POOLCLASS", "QueuePool")
|
|
with mock.patch("sqlalchemy.create_engine") as mock_create_engine:
|
|
utils.create_sqlalchemy_engine("mydb://host:port/")
|
|
mock_create_engine.assert_called_once_with(
|
|
"mydb://host:port/",
|
|
pool_pre_ping=True,
|
|
pool_size=2,
|
|
max_overflow=4,
|
|
pool_recycle=3600,
|
|
echo=True,
|
|
poolclass=QueuePool,
|
|
)
|
|
|
|
|
|
def test_create_sqlalchemy_engine_null_pool(monkeypatch):
|
|
monkeypatch.setenv("MLFLOW_SQLALCHEMYSTORE_POOLCLASS", "NullPool")
|
|
with mock.patch("sqlalchemy.create_engine") as mock_create_engine:
|
|
utils.create_sqlalchemy_engine("mydb://host:port/")
|
|
mock_create_engine.assert_called_once_with(
|
|
"mydb://host:port/",
|
|
pool_pre_ping=True,
|
|
poolclass=NullPool,
|
|
)
|
|
|
|
|
|
def test_create_sqlalchemy_engine_invalid_pool(monkeypatch):
|
|
monkeypatch.setenv("MLFLOW_SQLALCHEMYSTORE_POOLCLASS", "SomethingInvalid")
|
|
with mock.patch("sqlalchemy.create_engine"):
|
|
with pytest.raises(ValueError, match=r"Invalid poolclass parameter.*"):
|
|
utils.create_sqlalchemy_engine("mydb://host:port/")
|
|
|
|
|
|
def test_create_sqlalchemy_engine_no_pool_options(monkeypatch):
|
|
monkeypatch.delenv("MLFLOW_SQLALCHEMYSTORE_POOLCLASS", raising=False)
|
|
with mock.patch("sqlalchemy.create_engine") as mock_create_engine:
|
|
utils.create_sqlalchemy_engine("mydb://host:port/")
|
|
mock_create_engine.assert_called_once_with("mydb://host:port/", pool_pre_ping=True)
|
|
|
|
|
|
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
|
|
|
|
|
|
def test_create_sqlalchemy_engine_with_retry_success():
|
|
with (
|
|
mock.patch("sqlalchemy.inspect") as mock_sqlalchemy_inspect,
|
|
mock.patch(
|
|
"mlflow.store.db.utils.create_sqlalchemy_engine", return_value="Engine"
|
|
) as mock_create_sqlalchemy_engine,
|
|
mock.patch("time.sleep") as mock_sleep,
|
|
):
|
|
engine = utils.create_sqlalchemy_engine_with_retry("mydb://host:port/")
|
|
mock_create_sqlalchemy_engine.assert_called_once_with("mydb://host:port/")
|
|
mock_sqlalchemy_inspect.assert_called_once()
|
|
mock_sleep.assert_not_called()
|
|
assert engine == "Engine"
|
|
|
|
|
|
def test_create_sqlalchemy_engine_with_retry_success_after_third_call():
|
|
with (
|
|
mock.patch("sqlalchemy.inspect", side_effect=[Exception, Exception, "Inspect"]),
|
|
mock.patch(
|
|
"mlflow.store.db.utils.create_sqlalchemy_engine", return_value="Engine"
|
|
) as mock_create_sqlalchemy_engine,
|
|
mock.patch("time.sleep"),
|
|
):
|
|
engine = utils.create_sqlalchemy_engine_with_retry("mydb://host:port/")
|
|
assert mock_create_sqlalchemy_engine.mock_calls == [mock.call("mydb://host:port/")] * 3
|
|
assert engine == "Engine"
|
|
|
|
|
|
def test_create_sqlalchemy_engine_with_retry_fail():
|
|
with (
|
|
mock.patch(
|
|
"sqlalchemy.inspect",
|
|
side_effect=[Exception("failed")] * utils.MAX_RETRY_COUNT,
|
|
),
|
|
mock.patch(
|
|
"mlflow.store.db.utils.create_sqlalchemy_engine", return_value="Engine"
|
|
) as mock_create_sqlalchemy_engine,
|
|
mock.patch("time.sleep"),
|
|
):
|
|
with pytest.raises(Exception, match=r"failed"):
|
|
utils.create_sqlalchemy_engine_with_retry("mydb://host:port/")
|
|
assert (
|
|
mock_create_sqlalchemy_engine.mock_calls
|
|
== [mock.call("mydb://host:port/")] * utils.MAX_RETRY_COUNT
|
|
)
|
|
|
|
|
|
def test_mysql_ssl_params(monkeypatch):
|
|
monkeypatch.setenv("MLFLOW_MYSQL_SSL_CA", "/path/to/ca.pem")
|
|
monkeypatch.setenv("MLFLOW_MYSQL_SSL_CERT", "/path/to/cert.pem")
|
|
monkeypatch.setenv("MLFLOW_MYSQL_SSL_KEY", "/path/to/key.pem")
|
|
monkeypatch.setenv("MLFLOW_SQLALCHEMYSTORE_POOLCLASS", "NullPool")
|
|
|
|
with mock.patch("sqlalchemy.create_engine") as mock_create_engine:
|
|
utils.create_sqlalchemy_engine("mysql+pymysql://user@host:port/db")
|
|
|
|
# Check that create_engine was called with the right arguments
|
|
mock_create_engine.assert_called_once_with(
|
|
"mysql+pymysql://user@host:port/db",
|
|
pool_pre_ping=True,
|
|
connect_args={
|
|
"ssl_ca": "/path/to/ca.pem",
|
|
"ssl_cert": "/path/to/cert.pem",
|
|
"ssl_key": "/path/to/key.pem",
|
|
},
|
|
poolclass=NullPool,
|
|
)
|
|
|
|
|
|
def test_mysql_ssl_params_partial(monkeypatch):
|
|
monkeypatch.setenv("MLFLOW_MYSQL_SSL_CA", "/path/to/ca.pem")
|
|
monkeypatch.delenv("MLFLOW_MYSQL_SSL_CERT", raising=False)
|
|
monkeypatch.delenv("MLFLOW_MYSQL_SSL_KEY", raising=False)
|
|
monkeypatch.setenv("MLFLOW_SQLALCHEMYSTORE_POOLCLASS", "NullPool")
|
|
|
|
with mock.patch("sqlalchemy.create_engine") as mock_create_engine:
|
|
utils.create_sqlalchemy_engine("mysql+pymysql://user@host:port/db")
|
|
|
|
# Check that create_engine was called with the right arguments
|
|
mock_create_engine.assert_called_once_with(
|
|
"mysql+pymysql://user@host:port/db",
|
|
pool_pre_ping=True,
|
|
connect_args={
|
|
"ssl_ca": "/path/to/ca.pem",
|
|
},
|
|
poolclass=NullPool,
|
|
)
|
|
|
|
|
|
def test_non_mysql_no_ssl_params(monkeypatch):
|
|
monkeypatch.setenv("MLFLOW_MYSQL_SSL_CA", "/path/to/ca.pem")
|
|
monkeypatch.setenv("MLFLOW_MYSQL_SSL_CERT", "/path/to/cert.pem")
|
|
monkeypatch.setenv("MLFLOW_MYSQL_SSL_KEY", "/path/to/key.pem")
|
|
monkeypatch.setenv("MLFLOW_SQLALCHEMYSTORE_POOLCLASS", "NullPool")
|
|
|
|
with mock.patch("sqlalchemy.create_engine") as mock_create_engine:
|
|
utils.create_sqlalchemy_engine("postgresql://user@host:port/db")
|
|
|
|
# Check that create_engine was called without connect_args
|
|
mock_create_engine.assert_called_once_with(
|
|
"postgresql://user@host:port/db",
|
|
pool_pre_ping=True,
|
|
poolclass=NullPool,
|
|
)
|
|
|
|
|
|
def test_check_sqlite_version_too_old():
|
|
with mock.patch("mlflow.store.db.utils.sqlite3.sqlite_version", "3.30.1"):
|
|
with pytest.raises(MlflowException, match=r"MLflow requires SQLite"):
|
|
utils._check_sqlite_version("sqlite:///test.db")
|
|
|
|
|
|
def test_make_parent_dirs_if_sqlite(tmp_path: Path) -> None:
|
|
nested_path = tmp_path / "a" / "b" / "c" / "test.db"
|
|
db_uri = f"sqlite:///{nested_path}"
|
|
|
|
assert not nested_path.parent.exists()
|
|
utils._make_parent_dirs_if_sqlite(db_uri)
|
|
assert nested_path.parent.exists()
|
|
|
|
|
|
def test_make_parent_dirs_if_sqlite_skips_memory_db() -> None:
|
|
# Should not raise any errors for in-memory databases
|
|
utils._make_parent_dirs_if_sqlite("sqlite:///:memory:")
|
|
utils._make_parent_dirs_if_sqlite("sqlite://")
|
|
|
|
|
|
def test_make_parent_dirs_if_sqlite_skips_non_sqlite() -> None:
|
|
# Should not raise any errors for non-SQLite URIs
|
|
utils._make_parent_dirs_if_sqlite("postgresql://localhost/db")
|
|
utils._make_parent_dirs_if_sqlite("mysql://localhost/db")
|