works

Signed-off-by: Elena Kolevska <[email protected]>

Author: elena-kolevska

dapr/clients/exceptions.py

@@ -27,22 +27,15 @@
 class DaprInternalError(Exception):
     """DaprInternalError encapsulates all Dapr exceptions"""
 
-    def __init__(
-        self,
-        message: Optional[str],
-        error_code: Optional[str] = ERROR_CODE_UNKNOWN,
-        raw_response_bytes: Optional[bytes] = None,
-    ):
+    def __init__(self, message: Optional[str], error_code: Optional[str] = ERROR_CODE_UNKNOWN,
+            raw_response_bytes: Optional[bytes] = None, ):
         self._message = message
         self._error_code = error_code
         self._raw_response_bytes = raw_response_bytes
 
     def as_dict(self):
-        return {
-            'message': self._message,
-            'errorCode': self._error_code,
-            'raw_response_bytes': self._raw_response_bytes,
-        }
+        return {'message': self._message, 'errorCode': self._error_code,
+            'raw_response_bytes': self._raw_response_bytes, }
 
 
 class StatusDetails:
@@ -119,16 +112,16 @@ def get_grpc_status(self):
         return self._grpc_status
 
     def json(self):
-        error_details = {
-            'status_code': self.code().name,
-            'message': self.details(),
-            'error_code': self.error_code(),
-            'details': self._details.as_dict(),
-        }
+        error_details = {'status_code': self.code().name, 'message': self.details(),
+            'error_code': self.error_code(), 'details': self._details.as_dict(), }
         return json.dumps(error_details)
 
 
 def serialize_status_detail(status_detail):
     if not status_detail:
         return None
     return MessageToDict(status_detail, preserving_proto_field_name=True)
+
+
+class StreamInactiveError(Exception):
+    pass

dapr/clients/grpc/subscription.py

@@ -1,4 +1,7 @@
 import grpc
+
+from dapr.clients.exceptions import StreamInactiveError
+from dapr.clients.grpc._response import TopicEventResponse
 from dapr.proto import api_v1, appcallback_v1
 import queue
 import threading
@@ -24,9 +27,11 @@ def __init__(self, stub, pubsub_name, topic, metadata=None, dead_letter_topic=No
         self.metadata = metadata or {}
         self.dead_letter_topic = dead_letter_topic or ''
         self._stream = None
+        self._response_thread = None
         self._send_queue = queue.Queue()
         self._receive_queue = queue.Queue()
         self._stream_active = False
+        self._stream_lock = threading.Lock()  # Protects _stream_active
 
     def start(self):
         def request_iterator():
@@ -38,65 +43,112 @@ def request_iterator():
                         dead_letter_topic=self.dead_letter_topic or ''))
                 yield initial_request
 
-                while self._stream_active:
+                while self._is_stream_active():
                     try:
-                        request = self._send_queue.get()
-                        if request is None:
-                            break
-
-                        yield request
+                        yield self._send_queue.get()  # TODO Should I add a timeout?
                     except queue.Empty:
                         continue
             except Exception as e:
-                print(f"Exception in request_iterator: {e}")
-                raise e
+                raise Exception(f"Error in request iterator: {e}")
 
         # Create the bidirectional stream
         self._stream = self._stub.SubscribeTopicEventsAlpha1(request_iterator())
-        self._stream_active = True
+        self._set_stream_active()
 
         # Start a thread to handle incoming messages
-        threading.Thread(target=self._handle_responses, daemon=True).start()
+        self._response_thread = threading.Thread(target=self._handle_responses, daemon=True)
+        self._response_thread.start()
 
     def _handle_responses(self):
         try:
             # The first message dapr sends on the stream is for signalling only, so discard it
             next(self._stream)
 
-            for msg in self._stream:
-                print(f"Received message from dapr on stream: {msg.event_message.id}") # SubscribeTopicEventsResponseAlpha1
-                self._receive_queue.put(msg.event_message)
+            # Read messages from the stream and put them in the receive queue
+            for message in self._stream:
+                if self._is_stream_active():
+                    self._receive_queue.put(message.event_message)
+                else:
+                    break
         except grpc.RpcError as e:
-            print(f"gRPC error in stream: {e}")
+            if e.code() != grpc.StatusCode.CANCELLED:
+                print(f"gRPC error in stream: {e.details()}, Status Code: {e.code()}")
         except Exception as e:
-            print(f"Unexpected error in stream: {e}")
+            raise Exception(f"Error while handling responses: {e}")
         finally:
-            self._stream_active = False
+            self._set_stream_inactive()
 
-    def next_message(self, timeout=None):
-        print("in next_message")
-        try:
-            return self._receive_queue.get(timeout=timeout)
-        except queue.Empty as e :
-            print("queue empty", e)
-            return None
-        except Exception as e:
-            print(f"Exception in next_message: {e}")
-            return None
+    def next_message(self, timeout=1):
+        """
+        Gets the next message from the receive queue
+        @param timeout: Timeout in seconds
+        @return: The next message
+        """
+        return self.read_message_from_queue(self._receive_queue, timeout=timeout)
 
-    def respond(self, message, status):
+    def _respond(self, message, status):
         try:
             status = appcallback_v1.TopicEventResponse(status=status.value)
             response = api_v1.SubscribeTopicEventsRequestProcessedAlpha1(id=message.id,
                                                                          status=status)
             msg = api_v1.SubscribeTopicEventsRequestAlpha1(event_processed=response)
 
-            self._send_queue.put(msg)
+            self.send_message_to_queue(self._send_queue, msg)
         except Exception as e:
             print(f"Exception in send_message: {e}")
 
+    def respond_success(self, message):
+        self._respond(message, TopicEventResponse('success').status)
+
+    def respond_retry(self, message):
+        self._respond(message, TopicEventResponse('retry').status)
+
+    def respond_drop(self, message):
+        self._respond(message, TopicEventResponse('drop').status)
+
+    def send_message_to_queue(self, q, message):
+        if not self._is_stream_active():
+            raise StreamInactiveError("Stream is not active")
+        q.put(message)
+
+    def read_message_from_queue(self, q, timeout):
+        if not self._is_stream_active():
+            raise StreamInactiveError("Stream is not active")
+        try:
+            return q.get(timeout=timeout)
+        except queue.Empty:
+            return None
+
+    def _set_stream_active(self):
+        with self._stream_lock:
+            self._stream_active = True
+
+    def _set_stream_inactive(self):
+        with self._stream_lock:
+            self._stream_active = False
+
+    def _is_stream_active(self):
+        with self._stream_lock:
+            return self._stream_active
+
     def close(self):
-        self._stream_active = False
-        self._send_queue.put(None)
+        if not self._is_stream_active():
+            return
+
+        self._set_stream_inactive()
+
+        # Cancel the stream
         if self._stream:
-            self._stream.cancel()
+            try:
+                self._stream.cancel()
+            except grpc.RpcError as e:
+                if e.code() != grpc.StatusCode.CANCELLED:
+                    raise Exception(f"Error while closing stream: {e}")
+            except Exception as e:
+                raise Exception(f"Error while closing stream: {e}")
+
+        # Join the response-handling thread to ensure it has finished
+        if self._response_thread:
+            self._response_thread.join()
+            self._response_thread = None
+

examples/pubsub_streaming/subscriber.py

@@ -5,8 +5,8 @@
 
 def process_message(message):
     # Process the message here
-    print(f"Processing message: {message.data}")
-    return  TopicEventResponse('success').status
+    print(f"Processing message: {message}")
+    return "success"
 
 
 def main():
@@ -15,26 +15,25 @@ def main():
         subscription = client.subscribe(pubsub_name="pubsub", topic="TOPIC_A", dead_letter_topic="TOPIC_A_DEAD")
 
         try:
-            while True:
+            for i in range(5):
                 try:
-                    try:
-                        message = subscription.next_message(timeout=5)
-                    except Exception as e:
-                        print(f"An error occurred: {e}")
-
+                    message = subscription.next_message(timeout=0.1)
                     if message is None:
                         print("No message received within timeout period.")
                         continue
 
-                    print(f"Received message with ID: {message.id}")
-
                     # Process the message
-                    try:
-                        subscription.respond(message, process_message(message))
-                    except Exception as e:
-                        print(f"An error occurred while sending the message: {e}")
-                except KeyboardInterrupt:
-                    print("Received interrupt, shutting down...")
+                    response_status = process_message(message)
+
+                    if response_status == "success":
+                        subscription.respond_success(message)
+                    elif response_status == "retry":
+                        subscription.respond_retry(message)
+                    elif response_status == "drop":
+                        subscription.respond_drop(message)
+
+                except Exception as e:
+                    print(f"Error getting message: {e}")
                     break
 
         finally:

tests/clients/fake_dapr_server.py

@@ -7,7 +7,7 @@
 from grpc_status import rpc_status
 
 from dapr.clients.grpc._helpers import to_bytes
-from dapr.proto import api_service_v1, common_v1, api_v1
+from dapr.proto import api_service_v1, common_v1, api_v1, appcallback_v1
 from dapr.proto.common.v1.common_pb2 import ConfigurationItem
 from dapr.clients.grpc._response import WorkflowRuntimeStatus
 from dapr.proto.runtime.v1.dapr_pb2 import (
@@ -177,6 +177,14 @@ def PublishEvent(self, request, context):
         context.set_trailing_metadata(trailers)
         return empty_pb2.Empty()
 
+    def SubscribeTopicEventsAlpha1(self, request_iterator, context):
+        yield api_v1.SubscribeTopicEventsResponseAlpha1(
+            initial_response=api_v1.SubscribeTopicEventsResponseInitialAlpha1())
+        yield api_v1.SubscribeTopicEventsResponseAlpha1(
+            event_message=appcallback_v1.TopicEventRequest(id='123', topic="TOPIC_A", data=b'hello1'))
+        yield api_v1.SubscribeTopicEventsResponseAlpha1(
+            event_message=appcallback_v1.TopicEventRequest(id='456', topic="TOPIC_A", data=b'hello2'))
+
     def SaveState(self, request, context):
         self.check_for_exception(context)
 

tests/clients/test_dapr_grpc_client.py

@@ -24,7 +24,7 @@
 
 from google.rpc import status_pb2, code_pb2
 
-from dapr.clients.exceptions import DaprGrpcError
+from dapr.clients.exceptions import DaprGrpcError, StreamInactiveError
 from dapr.clients.grpc.client import DaprGrpcClient
 from dapr.clients import DaprClient
 from dapr.proto import common_v1
@@ -34,13 +34,9 @@
 from dapr.clients.grpc._request import TransactionalStateOperation
 from dapr.clients.grpc._state import StateOptions, Consistency, Concurrency, StateItem
 from dapr.clients.grpc._crypto import EncryptOptions, DecryptOptions
-from dapr.clients.grpc._response import (
-    ConfigurationItem,
-    ConfigurationResponse,
-    ConfigurationWatcher,
-    UnlockResponseStatus,
-    WorkflowRuntimeStatus,
-)
+from dapr.clients.grpc._response import (ConfigurationItem, ConfigurationResponse,
+                                         ConfigurationWatcher, UnlockResponseStatus,
+                                         WorkflowRuntimeStatus, TopicEventResponse, )
 
 
 class DaprGrpcClientTests(unittest.TestCase):
@@ -262,6 +258,35 @@ def test_publish_error(self):
                 data=111,
             )
 
+    def test_subscribe_topic(self):
+        dapr = DaprGrpcClient(f'{self.scheme}localhost:{self.grpc_port}')
+        subscription = dapr.subscribe(pubsub_name='pubsub', topic='example')
+
+        # First message
+        message1 = subscription.next_message(timeout=5)
+        subscription.respond_success(message1)
+
+        self.assertEqual('123', message1.id)
+        self.assertEqual(b'hello1', message1.data)
+        self.assertEqual('TOPIC_A', message1.topic)
+
+        # Second message
+        message2 = subscription.next_message(timeout=5)
+        subscription.respond_success(message2)
+
+        self.assertEqual('456', message2.id)
+        self.assertEqual(b'hello2', message2.data)
+        self.assertEqual('TOPIC_A', message2.topic)
+
+    def test_subscribe_topic_early_close(self):
+        dapr = DaprGrpcClient(f'{self.scheme}localhost:{self.grpc_port}')
+        subscription = dapr.subscribe(pubsub_name='pubsub', topic='example')
+        subscription.close()
+
+        with self.assertRaises(StreamInactiveError):
+            subscription.next_message(timeout=5)
+
+
     @patch.object(settings, 'DAPR_API_TOKEN', 'test-token')
     def test_dapr_api_token_insertion(self):
         dapr = DaprGrpcClient(f'{self.scheme}localhost:{self.grpc_port}')