Add create_or_update_stream to JetStream

nats-jetstream/src/nats/jetstream/__init__.py

@@ -462,6 +462,48 @@ 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:
+            response = await self._api.stream_create(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.
 

nats-jetstream/tests/test_stream.py

@@ -91,6 +91,41 @@ 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_update_nonexistent_stream_fails(jetstream: JetStream):
     """Test that updating a non-existent stream fails."""