✨ Kafka server

Author: aleksul

pyproject.toml

@@ -56,6 +56,7 @@ requires-python = ">=3.10"
 
 [project.optional-dependencies]
 amqp = []
+kafka = ["aiokafka>=0.12.0"]
 pubsub = [
   "google-auth>=2.43.0",
   "grpcio>=1.76.0"
@@ -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,5 @@
+from __future__ import annotations
+
+from repid.connections.kafka.message_broker import KafkaServer
+
+__all__ = ["KafkaServer"]

repid/connections/kafka/message.py

@@ -0,0 +1,150 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+from repid.connections.abc import MessageAction, ReceivedMessageT
+
+if TYPE_CHECKING:
+    from repid.connections.kafka.message_broker import KafkaServer
+
+logger = logging.getLogger("repid.connections.kafka")
+
+
+class KafkaReceivedMessage(ReceivedMessageT):
+    def __init__(
+        self,
+        server: KafkaServer,
+        record: Any,
+        mark_complete_callback: Any,
+    ) -> 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
+
+        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 not None:
+            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=self._channel,
+                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 not None:
+            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)

repid/connections/kafka/message_broker.py

@@ -0,0 +1,179 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import AsyncIterator, Callable, Coroutine, Mapping, Sequence
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, Any
+
+from aiokafka import AIOKafkaConsumer, AIOKafkaProducer
+
+from repid.connections.abc import CapabilitiesT, SentMessageT, ServerT, SubscriberT
+from repid.connections.kafka.subscriber import KafkaSubscriber
+
+if TYPE_CHECKING:
+    from repid.asyncapi.models.common import ServerBindingsObject
+    from repid.asyncapi.models.servers import ServerVariable
+    from repid.connections.abc import ReceivedMessageT
+    from repid.data import ExternalDocs, Tag
+
+logger = logging.getLogger("repid.connections.kafka")
+
+
+class KafkaServer(ServerT):
+    def __init__(
+        self,
+        dsn: str,
+        *,
+        dlq_topic_strategy: Callable[[str], str] | None = lambda channel: f"repid_{channel}_dlq",
+        title: str | None = None,
+        summary: str | None = None,
+        description: str | None = None,
+        variables: Mapping[str, ServerVariable] | None = None,
+        security: Sequence[Any] | None = None,
+        tags: Sequence[Tag] | None = None,
+        external_docs: ExternalDocs | None = None,
+        bindings: ServerBindingsObject | None = None,
+    ) -> None:
+        self.dsn = dsn
+        self._producer: AIOKafkaProducer | None = None  # type: ignore[no-any-unimported]
+        self._dlq_topic_strategy = dlq_topic_strategy
+
+        # AsyncAPI metadata
+        self._title = title
+        self._summary = summary
+        self._description = description
+        self._variables = variables
+        self._security = security
+        self._tags = tags
+        self._external_docs = external_docs
+        self._bindings = bindings
+
+    @property
+    def host(self) -> str:
+        return self.dsn
+
+    @property
+    def protocol(self) -> str:
+        return "kafka"
+
+    @property
+    def pathname(self) -> str | None:
+        return None
+
+    @property
+    def title(self) -> str | None:
+        return self._title
+
+    @property
+    def summary(self) -> str | None:
+        return self._summary
+
+    @property
+    def description(self) -> str | None:
+        return self._description
+
+    @property
+    def protocol_version(self) -> str | None:
+        return None
+
+    @property
+    def variables(self) -> Mapping[str, ServerVariable] | None:
+        return self._variables
+
+    @property
+    def security(self) -> Sequence[Any] | None:
+        return self._security
+
+    @property
+    def tags(self) -> Sequence[Tag] | None:
+        return self._tags
+
+    @property
+    def external_docs(self) -> ExternalDocs | None:
+        return self._external_docs
+
+    @property
+    def bindings(self) -> ServerBindingsObject | None:
+        return self._bindings
+
+    @property
+    def capabilities(self) -> CapabilitiesT:
+        return {
+            "supports_acknowledgments": True,
+            "supports_persistence": True,
+            "supports_reply": True,
+            "supports_lightweight_pause": False,
+        }
+
+    @property
+    def is_connected(self) -> bool:
+        return self._producer is not None
+
+    async def connect(self) -> None:
+        if self._producer is None:
+            self._producer = AIOKafkaProducer(bootstrap_servers=self.dsn)
+            await self._producer.start()
+            logger.info("server.connect", extra={"host": self.dsn})
+
+    async def disconnect(self) -> None:
+        if self._producer is not None:
+            await self._producer.stop()
+            self._producer = None
+            logger.info("server.disconnect")
+
+    @asynccontextmanager
+    async def connection(self) -> AsyncIterator[KafkaServer]:
+        await self.connect()
+        try:
+            yield self
+        finally:
+            await self.disconnect()
+
+    async def publish(
+        self,
+        *,
+        channel: str,
+        message: SentMessageT,
+        server_specific_parameters: dict[str, Any] | None = None,  # noqa: ARG002
+    ) -> None:
+        if self._producer is None:
+            raise RuntimeError("Kafka producer is not connected.")
+
+        logger.debug("channel.publish", extra={"channel": channel})
+
+        headers = []
+        if message.headers:
+            for k, v in message.headers.items():
+                headers.append((k, v.encode()))
+        if message.content_type:
+            headers.append(("content-type", message.content_type.encode()))
+
+        await self._producer.send_and_wait(
+            topic=channel,
+            value=message.payload,
+            headers=headers,
+        )
+
+    async def subscribe(
+        self,
+        *,
+        channels_to_callbacks: dict[str, Callable[[ReceivedMessageT], Coroutine[None, None, None]]],
+        concurrency_limit: int | None = None,
+    ) -> SubscriberT:
+        logger.debug("channel.subscribe", extra={"channels": list(channels_to_callbacks.keys())})
+
+        consumer = AIOKafkaConsumer(
+            *channels_to_callbacks.keys(),
+            bootstrap_servers=self.dsn,
+            group_id="repid-group",
+            enable_auto_commit=False,
+            auto_offset_reset="earliest",
+        )
+        await consumer.start()
+
+        return KafkaSubscriber(
+            server=self,
+            consumer=consumer,
+            channels_to_callbacks=channels_to_callbacks,
+            concurrency_limit=concurrency_limit,
+        )