4949# pylint: enable=no-name-in-module
5050DEFAULT_DISPATCH_PORT = 443
5151
52+ DEFAULT_GRPC_CALL_TIMEOUT = timedelta (seconds = 60 )
53+ """Default gRPC call timeout."""
54+
55+ DEFAULT_GRPC_STREAM_TIMEOUT = timedelta (minutes = 5 )
56+ """Default gRPC stream timeout."""
57+
5258
5359class DispatchApiClient (BaseApiClient [dispatch_pb2_grpc .MicrogridDispatchServiceStub ]):
5460 """Dispatch API client."""
5561
62+ # pylint: disable-next=too-many-arguments
5663 def __init__ (
5764 self ,
5865 * ,
5966 server_url : str ,
6067 key : str ,
6168 connect : bool = True ,
69+ call_timeout : timedelta | None = None ,
70+ stream_timeout : timedelta | None = None ,
6271 ) -> None :
6372 """Initialize the client.
6473
6574 Args:
6675 server_url: The URL of the server to connect to.
6776 key: API key to use for authentication.
6877 connect: Whether to connect to the service immediately.
78+ call_timeout: Timeout for gRPC calls, default is 60 seconds.
79+ stream_timeout: Timeout for gRPC streams, default is 5 minutes.
6980 """
7081 super ().__init__ (
7182 server_url ,
@@ -82,6 +93,11 @@ def __init__(
8293 ] = {}
8394 """A dictionary of streamers, keyed by microgrid_id."""
8495
96+ self ._call_timeout = (call_timeout or DEFAULT_GRPC_CALL_TIMEOUT ).total_seconds ()
97+ self ._stream_timeout = (
98+ stream_timeout or DEFAULT_GRPC_STREAM_TIMEOUT
99+ ).total_seconds ()
100+
85101 @property
86102 def stub (self ) -> dispatch_pb2_grpc .MicrogridDispatchServiceAsyncStub :
87103 """The stub for the service."""
@@ -177,7 +193,9 @@ def to_interval(
177193 while True :
178194 response = await cast (
179195 Awaitable [ListMicrogridDispatchesResponse ],
180- self .stub .ListMicrogridDispatches (request , metadata = self ._metadata ),
196+ self .stub .ListMicrogridDispatches (
197+ request , metadata = self ._metadata , timeout = self ._call_timeout
198+ ),
181199 )
182200
183201 yield (Dispatch .from_protobuf (dispatch ) for dispatch in response .dispatches )
@@ -234,7 +252,9 @@ def _get_stream(
234252 stream_method = lambda : cast (
235253 AsyncIterator [StreamMicrogridDispatchesResponse ],
236254 self .stub .StreamMicrogridDispatches (
237- request , metadata = self ._metadata
255+ request ,
256+ metadata = self ._metadata ,
257+ timeout = self ._stream_timeout ,
238258 ),
239259 ),
240260 transform = DispatchEvent .from_protobuf ,
@@ -303,7 +323,9 @@ async def create( # pylint: disable=too-many-positional-arguments
303323 response = await cast (
304324 Awaitable [CreateMicrogridDispatchResponse ],
305325 self .stub .CreateMicrogridDispatch (
306- request .to_protobuf (), metadata = self ._metadata
326+ request .to_protobuf (),
327+ metadata = self ._metadata ,
328+ timeout = self ._call_timeout ,
307329 ),
308330 )
309331
@@ -394,7 +416,9 @@ async def update(
394416
395417 response = await cast (
396418 Awaitable [UpdateMicrogridDispatchResponse ],
397- self .stub .UpdateMicrogridDispatch (msg , metadata = self ._metadata ),
419+ self .stub .UpdateMicrogridDispatch (
420+ msg , metadata = self ._metadata , timeout = self ._call_timeout
421+ ),
398422 )
399423
400424 return Dispatch .from_protobuf (response .dispatch )
@@ -414,7 +438,9 @@ async def get(self, *, microgrid_id: int, dispatch_id: int) -> Dispatch:
414438 )
415439 response = await cast (
416440 Awaitable [GetMicrogridDispatchResponse ],
417- self .stub .GetMicrogridDispatch (request , metadata = self ._metadata ),
441+ self .stub .GetMicrogridDispatch (
442+ request , metadata = self ._metadata , timeout = self ._call_timeout
443+ ),
418444 )
419445 return Dispatch .from_protobuf (response .dispatch )
420446
@@ -430,5 +456,7 @@ async def delete(self, *, microgrid_id: int, dispatch_id: int) -> None:
430456 )
431457 await cast (
432458 Awaitable [None ],
433- self .stub .DeleteMicrogridDispatch (request , metadata = self ._metadata ),
459+ self .stub .DeleteMicrogridDispatch (
460+ request , metadata = self ._metadata , timeout = self ._call_timeout
461+ ),
434462 )
0 commit comments