Merge pull request #344 from jumpstarter-dev/grpc-exception

Global exception handling for j command

Author: NickCao

packages/jumpstarter-cli-client/jumpstarter_cli_client/__init__.py

@@ -1,8 +1,10 @@
 import logging
+import sys
 from typing import Optional
 
 import asyncclick as click
 from jumpstarter_cli_common import AliasedGroup, opt_log_level, version
+from jumpstarter_cli_common.exceptions import handle_exceptions
 
 from .client_config import create_client_config, delete_client_config, list_client_configs, use_client_config
 from .client_exporter import list_client_exporters
@@ -24,7 +26,16 @@ def client(log_level: Optional[str]):
 
 def j():
     with env() as client:
-        client.cli()(standalone_mode=True)
+
+        @handle_exceptions
+        def cli():
+            client.cli()(standalone_mode=False)
+
+        try:
+            cli()
+        except click.ClickException as e:
+            e.show()
+            sys.exit(1)
 
 
 client.add_command(create_client_config)

packages/jumpstarter/jumpstarter/client/core.py

@@ -13,6 +13,7 @@
 from jumpstarter_protocol import jumpstarter_pb2, jumpstarter_pb2_grpc, router_pb2_grpc
 
 from jumpstarter.common import Metadata
+from jumpstarter.common.exceptions import JumpstarterException
 from jumpstarter.common.resources import ResourceMetadata
 from jumpstarter.common.serde import decode_value, encode_value
 from jumpstarter.common.streams import (
@@ -29,6 +30,24 @@
 )
 
 
+class DriverError(JumpstarterException):
+    """
+    Raised when a driver call returns an error
+    """
+
+
+class DriverMethodNotImplemented(DriverError, NotImplementedError):
+    """
+    Raised when a driver method is not implemented
+    """
+
+
+class DriverInvalidArgument(DriverError, ValueError):
+    """
+    Raised when a driver method is called with invalid arguments
+    """
+
+
 @dataclass(kw_only=True)
 class AsyncDriverClient(
     Metadata,
@@ -74,11 +93,13 @@ async def call_async(self, method, *args):
         except AioRpcError as e:
             match e.code():
                 case StatusCode.UNIMPLEMENTED:
-                    raise NotImplementedError(e.details()) from None
+                    raise DriverMethodNotImplemented(e.details()) from None
                 case StatusCode.INVALID_ARGUMENT:
-                    raise ValueError(e.details()) from None
+                    raise DriverInvalidArgument(e.details()) from None
+                case StatusCode.UNKNOWN:
+                    raise DriverError(e.details()) from None
                 case _:
-                    raise
+                    raise DriverError(e.details()) from e
 
         return decode_value(response.result)
 
@@ -97,11 +118,13 @@ async def streamingcall_async(self, method, *args):
         except AioRpcError as e:
             match e.code():
                 case StatusCode.UNIMPLEMENTED:
-                    raise NotImplementedError(e.details()) from None
+                    raise DriverMethodNotImplemented(e.details()) from None
                 case StatusCode.INVALID_ARGUMENT:
-                    raise ValueError(e.details()) from None
+                    raise DriverInvalidArgument(e.details()) from None
+                case StatusCode.UNKNOWN:
+                    raise DriverError(e.details()) from None
                 case _:
-                    raise
+                    raise DriverError(e.details()) from e
 
     @asynccontextmanager
     async def stream_async(self, method):

packages/jumpstarter/jumpstarter/driver/base.py

@@ -108,6 +108,8 @@ async def DriverCall(self, request, context):
             await context.abort(StatusCode.INVALID_ARGUMENT, str(e))
         except TimeoutError as e:
             await context.abort(StatusCode.DEADLINE_EXCEEDED, str(e))
+        except Exception as e:
+            await context.abort(StatusCode.UNKNOWN, str(e))
 
     async def StreamingDriverCall(self, request, context):
         """
@@ -136,6 +138,8 @@ async def StreamingDriverCall(self, request, context):
             await context.abort(StatusCode.INVALID_ARGUMENT, str(e))
         except TimeoutError as e:
             await context.abort(StatusCode.DEADLINE_EXCEEDED, str(e))
+        except Exception as e:
+            await context.abort(StatusCode.UNKNOWN, str(e))
 
     @asynccontextmanager
     async def Stream(self, request, context):