diff --git a/RELEASE.md b/RELEASE.md index e0b295425..66d6a51a2 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -17,6 +17,7 @@ Please follow the established format: - Improve `kedro viz build` usage documentation (#2126) - Fix unserializable parameters value (#2122) - Display full dataset type with library prefix in metadata panel (#2136) +- Enable SQLite WAL mode for Azure ML to fix database locking issues (#2131) # Release 10.0.0 diff --git a/package/kedro_viz/database.py b/package/kedro_viz/database.py index 4d97ca318..5a62e32bc 100644 --- a/package/kedro_viz/database.py +++ b/package/kedro_viz/database.py @@ -1,18 +1,41 @@ """Database management layer based on SQLAlchemy""" -from sqlalchemy import create_engine +import os + +from sqlalchemy import create_engine, text from sqlalchemy.orm import sessionmaker from kedro_viz.models.experiment_tracking import Base +def configure_wal_for_azure(engine): + """Applies WAL mode to SQLite if running in an Azure ML environment.""" + is_azure_ml = any( + var in os.environ + for var in [ + "AZUREML_ARM_SUBSCRIPTION", + "AZUREML_ARM_RESOURCEGROUP", + "AZUREML_RUN_ID", + ] + ) + if is_azure_ml: + with engine.connect() as conn: + conn.execute(text("PRAGMA journal_mode=WAL;")) + + def make_db_session_factory(session_store_location: str) -> sessionmaker: """SQLAlchemy connection to a SQLite DB""" database_url = f"sqlite:///{session_store_location}" engine = create_engine(database_url, connect_args={"check_same_thread": False}) - session_class = sessionmaker(engine) # TODO: making db session factory shouldn't depend on models. # So want to move the table creation elsewhere ideally. # But this means returning engine as well as session class. + + # Check if we are running in an Azure ML environment if so enable WAL mode. + configure_wal_for_azure(engine) + + # Create the database tables if they do not exist. Base.metadata.create_all(bind=engine) - return session_class + + # Return a session factory bound to the engine. + return sessionmaker(bind=engine) diff --git a/package/tests/test_integrations/test_sqlite_store.py b/package/tests/test_integrations/test_sqlite_store.py index a90d72b9b..ec14c6873 100644 --- a/package/tests/test_integrations/test_sqlite_store.py +++ b/package/tests/test_integrations/test_sqlite_store.py @@ -8,7 +8,7 @@ import boto3 import pytest from moto import mock_aws -from sqlalchemy import create_engine, func, select +from sqlalchemy import create_engine, func, select, text from sqlalchemy.orm import sessionmaker from kedro_viz.database import make_db_session_factory @@ -372,3 +372,22 @@ def test_sync_with_merge_error(self, mocker, store_path, remote_path, caplog): mock_merge.assert_called_once() mock_upload.assert_called_once() assert "Merge failed on sync: Merge failed" in caplog.text + + def test_make_db_session_factory_with_azure_env_var(self, mocker, tmp_path): + """Test that WAL mode is enabled when running in an Azure environment.""" + mocker.patch.dict( + os.environ, + { + "AZUREML_ARM_SUBSCRIPTION": "dummy_value", + "AZUREML_ARM_RESOURCEGROUP": "dummy_value", + }, + ) + db_location = str(tmp_path / "test_session_store.db") + session_class = make_db_session_factory(db_location) + + # Ensure that the session can be created without issues. + with session_class() as session: + assert session is not None + # Check if the database is using WAL mode by querying the PRAGMA + result = session.execute(text("PRAGMA journal_mode;")).scalar() + assert result == "wal"