Fix Kedro in Azure ML with SQLiteSessionStore (#2131)

* Update database.py

Signed-off-by: Sajid Alam <[email protected]>

* Update database.py

Signed-off-by: Sajid Alam <[email protected]>

* Update database.py

Signed-off-by: Sajid Alam <[email protected]>

* Update database.py

Signed-off-by: Sajid Alam <[email protected]>

* Update database.py

Signed-off-by: Sajid Alam <[email protected]>

* Update database.py

Signed-off-by: Sajid Alam <[email protected]>

* lint

Signed-off-by: Sajid Alam <[email protected]>

* add test

Signed-off-by: Sajid Alam <[email protected]>

* Update database.py

Signed-off-by: Sajid Alam <[email protected]>

* revert comments

Signed-off-by: Sajid Alam <[email protected]>

* separate azure logic

Signed-off-by: Sajid Alam <[email protected]>

* Update RELEASE.md

Signed-off-by: Sajid Alam <[email protected]>

* Update RELEASE.md

Signed-off-by: Sajid Alam <[email protected]>

---------

Signed-off-by: Sajid Alam <[email protected]>
Signed-off-by: Sajid Alam <[email protected]>

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

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)

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"