Generic smart stub

Author: NickCao

packages/jumpstarter/jumpstarter/client/client.py

@@ -8,8 +8,8 @@
 from google.protobuf import empty_pb2
 from jumpstarter_protocol import jumpstarter_pb2_grpc
 
+from .grpc import SmartExporterServiceStub
 from jumpstarter.client import DriverClient
-from jumpstarter.common.grpc import aio_secure_channel
 from jumpstarter.common.importlib import import_class
 
 
@@ -34,12 +34,12 @@ async def client_from_channel(
     reports = {}
     clients = OrderedDict()
 
-    response = await jumpstarter_pb2_grpc.ExporterServiceStub(channel).GetReport(empty_pb2.Empty())
+    response = await SmartExporterServiceStub([channel]).GetReport(empty_pb2.Empty())
 
     if use_alternative_endpoints:
         for endpoint in response.alternative_endpoints:
             if endpoint.certificate:
-                attempted_channel = aio_secure_channel(
+                attempted_channel = grpc.aio.secure_channel(
                     endpoint.endpoint,
                     grpc.ssl_channel_credentials(
                         root_certificates=endpoint.certificate.encode(),

packages/jumpstarter/jumpstarter/client/grpc.py

@@ -3,12 +3,14 @@
 from collections import OrderedDict
 from dataclasses import InitVar, dataclass, field
 from datetime import datetime, timedelta
+from functools import partial
+from typing import Generic, Type, TypeVar
 
 import yaml
 from google.protobuf import duration_pb2, field_mask_pb2, json_format
 from grpc import ChannelConnectivity
 from grpc.aio import Channel
-from jumpstarter_protocol import client_pb2, client_pb2_grpc, jumpstarter_pb2_grpc, kubernetes_pb2
+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.grpc import translate_grpc_exceptions
@@ -254,18 +256,19 @@ async def DeleteLease(self, *, name: str):
             )
 
 
+T = TypeVar("T")
+
+
 @dataclass(frozen=True, slots=True)
-class SmartExporterServiceStub:
+class SmartStub(Generic[T]):
     channels: InitVar[list[Channel]]
+    cls: InitVar[Type]
 
-    __stubs: dict[Channel, jumpstarter_pb2_grpc.ExporterServiceStub] = field(
-        init=False,
-        default_factory=OrderedDict,
-    )
+    __stubs: dict[Channel, T] = field(init=False, default_factory=OrderedDict)
 
-    def __post_init__(self, channels):
+    def __post_init__(self, channels, cls):
         for channel in channels:
-            self.__stubs[channel] = jumpstarter_pb2_grpc.ExporterServiceStub(channel)
+            self.__stubs[channel] = cls(channel)
 
     def __getattr__(self, name):
         for channel, stub in self.__stubs.items():
@@ -274,3 +277,13 @@ def __getattr__(self, name):
                 return getattr(stub, name)
         # or fallback to the last channel (via router)
         return getattr(next(reversed(self.__stubs.values())), name)
+
+
+SmartExporterServiceStub = partial(
+    SmartStub[jumpstarter_pb2_grpc.ExporterServiceStub],
+    cls=jumpstarter_pb2_grpc.ExporterServiceStub,
+)
+SmartRouterServiceStub = partial(
+    SmartStub[router_pb2_grpc.RouterServiceStub],
+    cls=router_pb2_grpc.RouterServiceStub,
+)