Merge pull request #10 from taskiq-python/feature/jetstream

Author: s3rius

.pre-commit-config.yaml

@@ -41,13 +41,15 @@ repos:
         language: system
         pass_filenames: false
         types: [python]
-        args: [--count, taskiq_nats]
+        args: [--count, taskiq_nats, tests]
 
       - id: mypy
         name: Validate types with MyPy
         entry: poetry run mypy
         language: system
+        pass_filenames: false
         types: [python]
+        args: [taskiq_nats, tests]
 
       - id: yesqa
         name: Remove usless noqa

README.md

@@ -1,6 +1,7 @@
 # Taskiq NATS
 
 Taskiq-nats is a plugin for taskiq that adds NATS broker.
+This package has support for NATS JetStream.
 
 ## Installation
 
@@ -16,7 +17,7 @@ Here's a minimal setup example with a broker and one task.
 
 ```python
 import asyncio
-from taskiq_nats import NatsBroker
+from taskiq_nats import NatsBroker, JetStreamBroker
 
 broker = NatsBroker(
     [
@@ -26,6 +27,17 @@ broker = NatsBroker(
     queue="random_queue_name",
 )
 
+# Or alternatively you can use a JetStream broker:
+broker = JetStreamBroker(
+    [
+        "nats://nats1:4222",
+        "nats://nats2:4222",
+    ],
+    queue="random_queue_name",
+    subject="my-subj",
+    stream_name="my-stream"
+)
+
 
 @broker.task
 async def my_lovely_task():

docker-compose.yml

@@ -10,6 +10,7 @@ services:
       - "-m"
       - "8222"
       - "--debug"
+      - "--jetstream"
     healthcheck:
       test:
         - "CMD"

poetry.lock

@@ -1,10 +1,10 @@
 """
 NATS integration for taskiq.
 
-This package contains broker that
+This package contains brokers that
 uses NATS as a message queue.
 """
 
-from taskiq_nats.broker import NatsBroker
+from taskiq_nats.broker import JetStreamBroker, NatsBroker
 
-__all__ = ["NatsBroker"]
+__all__ = ["NatsBroker", "JetStreamBroker"]

taskiq_nats/__init__.py

@@ -2,6 +2,8 @@
 from typing import Any, AsyncGenerator, Callable, List, Optional, TypeVar, Union
 
 from nats.aio.client import Client
+from nats.js import JetStreamContext
+from nats.js.api import StreamConfig
 from taskiq import AsyncBroker, AsyncResultBackend, BrokerMessage
 
 _T = TypeVar("_T")  # noqa: WPS111 (Too short)
@@ -76,3 +78,80 @@ async def shutdown(self) -> None:
         """Close connections to NATS."""
         await self.client.close()
         await super().shutdown()
+
+
+class JetStreamBroker(AsyncBroker):  # noqa: WPS230
+    """
+    JetStream broker for taskiq.
+
+    This broker creates a JetStream context
+    and uses it to send and receive messages.
+
+    This is useful for systems where you need to
+    be sure that messages are delivered to the workers.
+    """
+
+    def __init__(  # noqa: WPS211 (too many args)
+        self,
+        servers: Union[str, List[str]],
+        subject: str = "tasiq_tasks",
+        stream_name: str = "taskiq_jstream",
+        queue: Optional[str] = None,
+        result_backend: "Optional[AsyncResultBackend[_T]]" = None,
+        task_id_generator: Optional[Callable[[], str]] = None,
+        stream_config: Optional[StreamConfig] = None,
+        **connection_kwargs: Any,
+    ) -> None:
+        super().__init__(result_backend, task_id_generator)
+        self.servers = servers
+        self.client: Client = Client()
+        self.connection_kwargs = connection_kwargs
+        self.queue = queue
+        self.subject = subject
+        self.stream_name = stream_name
+        self.js: JetStreamContext
+        self.stream_config = stream_config or StreamConfig()
+
+    async def startup(self) -> None:
+        """
+        Startup event handler.
+
+        It simply connects to NATS cluster, and
+        setup JetStream.
+        """
+        await super().startup()
+        await self.client.connect(self.servers, **self.connection_kwargs)
+        self.js = self.client.jetstream()
+        if self.stream_config.name is None:
+            self.stream_config.name = self.stream_name
+        if not self.stream_config.subjects:
+            self.stream_config.subjects = [self.subject]
+        await self.js.add_stream(config=self.stream_config)
+
+    async def kick(self, message: BrokerMessage) -> None:
+        """
+        Send a message using NATS.
+
+        :param message: message to send.
+        """
+        await self.js.publish(
+            self.subject,
+            payload=message.message,
+            headers=message.labels,
+        )
+
+    async def listen(self) -> AsyncGenerator[bytes, None]:
+        """
+        Start listen to new messages.
+
+        :yield: incoming messages.
+        """
+        subscribe = await self.js.subscribe(self.subject, queue=self.queue or "")
+        async for message in subscribe.messages:
+            yield message.data
+            await message.ack()
+
+    async def shutdown(self) -> None:
+        """Close connections to NATS."""
+        await self.client.close()
+        await super().shutdown()

taskiq_nats/broker.py

@@ -1,4 +1,5 @@
 import os
+import uuid
 from typing import List
 
 import pytest
@@ -22,7 +23,7 @@ def nats_subject() -> str:
 
     :return: random string.
     """
-    return "12345"
+    return uuid.uuid4().hex
 
 
 @pytest.fixture

tests/conftest.py

@@ -6,17 +6,7 @@
 from taskiq import BrokerMessage
 
 from taskiq_nats import NatsBroker
-
-
-async def read_message(broker: NatsBroker) -> bytes:  # type: ignore
-    """
-    Read signle message from the broker's listen method.
-
-    :param broker: current broker.
-    :return: firs message.
-    """
-    async for message in broker.listen():  # noqa: WPS328
-        return message
+from tests.utils import read_message
 
 
 @pytest.mark.anyio
@@ -40,6 +30,8 @@ async def test_success_broadcast(nats_urls: List[str], nats_subject: str) -> Non
     for received_message in await asyncio.wait_for(asyncio.gather(*tasks), timeout=1):
         assert received_message == sent_message.message
 
+    await broker.shutdown()
+
 
 @pytest.mark.anyio
 async def test_success_queued(nats_urls: List[str], nats_subject: str) -> None:
@@ -58,3 +50,4 @@ async def test_success_queued(nats_urls: List[str], nats_subject: str) -> None:
     )
     asyncio.create_task(broker.kick(sent_message))
     assert await asyncio.wait_for(reading_task, timeout=1) == sent_message.message
+    await broker.shutdown()

tests/test_broker.py

@@ -0,0 +1,37 @@
+import asyncio
+import uuid
+from typing import List
+
+import pytest
+from taskiq import BrokerMessage
+
+from taskiq_nats import JetStreamBroker
+from tests.utils import read_message
+
+
+@pytest.mark.anyio
+async def test_success(nats_urls: List[str], nats_subject: str) -> None:
+    """
+    Tests that JetStream works.
+
+    This function sends a message to JetStream
+    before starting to listen to it.
+
+    It epexts to receive the same message.
+    """
+    broker = JetStreamBroker(
+        servers=nats_urls,
+        subject=nats_subject,
+        queue=uuid.uuid4().hex,
+        stream_name=uuid.uuid4().hex,
+    )
+    await broker.startup()
+    sent_message = BrokerMessage(
+        task_id=uuid.uuid4().hex,
+        task_name="meme",
+        message=b"some",
+        labels={},
+    )
+    await broker.kick(sent_message)
+    assert await asyncio.wait_for(read_message(broker), 0.5) == sent_message.message
+    await broker.shutdown()

tests/test_jstream.py

@@ -0,0 +1,13 @@
+from taskiq import AsyncBroker
+
+
+async def read_message(broker: AsyncBroker) -> bytes:  # type: ignore
+    """
+    Read signle message from the broker's listen method.
+
+    :param broker: current broker.
+    :return: firs message.
+    """
+    async for message in broker.listen():
+        assert isinstance(message, bytes)  # noqa: S101
+        return message