diff --git a/nats-jetstream/src/nats/jetstream/__init__.py b/nats-jetstream/src/nats/jetstream/__init__.py index a156ae0b..60bdb1fd 100644 --- a/nats-jetstream/src/nats/jetstream/__init__.py +++ b/nats-jetstream/src/nats/jetstream/__init__.py @@ -462,6 +462,53 @@ async def create_stream(self, config: StreamConfig | None = None, /, **kwargs) - info = StreamInfo.from_response(response, strict=self._strict) return Stream(self, config.name, info) + @overload + async def create_or_update_stream(self, config: StreamConfig, /) -> Stream: + """Create or update a stream from a StreamConfig object.""" + ... + + @overload + async def create_or_update_stream(self, *, name: str, **config) -> Stream: + """Create or update a stream with keyword arguments.""" + ... + + async def create_or_update_stream(self, config: StreamConfig | None = None, /, **kwargs) -> Stream: + """Create a stream, or update it if one with the same name exists. + + This is an idempotent operation. It attempts an update first and falls + back to create only when the stream does not yet exist; any other error + is propagated. + + Args: + config: A StreamConfig object (positional-only) + **kwargs: Stream configuration parameters as keyword arguments + + Returns: + The created or updated Stream object + + Raises: + ValueError: If stream name is not provided + JetStreamError: For JetStream API errors + """ + if config is None: + config = StreamConfig.from_kwargs(**kwargs) + + if config.name is None: + raise ValueError("StreamConfig must have a name") + + config_dict = config.to_request() + try: + response = await self._api.stream_update(config.name, **config_dict) + except StreamNotFoundError: + try: + response = await self._api.stream_create(config.name, **config_dict) + except StreamNameAlreadyInUseError: + # Lost a create race with a concurrent caller; the stream + # exists now, so updating keeps the operation idempotent. + response = await self._api.stream_update(config.name, **config_dict) + info = StreamInfo.from_response(response, strict=self._strict) + return Stream(self, config.name, info) + async def update_stream(self, **config) -> StreamInfo: """Update an existing stream. diff --git a/nats-jetstream/tests/test_stream.py b/nats-jetstream/tests/test_stream.py index dfeee335..6156bd08 100644 --- a/nats-jetstream/tests/test_stream.py +++ b/nats-jetstream/tests/test_stream.py @@ -91,6 +91,69 @@ async def test_update_stream_max_msgs_and_subjects(jetstream: JetStream): assert set(updated_info.config.subjects) == {"FOO.*", "BAR.*"} +@pytest.mark.asyncio +async def test_create_or_update_stream_creates_when_absent(jetstream: JetStream): + """create_or_update_stream creates the stream when it does not exist.""" + stream = await jetstream.create_or_update_stream(name="test", subjects=["FOO.*"]) + assert stream.name == "test" + info = await jetstream.get_stream_info("test") + assert info.config.subjects == ["FOO.*"] + + +@pytest.mark.asyncio +async def test_create_or_update_stream_updates_when_present(jetstream: JetStream): + """create_or_update_stream updates an existing stream in place (idempotent).""" + await jetstream.create_stream(name="test", subjects=["FOO.*"], max_msgs=100) + stream = await jetstream.create_or_update_stream(name="test", subjects=["FOO.*", "BAR.*"], max_msgs=200) + assert stream.name == "test" + assert stream.info.config.subjects == ["FOO.*", "BAR.*"] + assert stream.info.config.max_msgs == 200 + + +@pytest.mark.asyncio +async def test_create_or_update_stream_accepts_config_object(jetstream: JetStream): + """create_or_update_stream accepts a StreamConfig positionally, like create_stream.""" + from nats.jetstream import StreamConfig + + stream = await jetstream.create_or_update_stream(StreamConfig(name="test", subjects=["FOO.*"])) + assert stream.name == "test" + + +@pytest.mark.asyncio +async def test_create_or_update_stream_requires_name(jetstream: JetStream): + """create_or_update_stream raises ValueError when no name is provided.""" + with pytest.raises(ValueError): + await jetstream.create_or_update_stream(subjects=["FOO.*"]) + + +@pytest.mark.asyncio +async def test_create_or_update_stream_resolves_create_race(jetstream: JetStream): + """Losing a create race to a concurrent caller still resolves idempotently. + + Simulates the race window by having the first update report the stream + as missing while it actually exists: the fallback create then collides + on the server and must recover with a final update.""" + from nats.jetstream import StreamNotFoundError + + await jetstream.create_stream(name="test", subjects=["FOO.*"]) + + real_update = jetstream._api.stream_update + update_calls = 0 + + async def update_reporting_missing_once(name, /, **kwargs): + nonlocal update_calls + update_calls += 1 + if update_calls == 1: + raise StreamNotFoundError("stream not found") + return await real_update(name, **kwargs) + + jetstream._api.stream_update = update_reporting_missing_once + + stream = await jetstream.create_or_update_stream(name="test", subjects=["FOO.*", "BAR.*"]) + assert update_calls == 2 + assert set(stream.info.config.subjects) == {"FOO.*", "BAR.*"} + + @pytest.mark.asyncio async def test_update_nonexistent_stream_fails(jetstream: JetStream): """Test that updating a non-existent stream fails."""