From b7f635133a45af95861c409a7bc710fbbc0b1c69 Mon Sep 17 00:00:00 2001 From: Dmitry Anfimov Date: Mon, 28 Sep 2026 16:13:31 +0200 Subject: [PATCH] fix: reuse stream if it already exists --- taskiq_nats/broker.py | 16 +++++++++++- tests/test_jetstream.py | 57 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 72 insertions(+), 1 deletion(-) diff --git a/taskiq_nats/broker.py b/taskiq_nats/broker.py index 29652fc..ea9f4ab 100644 --- a/taskiq_nats/broker.py +++ b/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() diff --git a/tests/test_jetstream.py b/tests/test_jetstream.py index 49a4471..81b615f 100644 --- a/tests/test_jetstream.py +++ b/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()