1
1
import json
2
2
3
- import grpc
3
+ from grpc import StreamStreamMultiCallable , RpcError , StatusCode # type: ignore
4
4
5
5
from dapr .clients .exceptions import StreamInactiveError
6
6
from dapr .clients .grpc ._response import TopicEventResponse
7
7
from dapr .proto import api_v1 , appcallback_v1
8
8
import queue
9
9
import threading
10
+ from typing import Optional
10
11
11
12
12
13
def success ():
@@ -28,11 +29,11 @@ def __init__(self, stub, pubsub_name, topic, metadata=None, dead_letter_topic=No
28
29
self .topic = topic
29
30
self .metadata = metadata or {}
30
31
self .dead_letter_topic = dead_letter_topic or ''
31
- self ._stream = None
32
- self ._response_thread = None
33
- self ._send_queue = queue .Queue ()
34
- self ._receive_queue = queue .Queue ()
35
- self ._stream_active = False
32
+ self ._stream : Optional [ StreamStreamMultiCallable ] = None # Type annotation for gRPC stream
33
+ self ._response_thread : Optional [ threading . Thread ] = None # Type for thread
34
+ self ._send_queue : queue . Queue = queue .Queue () # Type annotation for send queue
35
+ self ._receive_queue : queue . Queue = queue .Queue () # Type annotation for receive queue
36
+ self ._stream_active : bool = False
36
37
self ._stream_lock = threading .Lock () # Protects _stream_active
37
38
38
39
def start (self ):
@@ -55,9 +56,8 @@ def outgoing_request_iterator():
55
56
# Start sending back acknowledgement messages from the send queue
56
57
while self ._is_stream_active ():
57
58
try :
58
- response = self ._send_queue .get ()
59
- # The above blocks until a message is available or the stream is closed
60
- # so that's why we need to check again if the stream is still active
59
+ response = self ._send_queue .get (timeout = 1 )
60
+ # Check again if the stream is still active
61
61
if not self ._is_stream_active ():
62
62
break
63
63
yield response
@@ -76,17 +76,19 @@ def outgoing_request_iterator():
76
76
77
77
def _handle_incoming_messages (self ):
78
78
try :
79
- # The first message dapr sends on the stream is for signalling only, so discard it
80
- next (self ._stream )
81
-
82
- # Read messages from the stream and put them in the receive queue
83
- for message in self ._stream :
84
- if self ._is_stream_active ():
85
- self ._receive_queue .put (message .event_message )
86
- else :
87
- break
88
- except grpc .RpcError as e :
89
- if e .code () != grpc .StatusCode .CANCELLED :
79
+ # Check if the stream is not None
80
+ if self ._stream is not None :
81
+ # The first message dapr sends on the stream is for signalling only, so discard it
82
+ next (self ._stream )
83
+
84
+ # Read messages from the stream and put them in the receive queue
85
+ for message in self ._stream :
86
+ if self ._is_stream_active ():
87
+ self ._receive_queue .put (message .event_message )
88
+ else :
89
+ break
90
+ except RpcError as e :
91
+ if e .code () != StatusCode .CANCELLED :
90
92
print (f'gRPC error in stream: { e .details ()} , Status Code: { e .code ()} ' )
91
93
except Exception as e :
92
94
raise Exception (f'Error while handling responses: { e } ' )
@@ -157,8 +159,8 @@ def close(self):
157
159
if self ._stream :
158
160
try :
159
161
self ._stream .cancel ()
160
- except grpc . RpcError as e :
161
- if e .code () != grpc . StatusCode .CANCELLED :
162
+ except RpcError as e :
163
+ if e .code () != StatusCode .CANCELLED :
162
164
raise Exception (f'Error while closing stream: { e } ' )
163
165
except Exception as e :
164
166
raise Exception (f'Error while closing stream: { e } ' )
0 commit comments