gracefully handle missing drivers with stub client

Author: evakhoni

packages/jumpstarter/jumpstarter/client/base.py

@@ -7,11 +7,13 @@
 from contextlib import ExitStack, contextmanager
 from dataclasses import field
 
+import click
 from anyio.from_thread import BlockingPortal
 from pydantic import ConfigDict
 from pydantic.dataclasses import dataclass
 
 from .core import AsyncDriverClient
+from jumpstarter.common.importlib import _format_missing_driver_message
 from jumpstarter.streams.blocking import BlockingStream
 
 
@@ -103,3 +105,52 @@ def close(self):
 
     def __del__(self):
         self.close()
+
+
+@dataclass(kw_only=True, config=ConfigDict(arbitrary_types_allowed=True))
+class MissingDriverClient(DriverClient):
+    """Stub client for drivers that are not installed.
+
+    This client is created when a driver client class cannot be imported.
+    It provides a placeholder that raises a clear error when the driver
+    is actually used.
+    """
+
+    def _get_missing_class_path(self) -> str:
+        """Get the missing class path from labels."""
+        return self.labels["jumpstarter.dev/client"]
+
+    def _raise_missing_error(self):
+        """Raise ImportError with installation instructions."""
+        class_path = self._get_missing_class_path()
+        message = _format_missing_driver_message(class_path)
+        raise ImportError(message)
+
+    def call(self, method, *args):
+        """Invoke driver call - raises ImportError since driver is not installed."""
+        self._raise_missing_error()
+
+    def streamingcall(self, method, *args):
+        """Invoke streaming driver call - raises ImportError since driver is not installed."""
+        self._raise_missing_error()
+        # Unreachable yield to make this a generator function for type checking
+        while False:  # noqa: SIM114
+            yield
+
+    @contextmanager
+    def stream(self, method="connect"):
+        """Open a stream - raises ImportError since driver is not installed."""
+        self._raise_missing_error()
+        yield
+
+    def cli(self):
+        """Return a Click command that shows an error when executed."""
+        @click.command()
+        @click.pass_context
+        def unavailable(ctx):
+            class_path = self._get_missing_class_path()
+            message = _format_missing_driver_message(class_path)
+            click.echo(f"Error: {message}", err=True)
+            ctx.exit(1)
+
+        return unavailable

packages/jumpstarter/jumpstarter/client/client.py

@@ -1,3 +1,4 @@
+import logging
 from collections import OrderedDict, defaultdict
 from contextlib import ExitStack, asynccontextmanager
 from graphlib import TopologicalSorter
@@ -9,6 +10,8 @@
 
 from .grpc import MultipathExporterStub
 from jumpstarter.client import DriverClient
+from jumpstarter.client.base import MissingDriverClient
+from jumpstarter.common.exceptions import MissingDriverError
 from jumpstarter.common.importlib import import_class
 
 
@@ -50,7 +53,13 @@ async def client_from_channel(
     for index in TopologicalSorter(topo).static_order():
         report = reports[index]
 
-        client_class = import_class(report.labels["jumpstarter.dev/client"], allow, unsafe)
+        try:
+            client_class = import_class(report.labels["jumpstarter.dev/client"], allow, unsafe)
+        except MissingDriverError as e:
+            # Create stub client instead of failing
+            logger = logging.getLogger(__name__)
+            logger.warning("Driver client '%s' is not available.", e.class_path)
+            client_class = MissingDriverClient
 
         client = client_class(
             uuid=UUID(report.uuid),

packages/jumpstarter/jumpstarter/common/exceptions.py

@@ -79,3 +79,15 @@ class EnvironmentVariableNotSetError(JumpstarterException):
     """Raised when a environment variable is not set."""
 
     pass
+
+
+class MissingDriverError(JumpstarterException):
+    """Raised when a driver module is not found but should be handled gracefully.
+
+    This exception is raised when a driver client class cannot be imported,
+    but the connection should continue with a stub client instead of failing.
+    """
+
+    def __init__(self, message: str, class_path: str):
+        super().__init__(message)
+        self.class_path = class_path

packages/jumpstarter/jumpstarter/common/importlib.py

@@ -5,9 +5,33 @@
 from fnmatch import fnmatchcase
 from importlib import import_module
 
+from jumpstarter.common.exceptions import MissingDriverError
+
 logger = logging.getLogger(__name__)
 
 
+def _format_missing_driver_message(class_path: str) -> str:
+    """Format error message depending on whether the class path is a Jumpstarter driver."""
+    # Extract package name from class path (first component)
+    package_name = class_path.split(".")[0]
+
+    if class_path.startswith("jumpstarter_driver_"):
+        return (
+            f"Driver '{class_path}' is not installed.\n\n"
+            "This usually indicates a version mismatch between your client and the exporter.\n"
+            "Please try to update your client to the latest version and ensure the exporter "
+            "has the correct version installed.\n"
+        )
+    else:
+        return (
+            f"Driver '{class_path}' is not installed.\n\n"
+            "Please install the missing module:\n"
+            f"  pip install {package_name}\n\n"
+            "or if using uv:\n"
+            f"  uv pip install {package_name}"
+        )
+
+
 def cached_import(module_path, class_name):
     # Check whether module is loaded and fully initialized.
     if not (
@@ -40,36 +64,9 @@ def import_class(class_path: str, allow: list[str], unsafe: bool):
     try:
         return cached_import(module_path, class_name)
     except ModuleNotFoundError as e:
-        module_name = str(e).split("'")[1] if "'" in str(e) else str(e).split()[-1]
-
-        is_jumpstarter_driver = unsafe or any(fnmatchcase(class_path, pattern) for pattern in allow)
-
-        if is_jumpstarter_driver:
-            logger.error(
-                "Missing Jumpstarter driver module '%s' for class '%s'. "
-                "This usually indicates a version mismatch between your client and the exporter.",
-                module_name,
-                class_path,
-            )
-            raise ConnectionError(
-                f"Missing Jumpstarter driver module '{module_name}'.\n\n"
-                "This usually indicates a version mismatch between your client and the exporter.\n"
-                "Please try to update your client to the latest version and ensure the exporter "
-                "has the correct version installed.\n"
-            ) from e
-        else:
-            logger.error(
-                "Missing Python module '%s' while importing '%s'. "
-                "This module needs to be installed in your environment.",
-                module_name,
-                class_path,
-            )
-            raise ConnectionError(
-                f"Missing Python module '{module_name}'.\n\n"
-                "Please install the missing module:\n"
-                f"  pip install {module_name}\n\n"
-                "or if using uv:\n"
-                f"  uv pip install {module_name}"
-            ) from e
+        raise MissingDriverError(
+            message=_format_missing_driver_message(class_path),
+            class_path=class_path,
+        ) from e
     except AttributeError as e:
         raise ImportError(f"{module_path} doesn't have specified class {class_name}") from e