feat(kafka): preserve the record key on retry and dead-letter writes

The original record key now travels with retry republishes, the retry-to-work
re-inject and both the envelope and raw dead-letter writes, as broker-bindings
section 6 recommends. A null key stays null and an empty key stays empty;
partition equality with the source record is not guaranteed.

Author: muhammetsafak

CHANGELOG.md

@@ -16,8 +16,17 @@ The envelope wire format is versioned separately by `meta.schema_version`
   implements it. `KafkaTransport` implements it.
 - `KafkaDeliveryError` (a `BabelQueueError`) and the optional `KafkaTransport(delivery_timeout=30.0)`
   constructor argument.
+- `ReceivedMessage.key` (`Optional[bytes]`, `None` = null key) and an optional `key=` argument on
+  `SettlePublisher.publish_settle()` / `publish_keeping_route()`. `KafkaTransport.pop()` fills the
+  key from the consumed record.
 
 ### Fixed
+- **Kafka: the original record key now survives the retry republish, the envelope DLQ and the
+  raw (poison / gate-rejection) DLQ.** The key was dropped on every settle write, so a keyed record
+  lost its key (and with it per-key partition affinity) after the first retry. A null key stays
+  null, an empty key stays empty, and a key is never synthesised. `redrive` restores to `<queue>.dlq`
+  also keep it. Partition equality is not guaranteed. The envelope and the `dead_letter` block are
+  unchanged.
 - **Kafka: retry / dead-letter writes are now broker-confirmed before the offset is committed.**
   `KafkaTransport.publish_settle()` waits for this write's own delivery report
   (`flush(delivery_timeout)`). A delivery error or timeout raises `KafkaDeliveryError`, nothing is

src/babelqueue/app.py

@@ -415,7 +415,7 @@ def _publish_for(self, received: ReceivedMessage, queue: str, body: str) -> None
         or its dead-letter record). Keeps the routing URN the broker property carried on this
         delivery when the transport can pin it (:class:`RoutedPublisher`), so a property/body
         mismatch does not flip to the body URN on the next attempt or in the DLQ."""
-        publish_keeping_route(self.transport, queue, body, received.route_urn)
+        publish_keeping_route(self.transport, queue, body, received.route_urn, received.key)
 
     def _ack_settled(self, received: ReceivedMessage) -> None:
         """Acknowledge a message whose outcome is already final (handled, dead-lettered,

src/babelqueue/kafka_transport.py

@@ -17,7 +17,9 @@
 
 Settlement writes (retry republish, DLQ) go through ``publish_settle`` and wait for the broker's
 delivery report (``delivery_timeout``); a failure raises ``KafkaDeliveryError`` and the offset is
-not committed. ``publish`` itself stays fire-and-forget. A record value that is not valid UTF-8 is handed to the runtime under a
+not committed. ``publish`` itself stays fire-and-forget. Those writes also carry the consumed
+record's key unchanged (``ReceivedMessage.key``; null stays null, empty stays empty, never
+synthesised). A record value that is not valid UTF-8 is handed to the runtime under a
 reserved route so its original bytes are dead-lettered raw (written back with no projected headers;
 with the DLQ off, or ``UnknownUrnStrategy.DELETE``, it is dropped and committed). The runtime logs
 it as a version-gated message because it reuses that settle path; the route in the log names it.
@@ -219,6 +221,13 @@ def _payload(message: Any) -> str:
             return bytes(value).decode("utf-8", "surrogateescape")
         return str(value) if value is not None else ""
 
+    @staticmethod
+    def _key(message: Any) -> Optional[bytes]:
+        """The record key as the client reports it (``None`` = null key); tolerant of an injected
+        consumer whose records expose no ``key()``."""
+        key_of = getattr(message, "key", None)
+        return key_of() if callable(key_of) else None
+
     @staticmethod
     def _is_non_utf8(message: Any) -> bool:
         value = message.value()
@@ -238,12 +247,21 @@ def publish(self, queue: str, body: str) -> None:
         producer.produce(queue, **kwargs)
         producer.poll(0)
 
-    def publish_settle(self, queue: str, body: str, route_urn: str = "") -> None:
+    def publish_settle(
+        self, queue: str, body: str, route_urn: str = "", key: Optional[bytes] = None
+    ) -> None:
         """:class:`~babelqueue.transport.SettlePublisher`: the runtime's settle write (retry
         republish, envelope / poison / gate DLQ, ``redrive`` restore). Unlike :meth:`publish` it
         waits for the broker's delivery report, because the offset may only be committed once the
-        write is confirmed; a delivery error or timeout raises :class:`KafkaDeliveryError`."""
-        self._produce_confirmed(self._producer_(), queue, self._record_kwargs(body, raw_ok=True))
+        write is confirmed; a delivery error or timeout raises :class:`KafkaDeliveryError`.
+
+        ``key`` is the consumed record's key (:attr:`ReceivedMessage.key`): it is written back
+        unchanged, so a retry or DLQ record keeps the original key (``None`` stays a null key, an
+        empty key stays empty, and no key is ever synthesised). §6."""
+        kwargs = self._record_kwargs(body, raw_ok=True)
+        if key is not None:
+            kwargs["key"] = key
+        self._produce_confirmed(self._producer_(), queue, kwargs)
 
     @classmethod
     def _record_kwargs(cls, body: str, raw_ok: bool = False) -> Dict[str, Any]:
@@ -310,6 +328,7 @@ def pop(self, queue: str, timeout: float = 1.0) -> Optional[ReceivedMessage]:
                 queue=queue,
                 handle=message,
                 route_urn=UNSUPPORTED_SCHEMA_ROUTE_PREFIX + quote(unsupported, safe=""),
+                key=self._key(message),
             )
             self._pending[queue] = message
             return received
@@ -321,12 +340,16 @@ def pop(self, queue: str, timeout: float = 1.0) -> Optional[ReceivedMessage]:
                 "Kafka record on %r is not valid UTF-8; dead-lettering its raw bytes.", queue
             )
             received = ReceivedMessage(
-                body=self._payload(message), queue=queue, handle=message, route_urn=NON_UTF8_ROUTE
+                body=self._payload(message),
+                queue=queue,
+                handle=message,
+                route_urn=NON_UTF8_ROUTE,
+                key=self._key(message),
             )
             self._pending[queue] = message
             return received
         body = self._reconcile(self._payload(message), headers)
-        received = ReceivedMessage(body=body, queue=queue, handle=message)
+        received = ReceivedMessage(body=body, queue=queue, handle=message, key=self._key(message))
         # Track the record as in flight only once ``pop`` has built it. A decode error above
         # escapes before this point and must not leave the record pending: the next ``pop`` would
         # seek back to it and fail identically, wedging the partition forever.

src/babelqueue/redrive.py

@@ -88,7 +88,7 @@ def redrive(
     result = RedriveResult()
     for message, envelope in batch:
         if envelope is None:
-            publish_keeping_route(transport, dlq, message.body, message.route_urn)  # restore the poison body; never drop it
+            publish_keeping_route(transport, dlq, message.body, message.route_urn, message.key)  # restore the poison body; never drop it
             transport.ack(message)
             result.skipped += 1
             result.items.append(RedriveItem("", "", "", "", dlq, "", False))
@@ -109,7 +109,7 @@ def redrive(
         )
 
         if select is not None and not select(envelope):
-            publish_keeping_route(transport, dlq, message.body, message.route_urn)  # not selected: restore unchanged
+            publish_keeping_route(transport, dlq, message.body, message.route_urn, message.key)  # not selected: restore unchanged
             transport.ack(message)
             result.skipped += 1
             result.items.append(item)
@@ -119,7 +119,7 @@ def redrive(
         item.to = target
 
         if dry_run:
-            publish_keeping_route(transport, dlq, message.body, message.route_urn)  # report the plan; restore unchanged
+            publish_keeping_route(transport, dlq, message.body, message.route_urn, message.key)  # report the plan; restore unchanged
             transport.ack(message)
             result.skipped += 1
             result.items.append(item)
@@ -133,7 +133,7 @@ def redrive(
                 transport, target, EnvelopeCodec.encode(reset), bypass, message.route_urn
             )
         except Exception:
-            publish_keeping_route(transport, dlq, message.body, message.route_urn)  # restore on a publish failure, then surface
+            publish_keeping_route(transport, dlq, message.body, message.route_urn, message.key)  # restore on a publish failure, then surface
             transport.ack(message)
             raise
         transport.ack(message)

src/babelqueue/transport.py

@@ -43,6 +43,11 @@ class ReceivedMessage:
     #: body is left untouched (handler, retry and dead-letter keep the producer's original body).
     #: Empty means "route on the body URN" (the default for every other transport).
     route_urn: str = ""
+    #: The broker record key of this delivery, byte-for-byte (``None`` = null key; an empty
+    #: ``bytes`` is an empty key), for transports that have one (Kafka). The runtime hands it back
+    #: to a :class:`SettlePublisher` on its own re-publishes (retry, DLQ) so the key is not lost;
+    #: it never synthesises one. Not part of the envelope; ``None`` for every other transport.
+    key: Optional[bytes] = None
 
 
 class Transport(ABC):
@@ -100,18 +105,33 @@ class SettlePublisher(Protocol):
     the delivery unsettled. A user's own :meth:`Transport.publish` is never routed here. A transport
     that does not implement it gets the :class:`RoutedPublisher` / plain-publish behaviour."""
 
-    def publish_settle(self, queue: str, body: str, route_urn: str = "") -> None:
-        """Append ``body`` to ``queue`` and return only once the broker confirmed the write."""
+    def publish_settle(
+        self, queue: str, body: str, route_urn: str = "", key: Optional[bytes] = None
+    ) -> None:
+        """Append ``body`` to ``queue`` and return only once the broker confirmed the write.
+
+        ``key`` is the consumed record's key (:attr:`ReceivedMessage.key`), carried unchanged
+        when the broker has record keys; ``None`` means a null key."""
         ...
 
 
-def publish_keeping_route(transport: "Transport", queue: str, body: str, route_urn: str = "") -> None:
+def publish_keeping_route(
+    transport: "Transport",
+    queue: str,
+    body: str,
+    route_urn: str = "",
+    key: Optional[bytes] = None,
+) -> None:
     """Publish ``body`` to ``queue``, carrying ``route_urn`` when it is non-empty and the
     transport implements :class:`RoutedPublisher`; otherwise a plain :meth:`Transport.publish`.
-    A :class:`SettlePublisher` transport gets the confirmed settle write instead.
+    A :class:`SettlePublisher` transport gets the confirmed settle write instead, with the
+    consumed record's ``key`` when there is one (a null key is simply not passed).
     Shared by the runtime's re-publishes (retry, release, dead-letter) and by ``redrive``."""
     if isinstance(transport, SettlePublisher):
-        transport.publish_settle(queue, body, route_urn)
+        if key is None:
+            transport.publish_settle(queue, body, route_urn)
+        else:
+            transport.publish_settle(queue, body, route_urn, key=key)
     elif route_urn and isinstance(transport, RoutedPublisher):
         transport.publish_routed(queue, body, route_urn)
     else:

tests/test_kafka_transport.py

@@ -48,10 +48,14 @@ def _assert_no_codec_warnings(case: unittest.TestCase, fn) -> None:
 
 
 class FakeMessage:
-    def __init__(self, value, headers=None, err=None) -> None:
+    def __init__(self, value, headers=None, err=None, key=None) -> None:
         self._value = value
         self._headers = headers or []
         self._err = err
+        self._key = key
+
+    def key(self):
+        return self._key
 
     def value(self):
         return self._value
@@ -72,7 +76,15 @@ def __init__(self) -> None:
         self._reports = []
 
     def produce(self, topic, value=None, headers=None, timestamp=None, on_delivery=None, **kw) -> None:
-        self.produced.append({"topic": topic, "value": value, "headers": headers, "timestamp": timestamp})
+        self.produced.append(
+            {
+                "topic": topic,
+                "value": value,
+                "headers": headers,
+                "timestamp": timestamp,
+                "key": kw.get("key"),
+            }
+        )
         if on_delivery is not None:
             self._reports.append(on_delivery)
 
@@ -104,8 +116,8 @@ def close(self) -> None:
 class OffsetMessage(FakeMessage):
     """A record that knows its place in a single-partition log."""
 
-    def __init__(self, value, headers, offset, topic="orders", partition=0) -> None:
-        super().__init__(value, headers=headers)
+    def __init__(self, value, headers, offset, topic="orders", partition=0, key=None) -> None:
+        super().__init__(value, headers=headers, key=key)
         self._offset, self._topic, self._partition = offset, topic, partition
 
     def topic(self):
@@ -638,7 +650,7 @@ def __init__(self, outcomes=(), stray=0) -> None:
         self._queued = []
 
     def produce(self, topic, value=None, headers=None, timestamp=None, on_delivery=None, **kw) -> None:
-        super().produce(topic, value=value, headers=headers, timestamp=timestamp)
+        super().produce(topic, value=value, headers=headers, timestamp=timestamp, **kw)
         outcome = self.outcomes.pop(0) if self.outcomes else "ok"
         if on_delivery is not None:
             self.confirmed.append(topic)
@@ -864,6 +876,104 @@ def test_non_utf8_attempts_header_does_not_raise(self):
         self.assertEqual(EnvelopeCodec.decode(message.body)["attempts"], 2)  # the body's own count
 
 
+class KafkaKeyPreservationTest(unittest.TestCase):
+    """The consumed record's key is carried, byte-for-byte, onto every record the runtime itself
+    writes for it (retry republish, envelope DLQ, gate DLQ, raw non-UTF-8 DLQ). Null stays null,
+    empty stays empty, and no key is ever made up (§6)."""
+
+    KEYS = (b"order-1042", b"", None)
+
+    def _run(self, record_value, key, headers=(), **app_kw):
+        consumer = OffsetConsumer([OffsetMessage(record_value, list(headers), offset=0, key=key)])
+        producer = ReportingProducer()
+        app = BabelQueue(transport=transport(producer, consumer), **app_kw)
+        app._rewind_pause = 0.0  # noqa: SLF001
+
+        @app.handler(URN)
+        def _on(data, meta):
+            raise RuntimeError("handler failure")
+
+        app.consume(queue="orders", max_messages=1)
+        self.assertEqual(consumer.commits, [0])
+        self.assertEqual(len(producer.produced), 1)
+        return producer.produced[0]
+
+    def test_pop_surfaces_the_record_key(self):
+        for key in self.KEYS:
+            with self.subTest(key=key):
+                got = transport(consumer=FakeConsumer([FakeMessage(body().encode(), key=key)])).pop("orders")
+                self.assertEqual(got.key, key)
+
+    def test_pop_tolerates_a_record_without_a_key_accessor(self):
+        class NoKey:
+            def value(self):
+                return body().encode()
+
+            def headers(self):
+                return []
+
+            def error(self):
+                return None
+
+        self.assertIsNone(transport(consumer=FakeConsumer([NoKey()])).pop("orders").key)
+
+    def test_retry_republish_keeps_the_key(self):
+        for key in self.KEYS:
+            with self.subTest(key=key):
+                record = self._run(body().encode(), key, max_attempts=3)
+                self.assertEqual(record["topic"], "orders")
+                self.assertEqual(record["key"], key)
+
+    def test_envelope_dead_letter_keeps_the_key(self):
+        for key in self.KEYS:
+            with self.subTest(key=key):
+                record = self._run(body().encode(), key, max_attempts=1, dead_letter=True)
+                self.assertEqual(record["topic"], "orders.dlq")
+                self.assertEqual(record["key"], key)
+
+    def test_gated_dead_letter_keeps_the_key(self):
+        for key in self.KEYS:
+            with self.subTest(key=key):
+                record = self._run(
+                    body().encode(), key, headers=[("bq-schema-version", b"2")], dead_letter=True
+                )
+                self.assertEqual(record["topic"], "orders.dlq")
+                self.assertEqual(record["key"], key)
+
+    def test_non_utf8_raw_dead_letter_keeps_the_key(self):
+        bad = b"\xff\xfe not utf-8 \x80"
+        for key in self.KEYS:
+            with self.subTest(key=key):
+                with self.assertLogs("babelqueue.kafka", level="WARNING"):
+                    record = self._run(bad, key, dead_letter=True)
+                self.assertEqual(record["topic"], "orders.dlq")
+                self.assertEqual(record["value"], bad)
+                self.assertEqual(record["key"], key)
+
+    def test_a_null_key_is_not_passed_to_produce(self):
+        class StrictProducer(FakeProducer):
+            def produce(self, topic, value=None, headers=None, timestamp=None, on_delivery=None, **kw):
+                assert "key" not in kw, "a null key must not be synthesised"
+                super().produce(topic, value=value, headers=headers, timestamp=timestamp, on_delivery=on_delivery)
+
+        transport(producer=StrictProducer()).publish_settle("orders", body(), key=None)
+
+    def test_plain_publish_does_not_set_a_key(self):
+        producer = FakeProducer()
+        transport(producer=producer).publish("orders", body())
+        self.assertIsNone(producer.produced[0]["key"])
+
+    def test_redrive_restore_keeps_the_key(self):
+        from babelqueue import redrive
+
+        consumer = FakeConsumer([FakeMessage(b"not json at all", key=b"k-7")])
+        producer = FakeProducer()
+        result = redrive.redrive(transport(producer, consumer), "orders.dlq", timeout=0)
+        self.assertEqual(result.skipped, 1)
+        self.assertEqual(producer.produced[0]["topic"], "orders.dlq")
+        self.assertEqual(producer.produced[0]["key"], b"k-7")
+
+
 class KafkaSchemeTest(unittest.TestCase):
     def test_make_transport_routes_kafka_scheme_lazily(self):
         # Construction is lazy (no client built), so it succeeds with no confluent-kafka.