improving ack implementation

Author: DanielePalaia

examples/getting_started/main.py

@@ -2,10 +2,10 @@
 from rabbitmq_amqp_python_client import (
     BindingSpecification,
     Connection,
-    Delivery,
     Event,
     ExchangeSpecification,
     Message,
+    MessageAck,
     MessagingHandler,
     QuorumQueueSpecification,
     exchange_address,
@@ -20,18 +20,19 @@ def __init__(self):
         self._count = 0
 
     def on_message(self, event: Event):
-        print("received message: " + event.message.body)
+        print("received message: " + str(event.message.annotations))
 
-        dlv = event.delivery
+        # annotations = {}
+        # annotations[symbol('x-opt-string')] = 'x-test1'
+        # MessageAck.requeue_with_annotations(event, annotations)
 
-        dlv.update(Delivery.ACCEPTED)
-        dlv.settle()
+        MessageAck.accept(event)
 
         print("count " + str(self._count))
 
         self._count = self._count + 1
 
-        if self._count == 100000:
+        if self._count == 100:
             print("closing receiver")
             event.receiver.close()
             event.connection.close()
@@ -42,9 +43,6 @@ def on_connection_closed(self, event: Event):
     def on_link_closed(self, event: Event) -> None:
         print("link closed")
 
-    def on_rejected(self, event: Event) -> None:
-        print("rejected")
-
 
 def create_connection() -> Connection:
     connection = Connection("amqp://guest:guest@localhost:5672/")
@@ -58,7 +56,7 @@ def main() -> None:
     exchange_name = "test-exchange"
     queue_name = "example-queue"
     routing_key = "routing-key"
-    messages_to_publish = 100000
+    messages_to_publish = 100
 
     print("connection to amqp server")
     connection = create_connection()
@@ -121,14 +119,14 @@ def main() -> None:
     management.unbind(bind_name)
 
     print("delete queue")
-    management.delete_queue(queue_name)
+    # management.delete_queue(queue_name)
 
     print("delete exchange")
     management.delete_exchange(exchange_name)
 
     print("closing connections")
     management.close()
-    consumer.close()
+    # consumer.close()
     print("after management closing")
     connection.close()
     print("after connection closing")

rabbitmq_amqp_python_client/__init__.py

@@ -9,7 +9,9 @@
     ExchangeSpecification,
 )
 from .management import Management
+from .message_ack import MessageAck
 from .publisher import Publisher
+from .qpid.proton._data import symbol  # noqa: E402
 from .qpid.proton._delivery import Delivery
 from .qpid.proton._events import Event
 from .qpid.proton._message import Message
@@ -46,4 +48,6 @@
     "MessagingHandler",
     "Event",
     "Delivery",
+    "MessageAck",
+    "symbol",
 ]

rabbitmq_amqp_python_client/message_ack.py

@@ -0,0 +1,52 @@
+from typing import Dict
+
+from .qpid.proton._data import PythonAMQPData
+from .qpid.proton._delivery import Delivery
+from .qpid.proton._events import Event
+
+
+class MessageAck:
+
+    @staticmethod
+    def accept(event: Event) -> None:
+        dlv = event.delivery
+        dlv.update(Delivery.ACCEPTED)
+        dlv.settle()
+
+    @staticmethod
+    def discard(event: Event) -> None:
+        dlv = event.delivery
+        dlv.update(Delivery.REJECTED)
+        dlv.settle()
+
+    @staticmethod
+    def discard_with_annotations(
+        event: Event, annotations: Dict[str, "PythonAMQPData"]
+    ) -> None:
+        dlv = event.delivery
+        dlv.local.annotations = annotations
+        dlv.local.failed = True
+        dlv.local.undeliverable = True
+
+        dlv.update(Delivery.MODIFIED)
+        dlv.settle()
+
+    @staticmethod
+    def requeue(event: Event) -> None:
+        dlv = event.delivery
+        dlv.update(Delivery.RELEASED)
+        dlv.settle()
+
+    @staticmethod
+    def requeue_with_annotations(
+        event: Event, annotations: Dict[str, "PythonAMQPData"]
+    ) -> None:
+        dlv = event.delivery
+        if dlv.local.annotations is None:
+            print("is none")
+        dlv.local.annotations = annotations
+        dlv.local.failed = False
+        dlv.local.undeliverable = False
+
+        dlv.update(Delivery.MODIFIED)
+        dlv.settle()

tests/conftest.py

@@ -2,8 +2,8 @@
 
 from rabbitmq_amqp_python_client import (
     Connection,
-    Delivery,
     Event,
+    MessageAck,
     MessagingHandler,
 )
 
@@ -44,10 +44,11 @@ def __str__(self) -> str:
 class MyMessageHandlerAccept(MessagingHandler):
 
     def __init__(self):
-        super().__init__(auto_accept=True, auto_settle=True)
+        super().__init__(auto_accept=False, auto_settle=False)
         self._received = 0
 
     def on_message(self, event: Event):
+        MessageAck.accept(event)
         self._received = self._received + 1
         if self._received == 1000:
             event.connection.close()
@@ -77,16 +78,14 @@ def on_rejected(self, event: Event) -> None:
         print("rejected")
 
 
-class MyMessageHandlerReject(MessagingHandler):
+class MyMessageHandlerDiscard(MessagingHandler):
 
     def __init__(self):
-        super().__init__(auto_accept=False, auto_settle=True)
+        super().__init__(auto_accept=False, auto_settle=False)
         self._received = 0
 
     def on_message(self, event: Event):
-        dlv = event.delivery
-        dlv.update(Delivery.REJECTED)
-        dlv.settle()
+        MessageAck.discard(event)
         self._received = self._received + 1
         if self._received == 1000:
             event.connection.close()
@@ -96,13 +95,11 @@ def on_message(self, event: Event):
 class MyMessageHandlerRequeue(MessagingHandler):
 
     def __init__(self):
-        super().__init__(auto_accept=False, auto_settle=True)
+        super().__init__(auto_accept=False, auto_settle=False)
         self._received = 0
 
     def on_message(self, event: Event):
-        dlv = event.delivery
-        dlv.update(Delivery.RELEASED)
-        dlv.settle()
+        MessageAck.requeue(event)
         self._received = self._received + 1
         if self._received == 1000:
             event.connection.close()

tests/test_consumer.py

@@ -8,8 +8,8 @@
 from .conftest import (
     ConsumerTestException,
     MyMessageHandlerAccept,
+    MyMessageHandlerDiscard,
     MyMessageHandlerNoack,
-    MyMessageHandlerReject,
     MyMessageHandlerRequeue,
     create_connection,
 )
@@ -132,10 +132,10 @@ def test_consumer_async_queue_noack(connection: Connection) -> None:
     assert message_count > 0
 
 
-def test_consumer_async_queue_with_reject(connection: Connection) -> None:
+def test_consumer_async_queue_with_discard(connection: Connection) -> None:
     messages_to_send = 1000
 
-    queue_name = "test-queue_async_reject"
+    queue_name = "test-queue_async_discard"
 
     management = connection.management()
 
@@ -155,7 +155,7 @@ def test_consumer_async_queue_with_reject(connection: Connection) -> None:
     connection_consumer = create_connection()
 
     consumer = connection_consumer.consumer(
-        addr_queue, handler=MyMessageHandlerReject()
+        addr_queue, handler=MyMessageHandlerDiscard()
     )
 
     try: