adding validations for annotations

Author: DanielePalaia

examples/getting_started/main.py

@@ -3,9 +3,9 @@
 
 from rabbitmq_amqp_python_client import (
     AddressHelper,
+    AMQPMessagingHandler,
     BindingSpecification,
     Connection,
-    AMQPMessagingHandler,
     Event,
     ExchangeSpecification,
     Message,

rabbitmq_amqp_python_client/__init__.py

@@ -1,16 +1,15 @@
 from importlib import metadata
 
 from .address_helper import AddressHelper
+from .amqp_consumer_handler import AMQPMessagingHandler
 from .common import ExchangeType, QueueType
 from .connection import Connection
 from .consumer import Consumer
-from .amqp_consumer_handler import (
-    AMQPMessagingHandler,
-)
 from .entities import (
     BindingSpecification,
     ExchangeSpecification,
 )
+from .exceptions import ArgumentOutOfRangeException
 from .management import Management
 from .publisher import Publisher
 from .qpid.proton._data import symbol  # noqa: E402
@@ -52,4 +51,5 @@
     "ExchangeType",
     "AddressHelper",
     "AMQPMessagingHandler",
+    "ArgumentOutOfRangeException",
 ]

rabbitmq_amqp_python_client/amqp_consumer_handler.py

@@ -1,16 +1,15 @@
 from .delivery_context import DeliveryContext
 from .qpid.proton.handlers import MessagingHandler
 
-
-''' 
-AMQPMessagingHandler extends the QPID MessagingHandler. 
+"""
+AMQPMessagingHandler extends the QPID MessagingHandler.
 It is an helper to set the default values needed for manually accepting and settling messages.
-
-self.delivery_context is an instance of DeliveryContext, which is used to accept, reject, 
+self.delivery_context is an instance of DeliveryContext, which is used to accept, reject,
 requeue or requeue with annotations a message.
+It is not mandatory to use this class, but it is a good practice to use it.
+"""
+
 
-It is not mandatory to use this class, but it is a good practice to use it. 
-'''
 class AMQPMessagingHandler(MessagingHandler):  # type: ignore
 
     def __init__(self, auto_accept: bool = False, auto_settle: bool = True):

rabbitmq_amqp_python_client/connection.py

@@ -32,9 +32,7 @@ def management(self) -> Management:
     # closes the connection to the AMQP 1.0 server.
     def close(self) -> None:
         logger.debug("Closing connection")
-        print("closing connection")
         self._conn.close()
-        print("after closing connection")
 
     def publisher(self, destination: str) -> Publisher:
         publisher = Publisher(self._conn, destination)

rabbitmq_amqp_python_client/delivery_context.py

@@ -1,13 +1,15 @@
 from typing import Dict
 
+from .exceptions import ArgumentOutOfRangeException
 from .qpid.proton._data import PythonAMQPData
 from .qpid.proton._delivery import Delivery
 from .qpid.proton._events import Event
+from .utils import validate_annotations
 
-'''
+"""
 DeliveryContext is a class that is used to accept, reject, requeue or requeue with annotations a message.
 It is an helper to set the default values needed for manually accepting and settling messages.
-'''
+"""
 
 
 class DeliveryContext:
@@ -39,32 +41,37 @@ def discard(self, event: Event) -> None:
     or dead-letter it if it is configured.
     Application-specific annotation keys must start with the <code>x-opt-</code> prefix.
     Annotation keys the broker understands starts with <code>x-</code>, but not with <code>x-opt-
-    
-        This maps to the AMQP 1.0 
+        This maps to the AMQP 1.0
         modified{delivery-failed = true, undeliverable-here = true}</code> outcome.
          <param name="annotations"> annotations message annotations to combine with existing ones </param>
         <a
             href="https://docs.oasis-open.org/amqp/core/v1.0/os/amqp-core-messaging-v1.0-os.html#type-modified">AMQP
-            1.0 <code>modified</code> outcome</a>
-        
+                        1.0 <code>modified</code> outcome</a>
          The annotations can be used only with Quorum queues, see https://www.rabbitmq.com/docs/amqp#modified-outcome
     """
 
     def discard_with_annotations(
-            self, event: Event, annotations: Dict[str, "PythonAMQPData"]
+        self, event: Event, annotations: Dict[str, "PythonAMQPData"]
     ) -> None:
         dlv = event.delivery
         dlv.local.annotations = annotations
         dlv.local.failed = True
         dlv.local.undeliverable = True
 
+        validated = validate_annotations(annotations.keys())
+
+        if validated is False:
+            raise ArgumentOutOfRangeException(
+                "Message annotation key must start with 'x-'"
+            )
+
         dlv.update(Delivery.MODIFIED)
         dlv.settle()
 
     """
-    Requeue the message (AMQP 1.0 <code>released</code> outcome).   
+    Requeue the message (AMQP 1.0 <code>released</code> outcome).
     This means the message has not been processed and the broker can requeue it and deliver it
-    to the same or a different consumer.    
+    to the same or a different consumer.
     """
 
     def requeue(self, event: Event) -> None:
@@ -74,31 +81,34 @@ def requeue(self, event: Event) -> None:
 
     """
      Requeue the message with annotations to combine with the existing message annotations.
-         
         This means the message has not been processed and the broker can requeue it and deliver it
          to the same or a different consumer.
          Application-specific annotation keys must start with the <code>x-opt-</code> prefix.
          Annotation keys the broker understands starts with <code>x-</code>, but not with <code>x-opt-
          </code>.
-        
          This maps to the AMQP 1.0 <code>
          modified{delivery-failed = false, undeliverable-here = false}</code> outcome.
-        
          <param name="annotations"> annotations message annotations to combine with existing ones </param>
         <a
              href="https://docs.oasis-open.org/amqp/core/v1.0/os/amqp-core-messaging-v1.0-os.html#type-modified">AMQP
              1.0 <code>modified</code> outcome</a>
-        
         The annotations can be used only with Quorum queues, see https://www.rabbitmq.com/docs/amqp#modified-outcome
     """
 
     def requeue_with_annotations(
-            self, event: Event, annotations: Dict[str, "PythonAMQPData"]
+        self, event: Event, annotations: Dict[str, "PythonAMQPData"]
     ) -> None:
         dlv = event.delivery
         dlv.local.annotations = annotations
         dlv.local.failed = False
         dlv.local.undeliverable = False
 
+        validated = validate_annotations(annotations.keys())
+
+        if validated is False:
+            raise ArgumentOutOfRangeException(
+                "Message annotation key must start with 'x-'"
+            )
+
         dlv.update(Delivery.MODIFIED)
         dlv.settle()

rabbitmq_amqp_python_client/exceptions.py

@@ -5,3 +5,12 @@ def __init__(self, msg: str):
 
     def __str__(self) -> str:
         return repr(self.msg)
+
+
+class ArgumentOutOfRangeException(BaseException):
+    # Constructor or Initializer
+    def __init__(self, msg: str):
+        self.msg = msg
+
+    def __str__(self) -> str:
+        return repr(self.msg)

rabbitmq_amqp_python_client/utils.py

@@ -0,0 +1,9 @@
+def validate_annotations(annotations: []) -> bool:  # type: ignore
+    validated = True
+    for annotation in annotations:
+        if len(annotation) > 0 and annotation[:2] == "x-":
+            pass
+        else:
+            validated = False
+            return validated
+    return validated

tests/conftest.py

@@ -2,8 +2,8 @@
 
 from rabbitmq_amqp_python_client import (
     AddressHelper,
-    Connection,
     AMQPMessagingHandler,
+    Connection,
     Event,
     symbol,
 )
@@ -144,3 +144,19 @@ def on_message(self, event: Event):
         if self._received == 1000:
             event.connection.close()
             raise ConsumerTestException("consumed")
+
+
+class MyMessageHandlerRequeueWithInvalidAnnotations(AMQPMessagingHandler):
+
+    def __init__(self):
+        super().__init__()
+        self._received = 0
+
+    def on_message(self, event: Event):
+        annotations = {}
+        annotations[symbol("invalid")] = "x-test1"
+        self.delivery_context.requeue_with_annotations(event, annotations)
+        self._received = self._received + 1
+        if self._received == 1000:
+            event.connection.close()
+            raise ConsumerTestException("consumed")

tests/test_consumer.py

@@ -1,5 +1,6 @@
 from rabbitmq_amqp_python_client import (
     AddressHelper,
+    ArgumentOutOfRangeException,
     Connection,
     QuorumQueueSpecification,
 )
@@ -12,6 +13,7 @@
     MyMessageHandlerNoack,
     MyMessageHandlerRequeue,
     MyMessageHandlerRequeueWithAnnotations,
+    MyMessageHandlerRequeueWithInvalidAnnotations,
 )
 from .utils import (
     cleanup_dead_lettering,
@@ -320,3 +322,43 @@ def test_consumer_async_queue_with_requeue_with_annotations(
     assert "x-opt-string" in message.annotations
 
     assert message_count > 0
+
+
+def test_consumer_async_queue_with_requeue_with_invalid_annotations(
+    connection: Connection,
+) -> None:
+    messages_to_send = 1000
+    test_failure = True
+
+    queue_name = "test-queue-async-requeue"
+
+    management = connection.management()
+
+    management.declare_queue(QuorumQueueSpecification(name=queue_name))
+
+    addr_queue = AddressHelper.queue_address(queue_name)
+
+    publish_messages(connection, messages_to_send, queue_name)
+
+    # we closed the connection so we need to open a new one
+    connection_consumer = create_connection()
+
+    try:
+        consumer = connection_consumer.consumer(
+            addr_queue, handler=MyMessageHandlerRequeueWithInvalidAnnotations()
+        )
+
+        consumer.run()
+    # ack to terminate the consumer
+    except ConsumerTestException:
+        pass
+
+    except ArgumentOutOfRangeException:
+        test_failure = False
+
+    consumer.close()
+
+    management.delete_queue(queue_name)
+    management.close()
+
+    assert test_failure is False