Accept AMQP field arrays in Message.headers

rejected/models.py

@@ -160,9 +160,7 @@ class Message(pydantic.BaseModel):
     correlation_id: str | None
     delivery_mode: int | None
     expiration: str | None
-    headers: dict[
-        str, bool | dict[str, typing.Any] | float | int | str | bytes
-    ]
+    headers: dict[str, typing.Any]
     message_id: str | None
     type: str | None
     priority: int | None

rejected/process.py

@@ -21,6 +21,7 @@
 import pika
 import pika.channel
 import pika.spec
+import pydantic
 
 try:
     import sentry_sdk
@@ -442,13 +443,6 @@ def on_message(
         body: bytes,
     ) -> None:
         """Process a message from Rabbit"""
-        timestamp = (
-            datetime.datetime.fromtimestamp(
-                properties.timestamp, tz=datetime.UTC
-            )
-            if properties.timestamp
-            else None
-        )
         if isinstance(method, pika.spec.Basic.Deliver):
             delivery_tag = method.delivery_tag
             redelivered = method.redelivered
@@ -458,11 +452,15 @@ def on_message(
             redelivered = False
             returned = True
 
-        ctx = models.ProcessingContext(
-            connection=self.connections[name],
-            channel=channel,
-            raw_body=body,
-            message=models.Message(
+        try:
+            timestamp = (
+                datetime.datetime.fromtimestamp(
+                    properties.timestamp, tz=datetime.UTC
+                )
+                if properties.timestamp
+                else None
+            )
+            message = models.Message(
                 delivery_tag=delivery_tag,
                 exchange=method.exchange,
                 routing_key=method.routing_key,
@@ -484,7 +482,28 @@ def on_message(
                 returned=returned,
                 timestamp=timestamp,
                 user_id=properties.user_id,
-            ),
+            )
+        except (
+            pydantic.ValidationError,
+            OverflowError,
+            OSError,
+            ValueError,
+        ) as error:
+            LOGGER.error(
+                'Rejecting message %s on %s: invalid message properties: %s',
+                delivery_tag,
+                name,
+                error,
+            )
+            if delivery_tag is not None:
+                channel.basic_reject(delivery_tag=delivery_tag, requeue=False)
+            return
+
+        ctx = models.ProcessingContext(
+            connection=self.connections[name],
+            channel=channel,
+            raw_body=body,
+            message=message,
         )
         self._schedule(self.invoke_consumer(ctx))
 

tests/test_models.py

@@ -0,0 +1,83 @@
+import datetime
+import decimal
+import typing
+import unittest
+
+import pydantic
+
+from rejected import models
+
+
+def _make_message(headers: typing.Any) -> models.Message:
+    return models.Message(
+        delivery_tag=1,
+        exchange='ex',
+        routing_key='k',
+        body=b'',
+        app_id=None,
+        content_encoding=None,
+        content_type=None,
+        correlation_id=None,
+        delivery_mode=None,
+        expiration=None,
+        headers=headers,
+        message_id=None,
+        type=None,
+        priority=None,
+        redelivered=False,
+        reply_to=None,
+        timestamp=None,
+        user_id=None,
+    )
+
+
+class MessageHeadersTestCase(unittest.TestCase):
+    def test_scalar_values(self) -> None:
+        headers = {
+            'bool': True,
+            'int': 1,
+            'float': 1.5,
+            'str': 'value',
+            'bytes': b'value',
+            'decimal': decimal.Decimal('1.5'),
+            'timestamp': datetime.datetime.now(tz=datetime.UTC),
+            'void': None,
+        }
+        self.assertEqual(_make_message(headers).headers, headers)
+
+    def test_nested_table(self) -> None:
+        headers = {'table': {'nested': {'key': 'value'}, 'count': 1}}
+        self.assertEqual(_make_message(headers).headers, headers)
+
+    def test_array_of_scalars(self) -> None:
+        headers = {'array': ['a', 1, True, None]}
+        self.assertEqual(_make_message(headers).headers, headers)
+
+    def test_x_death_array_of_tables(self) -> None:
+        """RabbitMQ sets ``x-death`` on every dead-lettered message as an
+        array of tables, which pika decodes to a list of dicts."""
+        headers = {
+            'x-death': [
+                {
+                    'count': 1,
+                    'reason': 'rejected',
+                    'queue': 'q',
+                    'time': datetime.datetime.now(tz=datetime.UTC),
+                    'exchange': 'ex',
+                    'routing-keys': ['k'],
+                }
+            ],
+            'x-first-death-exchange': 'ex',
+            'x-first-death-queue': 'q',
+            'x-first-death-reason': 'rejected',
+        }
+        self.assertEqual(_make_message(headers).headers, headers)
+
+    def test_empty_headers(self) -> None:
+        self.assertEqual(_make_message({}).headers, {})
+
+    def test_non_dict_headers_rejected(self) -> None:
+        for value in (['x-death'], 'value', 1, None):
+            with self.subTest(value=value):
+                with self.assertRaises(pydantic.ValidationError):
+                    _make_message(value)

tests/test_process.py

@@ -8,6 +8,8 @@
 import unittest
 from unittest import mock
 
+import pika.spec
+
 from rejected import __version__, codecs, consumer, models, process
 from rejected import config as config_module
 from rejected import connection as connection_mod
@@ -560,3 +562,96 @@ def test_on_sigprof_schedules_report(self) -> None:
         p.ioloop = mock.Mock()
         p.on_sigprof(signal.SIGPROF, None)
         p.ioloop.call_soon_threadsafe.assert_called_once_with(p._report_stats)
+
+    def _deliver(
+        self,
+        p: process.Process,
+        headers: dict[str, typing.Any] | None = None,
+        delivery_tag: int | None = 1,
+        **properties: typing.Any,
+    ) -> mock.Mock:
+        channel = mock.Mock(spec=['basic_reject', 'basic_nack'])
+        if delivery_tag is None:
+            method = pika.spec.Basic.Return(
+                exchange='ex', routing_key='k', reply_code=312
+            )
+        else:
+            method = pika.spec.Basic.Deliver(
+                delivery_tag=delivery_tag, exchange='ex', routing_key='k'
+            )
+        properties = pika.spec.BasicProperties(headers=headers, **properties)
+        p.on_message('MockConnection', channel, method, properties, b'{}')
+        return channel
+
+    async def test_on_message_with_x_death_header_schedules_consumer(
+        self,
+    ) -> None:
+        """A dead-lettered message carries an array-of-tables header, which
+        must build a Message and be handed to the consumer."""
+        p = self.mock_setup()
+        p.connections['MockConnection'] = mock.Mock(
+            spec=connection_mod.Connection
+        )
+        headers = {'x-death': [{'count': 1, 'routing-keys': ['k']}]}
+        with mock.patch.object(p, '_schedule') as schedule:
+            with mock.patch.object(
+                p, 'invoke_consumer', new_callable=mock.Mock
+            ) as invoke:
+                channel = self._deliver(p, headers)
+        schedule.assert_called_once()
+        invoke.assert_called_once()
+        ctx = invoke.call_args.args[0]
+        self.assertEqual(ctx.message.headers, headers)
+        channel.basic_reject.assert_not_called()
+
+    async def test_on_message_invalid_properties_rejects_delivery(
+        self,
+    ) -> None:
+        """A message that fails Message validation is rejected without
+        requeue instead of raising into pika's delivery callback."""
+        p = self.mock_setup()
+        p.connections['MockConnection'] = mock.Mock(
+            spec=connection_mod.Connection
+        )
+        with mock.patch.object(p, '_schedule') as schedule:
+            with self.assertLogs('rejected.process', level='ERROR') as logs:
+                channel = self._deliver(p, priority='not-an-int')
+        schedule.assert_not_called()
+        channel.basic_reject.assert_called_once_with(
+            delivery_tag=1, requeue=False
+        )
+        self.assertIn('Rejecting message 1', logs.output[0])
+
+    async def test_on_message_out_of_range_timestamp_rejects_delivery(
+        self,
+    ) -> None:
+        """A timestamp that cannot be converted to a datetime follows the
+        same reject-without-requeue path as other invalid properties."""
+        p = self.mock_setup()
+        p.connections['MockConnection'] = mock.Mock(
+            spec=connection_mod.Connection
+        )
+        with mock.patch.object(p, '_schedule') as schedule:
+            with self.assertLogs('rejected.process', level='ERROR'):
+                channel = self._deliver(p, timestamp=10**20)
+        schedule.assert_not_called()
+        channel.basic_reject.assert_called_once_with(
+            delivery_tag=1, requeue=False
+        )
+
+    async def test_on_message_invalid_properties_returned_message(
+        self,
+    ) -> None:
+        """A returned message has no delivery tag so there is nothing to
+        reject, but it still must not raise."""
+        p = self.mock_setup()
+        p.connections['MockConnection'] = mock.Mock(
+            spec=connection_mod.Connection
+        )
+        with mock.patch.object(p, '_schedule') as schedule:
+            with self.assertLogs('rejected.process', level='ERROR'):
+                channel = self._deliver(
+                    p, delivery_tag=None, priority='not-an-int'
+                )
+        schedule.assert_not_called()
+        channel.basic_reject.assert_not_called()