translate_grpc_exceptions in grpc wrapper

Author: NickCao

packages/jumpstarter/jumpstarter/client/grpc.py

@@ -9,6 +9,8 @@
 from jumpstarter_protocol import client_pb2, client_pb2_grpc, kubernetes_pb2
 from pydantic import BaseModel, ConfigDict, Field, field_serializer
 
+from jumpstarter.common.grpc import translate_grpc_exceptions
+
 
 def parse_identifier(identifier: str, kind: str) -> (str, str):
     segments = identifier.split("/")
@@ -143,11 +145,12 @@ def __post_init__(self):
         self.stub = client_pb2_grpc.ClientServiceStub(channel=self.channel)
 
     async def GetExporter(self, *, name: str):
-        exporter = await self.stub.GetExporter(
-            client_pb2.GetExporterRequest(
-                name="namespaces/{}/exporters/{}".format(self.namespace, name),
+        with translate_grpc_exceptions():
+            exporter = await self.stub.GetExporter(
+                client_pb2.GetExporterRequest(
+                    name="namespaces/{}/exporters/{}".format(self.namespace, name),
+                )
             )
-        )
         return Exporter.from_protobuf(exporter)
 
     async def ListExporters(
@@ -157,22 +160,24 @@ async def ListExporters(
         page_token: str | None = None,
         filter: str | None = None,
     ):
-        exporters = await self.stub.ListExporters(
-            client_pb2.ListExportersRequest(
-                parent="namespaces/{}".format(self.namespace),
-                page_size=page_size,
-                page_token=page_token,
-                filter=filter,
+        with translate_grpc_exceptions():
+            exporters = await self.stub.ListExporters(
+                client_pb2.ListExportersRequest(
+                    parent="namespaces/{}".format(self.namespace),
+                    page_size=page_size,
+                    page_token=page_token,
+                    filter=filter,
+                )
             )
-        )
         return ExporterList.from_protobuf(exporters)
 
     async def GetLease(self, *, name: str):
-        lease = await self.stub.GetLease(
-            client_pb2.GetLeaseRequest(
-                name="namespaces/{}/leases/{}".format(self.namespace, name),
+        with translate_grpc_exceptions():
+            lease = await self.stub.GetLease(
+                client_pb2.GetLeaseRequest(
+                    name="namespaces/{}/leases/{}".format(self.namespace, name),
+                )
             )
-        )
         return Lease.from_protobuf(lease)
 
     async def ListLeases(
@@ -182,14 +187,15 @@ async def ListLeases(
         page_token: str | None = None,
         filter: str | None = None,
     ):
-        leases = await self.stub.ListLeases(
-            client_pb2.ListLeasesRequest(
-                parent="namespaces/{}".format(self.namespace),
-                page_size=page_size,
-                page_token=page_token,
-                filter=filter,
+        with translate_grpc_exceptions():
+            leases = await self.stub.ListLeases(
+                client_pb2.ListLeasesRequest(
+                    parent="namespaces/{}".format(self.namespace),
+                    page_size=page_size,
+                    page_token=page_token,
+                    filter=filter,
+                )
             )
-        )
         return LeaseList.from_protobuf(leases)
 
     async def CreateLease(
@@ -201,15 +207,16 @@ async def CreateLease(
         duration_pb = duration_pb2.Duration()
         duration_pb.FromTimedelta(duration)
 
-        lease = await self.stub.CreateLease(
-            client_pb2.CreateLeaseRequest(
-                parent="namespaces/{}".format(self.namespace),
-                lease=client_pb2.Lease(
-                    duration=duration_pb,
-                    selector=selector,
-                ),
+        with translate_grpc_exceptions():
+            lease = await self.stub.CreateLease(
+                client_pb2.CreateLeaseRequest(
+                    parent="namespaces/{}".format(self.namespace),
+                    lease=client_pb2.Lease(
+                        duration=duration_pb,
+                        selector=selector,
+                    ),
+                )
             )
-        )
         return Lease.from_protobuf(lease)
 
     async def UpdateLease(
@@ -224,20 +231,22 @@ async def UpdateLease(
         update_mask = field_mask_pb2.FieldMask()
         update_mask.FromJsonString("duration")
 
-        lease = await self.stub.UpdateLease(
-            client_pb2.UpdateLeaseRequest(
-                lease=client_pb2.Lease(
-                    name="namespaces/{}/leases/{}".format(self.namespace, name),
-                    duration=duration_pb,
-                ),
-                update_mask=update_mask,
+        with translate_grpc_exceptions():
+            lease = await self.stub.UpdateLease(
+                client_pb2.UpdateLeaseRequest(
+                    lease=client_pb2.Lease(
+                        name="namespaces/{}/leases/{}".format(self.namespace, name),
+                        duration=duration_pb,
+                    ),
+                    update_mask=update_mask,
+                )
             )
-        )
         return Lease.from_protobuf(lease)
 
     async def DeleteLease(self, *, name: str):
-        await self.stub.DeleteLease(
-            client_pb2.DeleteLeaseRequest(
-                name="namespaces/{}/leases/{}".format(self.namespace, name),
+        with translate_grpc_exceptions():
+            await self.stub.DeleteLease(
+                client_pb2.DeleteLeaseRequest(
+                    name="namespaces/{}/leases/{}".format(self.namespace, name),
+                )
             )
-        )

packages/jumpstarter/jumpstarter/config/client.py

@@ -15,7 +15,7 @@
 from .tls import TLSConfigV1Alpha1
 from jumpstarter.client.grpc import ClientService
 from jumpstarter.common.exceptions import FileNotFoundError
-from jumpstarter.common.grpc import aio_secure_channel, ssl_channel_credentials, translate_grpc_exceptions
+from jumpstarter.common.grpc import aio_secure_channel, ssl_channel_credentials
 
 
 def _allow_from_env():
@@ -109,8 +109,7 @@ def update_lease(self, name, duration: timedelta):
 
     async def get_exporter_async(self, name: str):
         svc = ClientService(channel=await self.channel(), namespace=self.metadata.namespace)
-        with translate_grpc_exceptions():
-            return await svc.GetExporter(name=name)
+        return await svc.GetExporter(name=name)
 
     async def list_exporters_async(
         self,
@@ -119,37 +118,32 @@ async def list_exporters_async(
         filter: str | None = None,
     ):
         svc = ClientService(channel=await self.channel(), namespace=self.metadata.namespace)
-        with translate_grpc_exceptions():
-            return await svc.ListExporters(page_size=page_size, page_token=page_token, filter=filter)
+        return await svc.ListExporters(page_size=page_size, page_token=page_token, filter=filter)
 
     async def create_lease_async(
         self,
         selector: str,
         duration: timedelta,
     ):
         svc = ClientService(channel=await self.channel(), namespace=self.metadata.namespace)
-        with translate_grpc_exceptions():
-            return await svc.CreateLease(
-                selector=selector,
-                duration=duration,
-            )
+        return await svc.CreateLease(
+            selector=selector,
+            duration=duration,
+        )
 
     async def delete_lease_async(self, name: str):
         svc = ClientService(channel=await self.channel(), namespace=self.metadata.namespace)
-        with translate_grpc_exceptions():
-            await svc.DeleteLease(
-                name=name,
-            )
+        await svc.DeleteLease(
+            name=name,
+        )
 
     async def list_leases_async(self, filter: str):
         svc = ClientService(channel=await self.channel(), namespace=self.metadata.namespace)
-        with translate_grpc_exceptions():
-            return await svc.ListLeases(filter=filter)
+        return await svc.ListLeases(filter=filter)
 
     async def update_lease_async(self, name, duration: timedelta):
         svc = ClientService(channel=await self.channel(), namespace=self.metadata.namespace)
-        with translate_grpc_exceptions():
-            return await svc.UpdateLease(name=name, duration=duration)
+        return await svc.UpdateLease(name=name, duration=duration)
 
     @asynccontextmanager
     async def lease_async(