✨ Kafka server

Author: aleksul

.pre-commit-config.yaml

@@ -45,6 +45,6 @@ repos:
     name: mypy
     language: python
     pass_filenames: false
-    entry: uv run mypy .
+    entry: uv run --no-active mypy .
     types: [python]
     require_serial: true

pyproject.toml

@@ -56,6 +56,7 @@ requires-python = ">=3.10"
 
 [project.optional-dependencies]
 amqp = []
+kafka = ["aiokafka>=0.13.0,<0.14.0"]
 pubsub = [
   "google-auth>=2.43.0",
   "grpcio>=1.76.0"
@@ -80,7 +81,7 @@ exclude_lines = [
 fail_under = 100
 
 [tool.coverage.run]
-omit = ["repid/connections/rabbitmq/protocols.py"]
+omit = ["repid/connections/rabbitmq/protocols.py", "repid/connections/kafka/protocols.py"]
 source = ["repid"]
 
 [tool.hatch.build.targets.wheel]
@@ -119,7 +120,7 @@ module = "tests.*"
 
 [[tool.mypy.overrides]]
 ignore_missing_imports = true
-module = ["pytest_docker_tools"]
+module = ["pytest_docker_tools", "aiokafka.*"]
 
 [tool.pytest.ini_options]
 asyncio_mode = "auto"

repid/connections/__init__.py

@@ -29,6 +29,11 @@
 
     __all__ += ["RedisServer"]
 
+if is_installed("aiokafka"):
+    from repid.connections.kafka import KafkaServer
+
+    __all__ += ["KafkaServer"]
+
 # check via try-import because pubsub is a sub-package of gcloud.aio
 imported_pubsub = False
 try:

repid/connections/kafka/__init__.py

@@ -0,0 +1 @@
+from repid.connections.kafka.message_broker import KafkaServer as KafkaServer

repid/connections/kafka/message.py

@@ -0,0 +1,168 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable, Coroutine
+from typing import TYPE_CHECKING, Any
+
+from repid.connections.abc import MessageAction, ReceivedMessageT
+
+if TYPE_CHECKING:
+    from repid.connections.kafka.message_broker import KafkaServer
+    from repid.connections.kafka.protocols import ConsumerRecordProtocol
+
+logger = logging.getLogger("repid.connections.kafka")
+
+
+class KafkaReceivedMessage(ReceivedMessageT):
+    def __init__(
+        self,
+        server: KafkaServer,
+        record: ConsumerRecordProtocol,
+        mark_complete_callback: Callable[[ConsumerRecordProtocol], Coroutine[Any, Any, None]],
+    ) -> None:
+        self._server = server
+        self._record = record
+        self._mark_complete_callback = mark_complete_callback
+
+        self._action: MessageAction | None = None
+        self._payload = record.value
+        self._channel = record.topic
+        self._message_id = str(record.offset)
+
+        self._headers: dict[str, str] = {}
+        self._content_type: str | None = None
+
+        if record.headers:
+            for k, v in record.headers:
+                if k == "content-type":
+                    self._content_type = v.decode()
+                else:
+                    self._headers[k] = v.decode()
+
+    @property
+    def payload(self) -> bytes:
+        return self._payload if self._payload is not None else b""
+
+    @property
+    def headers(self) -> dict[str, str] | None:
+        return self._headers
+
+    @property
+    def content_type(self) -> str | None:
+        return self._content_type
+
+    @property
+    def channel(self) -> str:
+        return str(self._channel)
+
+    @property
+    def action(self) -> MessageAction | None:
+        return self._action
+
+    @property
+    def is_acted_on(self) -> bool:
+        return self._action is not None
+
+    @property
+    def message_id(self) -> str:
+        return self._message_id
+
+    async def ack(self) -> None:
+        if self._action is not None:
+            return
+        self._action = MessageAction.acked
+        await self._mark_complete_callback(self._record)
+
+    async def nack(self) -> None:
+        if self._action is not None:
+            return
+        self._action = MessageAction.nacked
+
+        if self._server._producer is None:  # pragma: no cover
+            # it should be impossible to get here since subscribing requires a connection,
+            # but we'll check just in case
+            raise ConnectionError("Kafka producer is not connected.")
+
+        dlq_topic_strategy = self._server._dlq_topic_strategy
+        if dlq_topic_strategy is not None and self._server._producer is not None:
+            dlq_topic = dlq_topic_strategy(self._channel)
+            headers = []
+            if self._headers:
+                for k, v in self._headers.items():
+                    headers.append((k, v.encode()))
+            if self._content_type:
+                headers.append(("content-type", self._content_type.encode()))
+
+            headers.append(("x-original-topic", self._channel.encode()))
+            headers.append(("x-original-offset", str(self._record.offset).encode()))
+
+            await self._server._producer.send_and_wait(
+                topic=dlq_topic,
+                value=self._payload,
+                headers=headers,
+            )
+
+        await self._mark_complete_callback(self._record)
+
+    async def reject(self) -> None:
+        if self._action is not None:
+            return
+        self._action = MessageAction.rejected
+
+        if self._server._producer is None:  # pragma: no cover
+            # it should be impossible to get here since subscribing requires a connection,
+            # but we'll check just in case
+            raise ConnectionError("Kafka producer is not connected.")
+
+        reject_topic_strategy = self._server._reject_topic_strategy
+        if reject_topic_strategy is not None:
+            reject_topic = reject_topic_strategy(self._channel)
+            headers = []
+            if self._headers:
+                for k, v in self._headers.items():
+                    headers.append((k, v.encode()))
+            if self._content_type:
+                headers.append(("content-type", self._content_type.encode()))
+
+            await self._server._producer.send_and_wait(
+                topic=reject_topic,
+                value=self._payload,
+                headers=headers,
+            )
+
+        await self._mark_complete_callback(self._record)
+
+    async def reply(
+        self,
+        *,
+        payload: bytes,
+        headers: dict[str, str] | None = None,
+        content_type: str | None = None,
+        channel: str | None = None,
+        server_specific_parameters: dict[str, Any] | None = None,  # noqa: ARG002
+    ) -> None:
+        if self._action is not None:
+            return
+        self._action = MessageAction.replied
+
+        if self._server._producer is None:  # pragma: no cover
+            # it should be impossible to get here since subscribing requires a connection,
+            # but we'll check just in case
+            raise ConnectionError("Kafka producer is not connected.")
+
+        reply_headers = []
+        if headers:
+            for k, v in headers.items():
+                reply_headers.append((k, v.encode()))
+        if content_type:
+            reply_headers.append(("content-type", content_type.encode()))
+
+        reply_channel = channel or self._channel
+
+        await self._server._producer.send_and_wait(
+            topic=reply_channel,
+            value=payload,
+            headers=reply_headers,
+        )
+
+        await self._mark_complete_callback(self._record)