5353class DispatchApiClient (BaseApiClient [dispatch_pb2_grpc .MicrogridDispatchServiceStub ]):
5454 """Dispatch API client."""
5555
56+ # pylint: disable-next=too-many-arguments
5657 def __init__ (
5758 self ,
5859 * ,
5960 server_url : str ,
6061 key : str ,
6162 connect : bool = True ,
63+ call_timeout : timedelta = timedelta (seconds = 60 ),
64+ stream_timeout : timedelta = timedelta (minutes = 5 ),
6265 ) -> None :
6366 """Initialize the client.
6467
6568 Args:
6669 server_url: The URL of the server to connect to.
6770 key: API key to use for authentication.
6871 connect: Whether to connect to the service immediately.
72+ call_timeout: Timeout for gRPC calls, default is 60 seconds.
73+ stream_timeout: Timeout for gRPC streams, default is 5 minutes.
6974 """
7075 super ().__init__ (
7176 server_url ,
@@ -82,6 +87,19 @@ def __init__(
8287 ] = {}
8388 """A dictionary of streamers, keyed by microgrid_id."""
8489
90+ self ._call_timeout_seconds = call_timeout .total_seconds ()
91+ self ._stream_timeout_seconds = stream_timeout .total_seconds ()
92+
93+ @property
94+ def call_timeout (self ) -> timedelta :
95+ """Get the call timeout."""
96+ return timedelta (seconds = self ._call_timeout_seconds )
97+
98+ @property
99+ def stream_timeout (self ) -> timedelta :
100+ """Get the stream timeout."""
101+ return timedelta (seconds = self ._stream_timeout_seconds )
102+
85103 @property
86104 def stub (self ) -> dispatch_pb2_grpc .MicrogridDispatchServiceAsyncStub :
87105 """The stub for the service."""
@@ -177,7 +195,9 @@ def to_interval(
177195 while True :
178196 response = await cast (
179197 Awaitable [ListMicrogridDispatchesResponse ],
180- self .stub .ListMicrogridDispatches (request , metadata = self ._metadata ),
198+ self .stub .ListMicrogridDispatches (
199+ request , metadata = self ._metadata , timeout = self ._call_timeout_seconds
200+ ),
181201 )
182202
183203 yield (Dispatch .from_protobuf (dispatch ) for dispatch in response .dispatches )
@@ -234,7 +254,9 @@ def _get_stream(
234254 stream_method = lambda : cast (
235255 AsyncIterator [StreamMicrogridDispatchesResponse ],
236256 self .stub .StreamMicrogridDispatches (
237- request , metadata = self ._metadata
257+ request ,
258+ metadata = self ._metadata ,
259+ timeout = self ._stream_timeout_seconds ,
238260 ),
239261 ),
240262 transform = DispatchEvent .from_protobuf ,
@@ -303,7 +325,9 @@ async def create( # pylint: disable=too-many-positional-arguments
303325 response = await cast (
304326 Awaitable [CreateMicrogridDispatchResponse ],
305327 self .stub .CreateMicrogridDispatch (
306- request .to_protobuf (), metadata = self ._metadata
328+ request .to_protobuf (),
329+ metadata = self ._metadata ,
330+ timeout = self ._call_timeout_seconds ,
307331 ),
308332 )
309333
@@ -394,7 +418,9 @@ async def update(
394418
395419 response = await cast (
396420 Awaitable [UpdateMicrogridDispatchResponse ],
397- self .stub .UpdateMicrogridDispatch (msg , metadata = self ._metadata ),
421+ self .stub .UpdateMicrogridDispatch (
422+ msg , metadata = self ._metadata , timeout = self ._call_timeout_seconds
423+ ),
398424 )
399425
400426 return Dispatch .from_protobuf (response .dispatch )
@@ -414,7 +440,9 @@ async def get(self, *, microgrid_id: int, dispatch_id: int) -> Dispatch:
414440 )
415441 response = await cast (
416442 Awaitable [GetMicrogridDispatchResponse ],
417- self .stub .GetMicrogridDispatch (request , metadata = self ._metadata ),
443+ self .stub .GetMicrogridDispatch (
444+ request , metadata = self ._metadata , timeout = self ._call_timeout_seconds
445+ ),
418446 )
419447 return Dispatch .from_protobuf (response .dispatch )
420448
@@ -430,5 +458,7 @@ async def delete(self, *, microgrid_id: int, dispatch_id: int) -> None:
430458 )
431459 await cast (
432460 Awaitable [None ],
433- self .stub .DeleteMicrogridDispatch (request , metadata = self ._metadata ),
461+ self .stub .DeleteMicrogridDispatch (
462+ request , metadata = self ._metadata , timeout = self ._call_timeout_seconds
463+ ),
434464 )
0 commit comments