49
49
# pylint: enable=no-name-in-module
50
50
DEFAULT_DISPATCH_PORT = 443
51
51
52
+ DEFAULT_GRPC_CALL_TIMEOUT = timedelta (seconds = 60 )
53
+ """Default gRPC call timeout."""
54
+
55
+ DEFAULT_GRPC_STREAM_TIMEOUT = timedelta (seconds = 5 )
56
+ """Default gRPC stream timeout."""
57
+
52
58
53
59
class DispatchApiClient (BaseApiClient [dispatch_pb2_grpc .MicrogridDispatchServiceStub ]):
54
60
"""Dispatch API client."""
@@ -59,13 +65,17 @@ def __init__(
59
65
server_url : str ,
60
66
key : str ,
61
67
connect : bool = True ,
68
+ call_timeout : timedelta | None = None ,
69
+ stream_timeout : timedelta | None = None ,
62
70
) -> None :
63
71
"""Initialize the client.
64
72
65
73
Args:
66
74
server_url: The URL of the server to connect to.
67
75
key: API key to use for authentication.
68
76
connect: Whether to connect to the service immediately.
77
+ call_timeout: Timeout for gRPC calls, default is 60 seconds.
78
+ stream_timeout: Timeout for gRPC streams, default is 5 minutes.
69
79
"""
70
80
super ().__init__ (
71
81
server_url ,
@@ -82,6 +92,11 @@ def __init__(
82
92
] = {}
83
93
"""A dictionary of streamers, keyed by microgrid_id."""
84
94
95
+ self ._call_timeout = (call_timeout or DEFAULT_GRPC_CALL_TIMEOUT ).total_seconds ()
96
+ self ._stream_timeout = (
97
+ stream_timeout or DEFAULT_GRPC_STREAM_TIMEOUT
98
+ ).total_seconds ()
99
+
85
100
@property
86
101
def stub (self ) -> dispatch_pb2_grpc .MicrogridDispatchServiceAsyncStub :
87
102
"""The stub for the service."""
@@ -177,7 +192,9 @@ def to_interval(
177
192
while True :
178
193
response = await cast (
179
194
Awaitable [ListMicrogridDispatchesResponse ],
180
- self .stub .ListMicrogridDispatches (request , metadata = self ._metadata ),
195
+ self .stub .ListMicrogridDispatches (
196
+ request , metadata = self ._metadata , timeout = self ._call_timeout
197
+ ),
181
198
)
182
199
183
200
yield (Dispatch .from_protobuf (dispatch ) for dispatch in response .dispatches )
@@ -234,7 +251,9 @@ def _get_stream(
234
251
stream_method = lambda : cast (
235
252
AsyncIterator [StreamMicrogridDispatchesResponse ],
236
253
self .stub .StreamMicrogridDispatches (
237
- request , metadata = self ._metadata
254
+ request ,
255
+ metadata = self ._metadata ,
256
+ timeout = self ._stream_timeout ,
238
257
),
239
258
),
240
259
transform = DispatchEvent .from_protobuf ,
@@ -303,7 +322,9 @@ async def create( # pylint: disable=too-many-positional-arguments
303
322
response = await cast (
304
323
Awaitable [CreateMicrogridDispatchResponse ],
305
324
self .stub .CreateMicrogridDispatch (
306
- request .to_protobuf (), metadata = self ._metadata
325
+ request .to_protobuf (),
326
+ metadata = self ._metadata ,
327
+ timeout = self ._call_timeout ,
307
328
),
308
329
)
309
330
@@ -394,7 +415,9 @@ async def update(
394
415
395
416
response = await cast (
396
417
Awaitable [UpdateMicrogridDispatchResponse ],
397
- self .stub .UpdateMicrogridDispatch (msg , metadata = self ._metadata ),
418
+ self .stub .UpdateMicrogridDispatch (
419
+ msg , metadata = self ._metadata , timeout = self ._call_timeout
420
+ ),
398
421
)
399
422
400
423
return Dispatch .from_protobuf (response .dispatch )
@@ -414,7 +437,9 @@ async def get(self, *, microgrid_id: int, dispatch_id: int) -> Dispatch:
414
437
)
415
438
response = await cast (
416
439
Awaitable [GetMicrogridDispatchResponse ],
417
- self .stub .GetMicrogridDispatch (request , metadata = self ._metadata ),
440
+ self .stub .GetMicrogridDispatch (
441
+ request , metadata = self ._metadata , timeout = self ._call_timeout
442
+ ),
418
443
)
419
444
return Dispatch .from_protobuf (response .dispatch )
420
445
@@ -430,5 +455,7 @@ async def delete(self, *, microgrid_id: int, dispatch_id: int) -> None:
430
455
)
431
456
await cast (
432
457
Awaitable [None ],
433
- self .stub .DeleteMicrogridDispatch (request , metadata = self ._metadata ),
458
+ self .stub .DeleteMicrogridDispatch (
459
+ request , metadata = self ._metadata , timeout = self ._call_timeout
460
+ ),
434
461
)
0 commit comments