Merge pull request #30 from taskiq-python/fix-reuse-stream-if-it-already-exists

fix: reuse stream if it already exists

Author: danfimov

taskiq_nats/broker.py

@@ -7,6 +7,7 @@
 from nats.errors import TimeoutError as NatsTimeoutError
 from nats.js import JetStreamContext
 from nats.js.api import ConsumerConfig, StreamConfig
+from nats.js.errors import BadRequestError
 from taskiq import AckableMessage, AsyncBroker, AsyncResultBackend, BrokerMessage
 
 _T = typing.TypeVar("_T")  # (Too short)
@@ -152,9 +153,22 @@ async def startup(self) -> 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)
+        await self._add_or_reuse_stream()
         await self._startup_consumer()
 
+    async def _add_or_reuse_stream(self) -> None:
+        """Create a stream or reuse an already existing one."""
+        try:
+            await self.js.add_stream(config=self.stream_config)
+        except BadRequestError as exc:
+            if exc.err_code != 10058:
+                raise
+            logger.info(
+                "Stream %s already exists with a different configuration. "
+                "Reusing the existing stream.",
+                self.stream_config.name,
+            )
+
     async def shutdown(self) -> None:
         """Close connections to NATS."""
         await self.client.close()

tests/test_jetstream.py

@@ -1,6 +1,11 @@
 import asyncio
 import uuid
+from unittest.mock import AsyncMock, MagicMock
 
+import pytest
+from nats.js import JetStreamContext
+from nats.js.api import RetentionPolicy, StreamConfig
+from nats.js.errors import BadRequestError
 from taskiq import AckableMessage, BrokerMessage
 
 from taskiq_nats import PullBasedJetStreamBroker, PushBasedJetStreamBroker
@@ -87,3 +92,55 @@ async def test_pull_based_broker_success(
         broker.stream_name,
     )
     await broker.shutdown()
+
+
+async def test_broker_reuses_existing_stream(
+    nats_urls: list[str],
+    nats_subject: str,
+    nats_jetstream: JetStreamContext,
+) -> None:
+    stream_name = uuid.uuid4().hex
+    await nats_jetstream.add_stream(
+        config=StreamConfig(
+            name=stream_name,
+            subjects=[nats_subject],
+            retention=RetentionPolicy.WORK_QUEUE,
+            max_msgs=1000,
+        ),
+    )
+
+    broker = PullBasedJetStreamBroker(
+        servers=nats_urls,
+        subject=nats_subject,
+        stream_name=stream_name,
+    )
+    await broker.startup()
+
+    stream_info = await broker.js.stream_info(stream_name)
+    assert stream_info.config.retention == RetentionPolicy.WORK_QUEUE
+    assert stream_info.config.max_msgs == 1000
+
+    await broker.js.delete_consumer(
+        stream=stream_name,
+        consumer=broker.durable,
+    )
+    await broker.js.delete_stream(stream_name)
+    await broker.shutdown()
+
+
+async def test_broker_startup_reraises_unknown_stream_error(
+    nats_urls: list[str],
+    nats_subject: str,
+) -> None:
+    broker = PullBasedJetStreamBroker(
+        servers=nats_urls,
+        subject=nats_subject,
+        stream_name=uuid.uuid4().hex,
+    )
+    broker.js = MagicMock()
+    broker.js.add_stream = AsyncMock(
+        side_effect=BadRequestError(code=400, err_code=10052),
+    )
+
+    with pytest.raises(BadRequestError):
+        await broker._add_or_reuse_stream()