53
53
class DispatchApiClient (BaseApiClient [dispatch_pb2_grpc .MicrogridDispatchServiceStub ]):
54
54
"""Dispatch API client."""
55
55
56
+ # pylint: disable-next=too-many-arguments
56
57
def __init__ (
57
58
self ,
58
59
* ,
59
60
server_url : str ,
60
61
key : str ,
61
62
connect : bool = True ,
63
+ call_timeout : timedelta = timedelta (seconds = 60 ),
64
+ stream_timeout : timedelta = timedelta (minutes = 5 ),
62
65
) -> None :
63
66
"""Initialize the client.
64
67
65
68
Args:
66
69
server_url: The URL of the server to connect to.
67
70
key: API key to use for authentication.
68
71
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.
69
74
"""
70
75
super ().__init__ (
71
76
server_url ,
@@ -82,6 +87,19 @@ def __init__(
82
87
] = {}
83
88
"""A dictionary of streamers, keyed by microgrid_id."""
84
89
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
+
85
103
@property
86
104
def stub (self ) -> dispatch_pb2_grpc .MicrogridDispatchServiceAsyncStub :
87
105
"""The stub for the service."""
@@ -177,7 +195,9 @@ def to_interval(
177
195
while True :
178
196
response = await cast (
179
197
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
+ ),
181
201
)
182
202
183
203
yield (Dispatch .from_protobuf (dispatch ) for dispatch in response .dispatches )
@@ -234,7 +254,9 @@ def _get_stream(
234
254
stream_method = lambda : cast (
235
255
AsyncIterator [StreamMicrogridDispatchesResponse ],
236
256
self .stub .StreamMicrogridDispatches (
237
- request , metadata = self ._metadata
257
+ request ,
258
+ metadata = self ._metadata ,
259
+ timeout = self ._stream_timeout_seconds ,
238
260
),
239
261
),
240
262
transform = DispatchEvent .from_protobuf ,
@@ -303,7 +325,9 @@ async def create( # pylint: disable=too-many-positional-arguments
303
325
response = await cast (
304
326
Awaitable [CreateMicrogridDispatchResponse ],
305
327
self .stub .CreateMicrogridDispatch (
306
- request .to_protobuf (), metadata = self ._metadata
328
+ request .to_protobuf (),
329
+ metadata = self ._metadata ,
330
+ timeout = self ._call_timeout_seconds ,
307
331
),
308
332
)
309
333
@@ -394,7 +418,9 @@ async def update(
394
418
395
419
response = await cast (
396
420
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
+ ),
398
424
)
399
425
400
426
return Dispatch .from_protobuf (response .dispatch )
@@ -414,7 +440,9 @@ async def get(self, *, microgrid_id: int, dispatch_id: int) -> Dispatch:
414
440
)
415
441
response = await cast (
416
442
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
+ ),
418
446
)
419
447
return Dispatch .from_protobuf (response .dispatch )
420
448
@@ -430,5 +458,7 @@ async def delete(self, *, microgrid_id: int, dispatch_id: int) -> None:
430
458
)
431
459
await cast (
432
460
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
+ ),
434
464
)
0 commit comments