Merge pull request #4139 from zenml-io/bug/4120-fix-skip-migration-logic

bug:4120 Fix skip migration logic

Author: Json-Andriopoulos

src/zenml/zen_stores/sql_zen_store.py

@@ -1248,6 +1248,14 @@ def filter_and_paginate(
     # Initialization and configuration
     # --------------------------------
 
+    def _run_migrations(self) -> None:
+        if self.skip_migrations or handle_bool_env_var(
+            ENV_ZENML_DISABLE_DATABASE_MIGRATION
+        ):
+            logger.debug("Skipping database migration.")
+        else:
+            self.migrate_database()
+
     def _initialize(self) -> None:
         """Initialize the SQL store."""
         logger.debug("Initializing SqlZenStore at %s", self.config.url)
@@ -1284,11 +1292,7 @@ def _initialize(self) -> None:
 
         self._alembic = Alembic(self.engine)
 
-        if (
-            not self.skip_migrations
-            and ENV_ZENML_DISABLE_DATABASE_MIGRATION not in os.environ
-        ):
-            self.migrate_database()
+        self._run_migrations()
 
         if self.config.driver == SQLDatabaseDriver.SQLITE:
             # Enable foreign key checks at the SQLite database level, but only

tests/unit/zen_stores/__init__.py

@@ -0,0 +1,60 @@
+from unittest.mock import MagicMock
+
+import pytest
+
+from zenml.client import Client
+from zenml.constants import ENV_ZENML_DISABLE_DATABASE_MIGRATION
+from zenml.zen_stores.sql_zen_store import SqlZenStore
+
+
+def test_run_migrations_helper_func(monkeypatch):
+    store = Client().zen_store
+
+    if not isinstance(store, SqlZenStore):
+        pytest.skip(
+            "Run migration helper function is testable only for SQL ZenML store"
+        )
+
+    fake_migrate = MagicMock(return_value=None)
+
+    with monkeypatch.context() as m:
+        m.setattr(SqlZenStore, "migrate_database", fake_migrate)
+        store.skip_migrations = False
+        m.setenv(ENV_ZENML_DISABLE_DATABASE_MIGRATION, "false")
+
+        store._run_migrations()
+
+        fake_migrate.assert_called_once()
+
+
+def test_run_run_migrations_skipped(monkeypatch):
+    store = Client().zen_store
+
+    if not isinstance(store, SqlZenStore):
+        pytest.skip(
+            "Run migration helper function is testable only for SQL ZenML store"
+        )
+
+    fake_migrate = MagicMock(return_value=None)
+
+    # check skip migrations via store.skip_migrations works
+
+    with monkeypatch.context() as m:
+        m.setattr(SqlZenStore, "migrate_database", fake_migrate)
+        store.skip_migrations = True
+        m.setenv(ENV_ZENML_DISABLE_DATABASE_MIGRATION, "false")
+
+        store._run_migrations()
+
+        fake_migrate.assert_not_called()
+
+    # check skip migrations via env var works
+
+    with monkeypatch.context() as m:
+        m.setattr(SqlZenStore, "migrate_database", fake_migrate)
+        store.skip_migrations = False
+        m.setenv(ENV_ZENML_DISABLE_DATABASE_MIGRATION, "true")
+
+        store._run_migrations()
+
+        fake_migrate.assert_not_called()