Fix ImportError when using zenml login without SQL dependencies (#4252)

Fixes #4251
Fixes tests that were broken on Python 3.10 since PR #4224.

Author: strickvl

src/zenml/cli/login.py

@@ -952,8 +952,8 @@ def login(
     # Get the server that the client is currently connected to, if any
     current_non_local_server: Optional[str] = None
     gc = GlobalConfiguration()
-    store_cfg = gc.store_configuration
-    if store_cfg.type == StoreType.REST:
+    store_cfg = gc.get_store_configuration(allow_default=False)
+    if store_cfg and store_cfg.type == StoreType.REST:
         if not connected_to_local_server():
             current_non_local_server = store_cfg.url
 

src/zenml/config/global_config.py

@@ -16,7 +16,7 @@
 import os
 import uuid
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, Optional, cast
+from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, cast, overload
 from uuid import UUID
 
 from packaging import version
@@ -488,9 +488,25 @@ def get_config_environment_vars(self) -> Dict[str, str]:
 
         return environment_vars
 
-    def _get_store_configuration(
-        self, baseline: Optional[StoreConfiguration] = None
-    ) -> StoreConfiguration:
+    @overload
+    def get_store_configuration(
+        self,
+        baseline: Optional[StoreConfiguration] = ...,
+        allow_default: Literal[True] = ...,
+    ) -> StoreConfiguration: ...
+
+    @overload
+    def get_store_configuration(
+        self,
+        baseline: Optional[StoreConfiguration] = ...,
+        allow_default: Literal[False] = ...,
+    ) -> Optional[StoreConfiguration]: ...
+
+    def get_store_configuration(
+        self,
+        baseline: Optional[StoreConfiguration] = None,
+        allow_default: bool = True,
+    ) -> Optional[StoreConfiguration]:
         """Get the store configuration.
 
         This method computes a store configuration starting from a baseline and
@@ -504,9 +520,12 @@ def _get_store_configuration(
 
         Args:
             baseline: Optional baseline store configuration to use.
+            allow_default: Whether to fall back to the default store
+                configuration if none is set.
 
         Returns:
-            The store configuration.
+            The store configuration or `None` if defaults are disallowed and no
+            configuration is available.
         """
         from zenml.zen_stores.base_zen_store import BaseZenStore
 
@@ -555,15 +574,20 @@ def _get_store_configuration(
                 logger.debug(
                     "Using environment variables to update store config"
                 )
-                if not store:
+                if not store and allow_default:
                     store = self.get_default_store()
-                store = store.model_copy(update=env_store_config, deep=True)
+                if store:
+                    store = store.model_copy(
+                        update=env_store_config, deep=True
+                    )
 
         # Step 2: Only after we've applied the environment variables, we
         # fallback to the default store if no store configuration is set. This
         # is to avoid importing the SQL store config in cases where a rest store
         # is configured with environment variables.
         if not store:
+            if not allow_default:
+                return None
             store = self.get_default_store()
 
         # Step 3: Replace or update the baseline secrets store configuration
@@ -631,7 +655,7 @@ def store_configuration(self) -> StoreConfiguration:
         # configuration from there and disregard the global configuration.
         if self._zen_store is not None:
             return self._zen_store.config
-        return self._get_store_configuration()
+        return self.get_store_configuration()
 
     def get_default_store(self) -> StoreConfiguration:
         """Get the default SQLite store configuration.
@@ -655,7 +679,7 @@ def set_default_store(self) -> None:
         default store.
         """
         # Apply the environment variables to the default store configuration
-        default_store_cfg = self._get_store_configuration(
+        default_store_cfg = self.get_store_configuration(
             baseline=self.get_default_store()
         )
         self._configure_store(default_store_cfg)
@@ -697,8 +721,10 @@ def set_store(
                 constructor.
         """
         # Apply the environment variables to the custom store configuration
-        config = self._get_store_configuration(baseline=config)
-        self._configure_store(config, skip_default_registrations, **kwargs)
+        resolved_config = self.get_store_configuration(baseline=config)
+        self._configure_store(
+            resolved_config, skip_default_registrations, **kwargs
+        )
         logger.info("Updated the global store configuration.")
 
     @property

tests/integration/functional/cli/test_cli.py

@@ -84,31 +84,43 @@ def test_cli_does_not_set_custom_source_root_if_inside_repository(
     mock_set_custom_source_root.assert_not_called()
 
 
-def _mock_rest_store_environment(mocker):
-    """Mock dependencies needed for connect_to_server REST flow."""
+def _mock_rest_store_environment(mocker, login_module):
+    """Mock dependencies needed for connect_to_server REST flow.
+
+    Args:
+        mocker: pytest-mock fixture.
+        login_module: The imported zenml.cli.login module. We pass this
+            explicitly and use patch.object() because string-based patches
+            like "zenml.cli.login.web_login" fail on Python 3.10 due to
+            the star import in zenml/cli/__init__.py shadowing the submodule.
+    """
     credentials_store = mocker.Mock()
     credentials_store.has_valid_credentials.return_value = True
     mocker.patch(
-        "zenml.cli.login.get_credentials_store", return_value=credentials_store
+        "zenml.login.credentials_store.get_credentials_store",
+        return_value=credentials_store,
     )
-    mocker.patch(
-        "zenml.cli.login.BaseZenStore.get_store_type",
+    mocker.patch.object(
+        login_module.BaseZenStore,
+        "get_store_type",
         return_value=StoreType.REST,
     )
-    mocker.patch(
-        "zenml.cli.login.RestZenStoreConfiguration", return_value="rest-config"
+    mocker.patch.object(
+        login_module,
+        "RestZenStoreConfiguration",
+        return_value="rest-config",
     )
-    mocker.patch("zenml.cli.login.cli_utils.declare")
-    mocker.patch("zenml.cli.login.web_login")
+    mocker.patch.object(login_module.cli_utils, "declare")
+    mocker.patch.object(login_module, "web_login")
 
 
 def test_connect_to_server_sets_project_after_success(mocker):
     """Project flag should set the active project after connecting."""
-    _mock_rest_store_environment(mocker)
     login_module = importlib.import_module("zenml.cli.login")
-    mock_gc = mocker.patch("zenml.cli.login.GlobalConfiguration")
+    _mock_rest_store_environment(mocker, login_module)
+    mock_gc = mocker.patch.object(login_module, "GlobalConfiguration")
     mock_gc.return_value.set_store.return_value = None
-    mock_set_project = mocker.patch("zenml.cli.login._set_active_project")
+    mock_set_project = mocker.patch.object(login_module, "_set_active_project")
 
     login_module.connect_to_server(
         url="https://example.com",
@@ -120,13 +132,15 @@ def test_connect_to_server_sets_project_after_success(mocker):
 
 def test_connect_to_server_does_not_set_project_on_failure(mocker):
     """Project change should be skipped if connecting to the store fails."""
-    _mock_rest_store_environment(mocker)
     login_module = importlib.import_module("zenml.cli.login")
-    mock_gc = mocker.patch("zenml.cli.login.GlobalConfiguration")
+    _mock_rest_store_environment(mocker, login_module)
+    mock_gc = mocker.patch.object(login_module, "GlobalConfiguration")
     mock_gc.return_value.set_store.side_effect = IllegalOperationError("boom")
-    mock_set_project = mocker.patch("zenml.cli.login._set_active_project")
-    mocker.patch(
-        "zenml.cli.login.cli_utils.error", side_effect=RuntimeError("exit")
+    mock_set_project = mocker.patch.object(login_module, "_set_active_project")
+    mocker.patch.object(
+        login_module.cli_utils,
+        "error",
+        side_effect=RuntimeError("exit"),
     )
 
     with pytest.raises(RuntimeError):