Implement SmartExporterServiceStub

Author: NickCao

packages/jumpstarter/jumpstarter/client/grpc.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
-from dataclasses import dataclass, field
+from collections import OrderedDict
+from dataclasses import InitVar, dataclass, field
 from datetime import datetime, timedelta
 
 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, kubernetes_pb2
+from jumpstarter_protocol import client_pb2, client_pb2_grpc, jumpstarter_pb2_grpc, kubernetes_pb2
 from pydantic import BaseModel, ConfigDict, Field, field_serializer
 
 from jumpstarter.common.grpc import translate_grpc_exceptions
@@ -250,3 +252,25 @@ async def DeleteLease(self, *, name: str):
                     name="namespaces/{}/leases/{}".format(self.namespace, name),
                 )
             )
+
+
+@dataclass(frozen=True, slots=True)
+class SmartExporterServiceStub:
+    channels: InitVar[list[Channel]]
+
+    __stubs: dict[Channel, jumpstarter_pb2_grpc.ExporterServiceStub] = field(
+        init=False,
+        default_factory=OrderedDict,
+    )
+
+    def __post_init__(self, channels):
+        for channel in channels:
+            self.__stubs[channel] = jumpstarter_pb2_grpc.ExporterServiceStub(channel)
+
+    def __getattr__(self, name):
+        for channel, stub in self.__stubs.items():
+            # find the first channel that's ready
+            if channel.get_state(try_to_connect=True) == ChannelConnectivity.READY:
+                return getattr(stub, name)
+        # or fallback to the last channel (via router)
+        return getattr(next(reversed(self.__stubs.values())), name)