1+ import json
2+
13import grpc
24
35from dapr .clients .exceptions import StreamInactiveError
@@ -34,32 +36,45 @@ def __init__(self, stub, pubsub_name, topic, metadata=None, dead_letter_topic=No
3436 self ._stream_lock = threading .Lock () # Protects _stream_active
3537
3638 def start (self ):
37- def request_iterator ():
39+ def outgoing_request_iterator ():
40+ """
41+ Generator function to create the request iterator for the stream
42+ """
3843 try :
3944 # Send InitialRequest needed to establish the stream
4045 initial_request = api_v1 .SubscribeTopicEventsRequestAlpha1 (
4146 initial_request = api_v1 .SubscribeTopicEventsRequestInitialAlpha1 (
42- pubsub_name = self .pubsub_name , topic = self .topic , metadata = self .metadata or {},
43- dead_letter_topic = self .dead_letter_topic or '' ))
47+ pubsub_name = self .pubsub_name ,
48+ topic = self .topic ,
49+ metadata = self .metadata or {},
50+ dead_letter_topic = self .dead_letter_topic or '' ,
51+ )
52+ )
4453 yield initial_request
4554
55+ # Start sending back acknowledgement messages from the send queue
4656 while self ._is_stream_active ():
4757 try :
48- yield self ._send_queue .get () # TODO Should I add a timeout?
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
61+ if not self ._is_stream_active ():
62+ break
63+ yield response
4964 except queue .Empty :
5065 continue
5166 except Exception as e :
52- raise Exception (f" Error in request iterator: { e } " )
67+ raise Exception (f' Error in request iterator: { e } ' )
5368
5469 # Create the bidirectional stream
55- self ._stream = self ._stub .SubscribeTopicEventsAlpha1 (request_iterator ())
70+ self ._stream = self ._stub .SubscribeTopicEventsAlpha1 (outgoing_request_iterator ())
5671 self ._set_stream_active ()
5772
5873 # Start a thread to handle incoming messages
59- self ._response_thread = threading .Thread (target = self ._handle_responses , daemon = True )
74+ self ._response_thread = threading .Thread (target = self ._handle_incoming_messages , daemon = True )
6075 self ._response_thread .start ()
6176
62- def _handle_responses (self ):
77+ def _handle_incoming_messages (self ):
6378 try :
6479 # The first message dapr sends on the stream is for signalling only, so discard it
6580 next (self ._stream )
@@ -72,30 +87,31 @@ def _handle_responses(self):
7287 break
7388 except grpc .RpcError as e :
7489 if e .code () != grpc .StatusCode .CANCELLED :
75- print (f" gRPC error in stream: { e .details ()} , Status Code: { e .code ()} " )
90+ print (f' gRPC error in stream: { e .details ()} , Status Code: { e .code ()} ' )
7691 except Exception as e :
77- raise Exception (f" Error while handling responses: { e } " )
92+ raise Exception (f' Error while handling responses: { e } ' )
7893 finally :
7994 self ._set_stream_inactive ()
8095
81- def next_message (self , timeout = 1 ):
82- """
83- Gets the next message from the receive queue
84- @param timeout: Timeout in seconds
85- @return: The next message
86- """
87- return self . read_message_from_queue ( self . _receive_queue , timeout = timeout )
96+ def next_message (self , timeout = None ):
97+ msg = self . read_message_from_queue ( self . _receive_queue , timeout = timeout )
98+
99+ if msg is None :
100+ return None
101+
102+ return SubscriptionMessage ( msg )
88103
89104 def _respond (self , message , status ):
90105 try :
91106 status = appcallback_v1 .TopicEventResponse (status = status .value )
92- response = api_v1 .SubscribeTopicEventsRequestProcessedAlpha1 (id = message .id ,
93- status = status )
107+ response = api_v1 .SubscribeTopicEventsRequestProcessedAlpha1 (
108+ id = message .id (), status = status
109+ )
94110 msg = api_v1 .SubscribeTopicEventsRequestAlpha1 (event_processed = response )
95111
96112 self .send_message_to_queue (self ._send_queue , msg )
97113 except Exception as e :
98- print (f" Exception in send_message: { e } " )
114+ print (f' Exception in send_message: { e } ' )
99115
100116 def respond_success (self , message ):
101117 self ._respond (message , TopicEventResponse ('success' ).status )
@@ -108,12 +124,12 @@ def respond_drop(self, message):
108124
109125 def send_message_to_queue (self , q , message ):
110126 if not self ._is_stream_active ():
111- raise StreamInactiveError (" Stream is not active" )
127+ raise StreamInactiveError (' Stream is not active' )
112128 q .put (message )
113129
114- def read_message_from_queue (self , q , timeout ):
130+ def read_message_from_queue (self , q , timeout = None ):
115131 if not self ._is_stream_active ():
116- raise StreamInactiveError (" Stream is not active" )
132+ raise StreamInactiveError (' Stream is not active' )
117133 try :
118134 return q .get (timeout = timeout )
119135 except queue .Empty :
@@ -143,12 +159,84 @@ def close(self):
143159 self ._stream .cancel ()
144160 except grpc .RpcError as e :
145161 if e .code () != grpc .StatusCode .CANCELLED :
146- raise Exception (f" Error while closing stream: { e } " )
162+ raise Exception (f' Error while closing stream: { e } ' )
147163 except Exception as e :
148- raise Exception (f" Error while closing stream: { e } " )
164+ raise Exception (f' Error while closing stream: { e } ' )
149165
150166 # Join the response-handling thread to ensure it has finished
151167 if self ._response_thread :
152168 self ._response_thread .join ()
153169 self ._response_thread = None
154170
171+
172+ class SubscriptionMessage :
173+ def __init__ (self , msg ):
174+ self ._id = msg .id
175+ self ._source = msg .source
176+ self ._type = msg .type
177+ self ._spec_version = msg .spec_version
178+ self ._data_content_type = msg .data_content_type
179+ self ._topic = msg .topic
180+ self ._pubsub_name = msg .pubsub_name
181+ self ._raw_data = msg .data
182+ self ._extensions = msg .extensions
183+ self ._data = None
184+
185+ # Parse the content based on its media type
186+ if self ._raw_data and len (self ._raw_data ) > 0 :
187+ self ._parse_data_content ()
188+
189+ def id (self ):
190+ return self ._id
191+
192+ def source (self ):
193+ return self ._source
194+
195+ def type (self ):
196+ return self ._type
197+
198+ def spec_version (self ):
199+ return self ._spec_version
200+
201+ def data_content_type (self ):
202+ return self ._data_content_type
203+
204+ def topic (self ):
205+ return self ._topic
206+
207+ def pubsub_name (self ):
208+ return self ._pubsub_name
209+
210+ def raw_data (self ):
211+ return self ._raw_data
212+
213+ def extensions (self ):
214+ return self ._extensions
215+
216+ def data (self ):
217+ return self ._data
218+
219+ def _parse_data_content (self ):
220+ try :
221+ if self ._data_content_type == 'application/json' :
222+ try :
223+ self ._data = json .loads (self ._raw_data )
224+ except json .JSONDecodeError :
225+ pass # If JSON parsing fails, keep `data` as None
226+ elif self ._data_content_type == 'text/plain' :
227+ # Assume UTF-8 encoding
228+ try :
229+ self ._data = self ._raw_data .decode ('utf-8' )
230+ except UnicodeDecodeError :
231+ pass
232+ elif self ._data_content_type .startswith (
233+ 'application/'
234+ ) and self ._data_content_type .endswith ('+json' ):
235+ # Handle custom JSON-based media types (e.g., application/vnd.api+json)
236+ try :
237+ self ._data = json .loads (self ._raw_data )
238+ except json .JSONDecodeError :
239+ pass # If JSON parsing fails, keep `data` as None
240+ except Exception as e :
241+ # Log or handle any unexpected exceptions
242+ print (f'Error parsing media type: { e } ' )
0 commit comments