Add enums and exporter status reporting

Author: kirkbrauer

packages/jumpstarter-cli/jumpstarter_cli/get.py

@@ -21,8 +21,8 @@ def get():
 @opt_output_all
 @opt_comma_separated(
     "with",
-    {"leases", "online"},
-    help_text="Include fields: leases, online (comma-separated or repeated)"
+    {"leases", "online", "status"},
+    help_text="Include fields: leases, online, status (comma-separated or repeated)",
 )
 @handle_exceptions_with_reauthentication(relogin_client)
 def get_exporters(config, selector: str | None, output: OutputType, with_options: list[str]):
@@ -32,7 +32,10 @@ def get_exporters(config, selector: str | None, output: OutputType, with_options
 
     include_leases = "leases" in with_options
     include_online = "online" in with_options
-    exporters = config.list_exporters(filter=selector, include_leases=include_leases, include_online=include_online)
+    include_status = "status" in with_options
+    exporters = config.list_exporters(
+        filter=selector, include_leases=include_leases, include_online=include_online, include_status=include_status
+    )
 
     model_print(exporters, output)
 

packages/jumpstarter/jumpstarter/client/core.py

@@ -14,7 +14,7 @@
 from jumpstarter_protocol import jumpstarter_pb2, jumpstarter_pb2_grpc, router_pb2_grpc
 from rich.logging import RichHandler
 
-from jumpstarter.common import Metadata
+from jumpstarter.common import ExporterStatus, Metadata
 from jumpstarter.common.exceptions import JumpstarterException
 from jumpstarter.common.resources import ResourceMetadata
 from jumpstarter.common.serde import decode_value, encode_value
@@ -48,6 +48,12 @@ class DriverInvalidArgument(DriverError, ValueError):
     """
 
 
+class ExporterNotReady(DriverError):
+    """
+    Raised when the exporter is not ready to accept driver calls
+    """
+
+
 @dataclass(kw_only=True)
 class AsyncDriverClient(
     Metadata,
@@ -76,9 +82,28 @@ def __post_init__(self):
             handler = RichHandler()
             self.logger.addHandler(handler)
 
+    async def check_exporter_status(self):
+        """Check if the exporter is ready to accept driver calls"""
+        try:
+            response = await self.stub.GetStatus(jumpstarter_pb2.GetStatusRequest())
+            status = ExporterStatus.from_proto(response.status)
+
+            if status != ExporterStatus.LEASE_READY:
+                raise ExporterNotReady(f"Exporter status is {status}: {response.status_message}")
+
+        except AioRpcError as e:
+            # If GetStatus is not implemented, assume ready for backward compatibility
+            if e.code() == StatusCode.UNIMPLEMENTED:
+                self.logger.debug("GetStatus not implemented, assuming exporter is ready")
+                return
+            raise DriverError(f"Failed to check exporter status: {e.details()}") from e
+
     async def call_async(self, method, *args):
         """Make DriverCall by method name and arguments"""
 
+        # Check exporter status before making the call
+        await self.check_exporter_status()
+
         request = jumpstarter_pb2.DriverCallRequest(
             uuid=str(self.uuid),
             method=method,
@@ -105,6 +130,9 @@ async def call_async(self, method, *args):
     async def streamingcall_async(self, method, *args):
         """Make StreamingDriverCall by method name and arguments"""
 
+        # Check exporter status before making the call
+        await self.check_exporter_status()
+
         request = jumpstarter_pb2.StreamingDriverCallRequest(
             uuid=str(self.uuid),
             method=method,

packages/jumpstarter/jumpstarter/client/grpc.py

@@ -13,13 +13,15 @@
 from jumpstarter_protocol import client_pb2, client_pb2_grpc, jumpstarter_pb2_grpc, kubernetes_pb2, router_pb2_grpc
 from pydantic import BaseModel, ConfigDict, Field, field_serializer
 
+from jumpstarter.common import ExporterStatus
 from jumpstarter.common.grpc import translate_grpc_exceptions
 
 
 @dataclass
 class WithOptions:
     show_online: bool = False
     show_leases: bool = False
+    show_status: bool = False
 
 
 def add_display_columns(table, options: WithOptions = None):
@@ -28,6 +30,8 @@ def add_display_columns(table, options: WithOptions = None):
     table.add_column("NAME")
     if options.show_online:
         table.add_column("ONLINE")
+    if options.show_status:
+        table.add_column("STATUS")
     table.add_column("LABELS")
     if options.show_leases:
         table.add_column("LEASED BY")
@@ -42,6 +46,9 @@ def add_exporter_row(table, exporter, options: WithOptions = None, lease_info: t
     row_data.append(exporter.name)
     if options.show_online:
         row_data.append("yes" if exporter.online else "no")
+    if options.show_status:
+        status_str = str(exporter.status) if exporter.status else "UNKNOWN"
+        row_data.append(status_str)
     row_data.append(",".join(("{}={}".format(k, v) for k, v in sorted(exporter.labels.items()))))
     if options.show_leases:
         if lease_info:
@@ -81,12 +88,16 @@ class Exporter(BaseModel):
     name: str
     labels: dict[str, str]
     online: bool = False
+    status: ExporterStatus | None = None
     lease: Lease | None = None
 
     @classmethod
     def from_protobuf(cls, data: client_pb2.Exporter) -> Exporter:
         namespace, name = parse_exporter_identifier(data.name)
-        return cls(namespace=namespace, name=name, labels=data.labels, online=data.online)
+        status = None
+        if hasattr(data, "status") and data.status:
+            status = ExporterStatus.from_proto(data.status)
+        return cls(namespace=namespace, name=name, labels=data.labels, online=data.online, status=status)
 
     @classmethod
     def rich_add_columns(cls, table, options: WithOptions = None):
@@ -244,6 +255,7 @@ class ExporterList(BaseModel):
     next_page_token: str | None = Field(exclude=True)
     include_online: bool = Field(default=False, exclude=True)
     include_leases: bool = Field(default=False, exclude=True)
+    include_status: bool = Field(default=False, exclude=True)
 
     @classmethod
     def from_protobuf(cls, data: client_pb2.ListExportersResponse) -> ExporterList:
@@ -253,11 +265,15 @@ def from_protobuf(cls, data: client_pb2.ListExportersResponse) -> ExporterList:
         )
 
     def rich_add_columns(self, table):
-        options = WithOptions(show_online=self.include_online, show_leases=self.include_leases)
+        options = WithOptions(
+            show_online=self.include_online, show_leases=self.include_leases, show_status=self.include_status
+        )
         Exporter.rich_add_columns(table, options)
 
     def rich_add_rows(self, table):
-        options = WithOptions(show_online=self.include_online, show_leases=self.include_leases)
+        options = WithOptions(
+            show_online=self.include_online, show_leases=self.include_leases, show_status=self.include_status
+        )
         for exporter in self.exporters:
             exporter.rich_add_rows(table, options)
 
@@ -274,6 +290,8 @@ def model_dump_json(self, **kwargs):
             exclude_fields.add("lease")
         if not self.include_online:
             exclude_fields.add("online")
+        if not self.include_status:
+            exclude_fields.add("status")
 
         data = {"exporters": [exporter.model_dump(mode="json", exclude=exclude_fields) for exporter in self.exporters]}
         return json.dumps(data, **json_kwargs)
@@ -284,6 +302,8 @@ def model_dump(self, **kwargs):
             exclude_fields.add("lease")
         if not self.include_online:
             exclude_fields.add("online")
+        if not self.include_status:
+            exclude_fields.add("status")
 
         return {"exporters": [exporter.model_dump(mode="json", exclude=exclude_fields) for exporter in self.exporters]}
 

packages/jumpstarter/jumpstarter/common/__init__.py

@@ -1,4 +1,12 @@
+from .enums import ExporterStatus, LogSource
 from .metadata import Metadata
 from .tempfile import TemporarySocket, TemporaryTcpListener, TemporaryUnixListener
 
-__all__ = ["Metadata", "TemporarySocket", "TemporaryUnixListener", "TemporaryTcpListener"]
+__all__ = [
+    "ExporterStatus",
+    "LogSource",
+    "Metadata",
+    "TemporarySocket",
+    "TemporaryUnixListener",
+    "TemporaryTcpListener",
+]

packages/jumpstarter/jumpstarter/common/enums.py

@@ -0,0 +1,76 @@
+"""Human-readable enum wrappers for protobuf-generated constants."""
+
+from enum import IntEnum
+
+from jumpstarter_protocol.jumpstarter.v1 import common_pb2
+
+
+class ExporterStatus(IntEnum):
+    """Exporter status states."""
+
+    UNSPECIFIED = common_pb2.EXPORTER_STATUS_UNSPECIFIED
+    """Unknown/unspecified exporter status"""
+
+    OFFLINE = common_pb2.EXPORTER_STATUS_OFFLINE
+    """The exporter is currently offline"""
+
+    AVAILABLE = common_pb2.EXPORTER_STATUS_AVAILABLE
+    """Exporter is available to be leased"""
+
+    BEFORE_LEASE_HOOK = common_pb2.EXPORTER_STATUS_BEFORE_LEASE_HOOK
+    """Exporter is leased, but currently executing before lease hook"""
+
+    LEASE_READY = common_pb2.EXPORTER_STATUS_LEASE_READY
+    """Exporter is leased and ready to accept commands"""
+
+    AFTER_LEASE_HOOK = common_pb2.EXPORTER_STATUS_AFTER_LEASE_HOOK
+    """Lease was releaseed, but exporter is executing after lease hook"""
+
+    BEFORE_LEASE_HOOK_FAILED = common_pb2.EXPORTER_STATUS_BEFORE_LEASE_HOOK_FAILED
+    """The before lease hook failed and the exporter is no longer available"""
+
+    AFTER_LEASE_HOOK_FAILED = common_pb2.EXPORTER_STATUS_AFTER_LEASE_HOOK_FAILED
+    """The after lease hook failed and the exporter is no longer available"""
+
+    def __str__(self):
+        return self.name
+
+    @classmethod
+    def from_proto(cls, value: int) -> "ExporterStatus":
+        """Convert from protobuf integer to enum."""
+        return cls(value)
+
+    def to_proto(self) -> int:
+        """Convert to protobuf integer."""
+        return self.value
+
+
+class LogSource(IntEnum):
+    """Log source types."""
+
+    UNSPECIFIED = common_pb2.LOG_SOURCE_UNSPECIFIED
+    """Unspecified/unknown log source"""
+
+    DRIVER = common_pb2.LOG_SOURCE_DRIVER
+    """Logs produced by a Jumpstarter driver"""
+
+    BEFORE_LEASE_HOOK = common_pb2.LOG_SOURCE_BEFORE_LEASE_HOOK
+    """Logs produced by a before lease hook"""
+
+    AFTER_LEASE_HOOK = common_pb2.LOG_SOURCE_AFTER_LEASE_HOOK
+    """Logs produced by an after lease hook"""
+
+    SYSTEM = common_pb2.LOG_SOURCE_SYSTEM
+    """System/exporter logs"""
+
+    def __str__(self):
+        return self.name
+
+    @classmethod
+    def from_proto(cls, value: int) -> "LogSource":
+        """Convert from protobuf integer to enum."""
+        return cls(value)
+
+    def to_proto(self) -> int:
+        """Convert to protobuf integer."""
+        return self.value

packages/jumpstarter/jumpstarter/config/client.py

@@ -160,12 +160,14 @@ async def list_exporters(
         filter: str | None = None,
         include_leases: bool = False,
         include_online: bool = False,
+        include_status: bool = False,
     ):
         svc = ClientService(channel=await self.channel(), namespace=self.metadata.namespace)
         exporters_response = await svc.ListExporters(page_size=page_size, page_token=page_token, filter=filter)
 
-        # Set the include_online flag for display purposes
+        # Set the include flags for display purposes
         exporters_response.include_online = include_online
+        exporters_response.include_status = include_status
 
         if not include_leases:
             return exporters_response