Add create_or_update_consumer to JetStream and Stream

Signed-off-by: Casper Beyer <[email protected]>

Author: caspervonb

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

@@ -535,6 +535,28 @@ async def update_consumer(self, stream_name: str, consumer_name: str, **config)
         stream = await self.get_stream(stream_name)
         return await stream.update_consumer(name=consumer_name, **config)
 
+    async def create_or_update_consumer(self, stream_name: str, **config) -> Consumer:
+        """Create or update a consumer.
+
+        This method will either create a consumer if it does not exist or update
+        an existing consumer (if possible). This is an idempotent operation.
+
+        Note: Some consumer configuration fields cannot be updated after creation
+        (e.g., max_waiting must be set during creation).
+
+        Args:
+            stream_name: Name of the stream
+            **config: Consumer configuration (must include 'name')
+
+        Returns:
+            The created or updated consumer
+
+        Raises:
+            ValueError: If 'name' is not provided in config
+        """
+        stream = await self.get_stream(stream_name)
+        return await stream.create_or_update_consumer(**config)
+
     async def consumer_names(self, stream_name: str) -> AsyncIterator[str]:
         """Get an async iterator over all consumer names for a stream.
 

nats-jetstream/src/nats/jetstream/stream.py

@@ -1390,6 +1390,46 @@ async def update_consumer(self, config: ConsumerConfig | None = None, /, **kwarg
         # Delegate to _upsert_consumer with "update" action
         return await self._upsert_consumer(action=CONSUMER_ACTION_UPDATE, **config_dict)
 
+    @overload
+    async def create_or_update_consumer(self, config: ConsumerConfig, /) -> Consumer:
+        """Create or update a consumer from a ConsumerConfig object."""
+        ...
+
+    @overload
+    async def create_or_update_consumer(self, *, name: str, **config) -> Consumer:
+        """Create or update a consumer with keyword arguments."""
+        ...
+
+    async def create_or_update_consumer(self, config: ConsumerConfig | None = None, /, **kwargs) -> Consumer:
+        """Create or update a consumer for this stream.
+
+        This method will either create a consumer if it does not exist or update
+        an existing consumer (if possible). This is an idempotent operation.
+
+        Note: Some consumer configuration fields cannot be updated after creation
+        (e.g., max_waiting must be set during creation).
+
+        Args:
+            config: A ConsumerConfig object (positional-only)
+            **kwargs: Consumer configuration parameters as keyword arguments
+
+        Returns:
+            The created or updated consumer
+        """
+        if config is None:
+            # Create ConsumerConfig from kwargs
+            config = ConsumerConfig.from_kwargs(**kwargs)
+
+        # Validate consumer name
+        if config.name is None:
+            raise ValueError("ConsumerConfig must have a name")
+
+        # Convert ConsumerConfig to API request format
+        config_dict = config.to_request()
+
+        # Delegate to _upsert_consumer with empty action (create-or-update)
+        return await self._upsert_consumer(action=CONSUMER_ACTION_CREATE_OR_UPDATE, **config_dict)
+
     async def pause_consumer(self, consumer_name: str, pause_until: float) -> None:
         """Pause a consumer until a specific time.
 

nats-jetstream/tests/test_jetstream.py

@@ -519,6 +519,56 @@ async def test_update_nonexistent_consumer_fails(jetstream: JetStream):
         await jetstream.update_consumer(stream_name="test_stream", consumer_name="nonexistent", max_deliver=20)
 
 
+@pytest.mark.asyncio
+async def test_create_or_update_consumer_creates_new(jetstream: JetStream):
+    """Test that create_or_update_consumer creates a new consumer when it doesn't exist."""
+    # Create a stream
+    await jetstream.create_stream(name="test_stream", subjects=["FOO.*"])
+
+    # Create a consumer using create_or_update
+    consumer = await jetstream.create_or_update_consumer(
+        stream_name="test_stream", name="test_consumer", durable_name="test_consumer", max_deliver=10
+    )
+
+    assert consumer.info.config.name == "test_consumer"
+    assert consumer.info.config.max_deliver == 10
+
+
+@pytest.mark.asyncio
+async def test_create_or_update_consumer_updates_existing(jetstream: JetStream):
+    """Test that create_or_update_consumer updates an existing consumer."""
+    # Create a stream and consumer
+    await jetstream.create_stream(name="test_stream", subjects=["FOO.*"])
+    await jetstream.create_consumer(
+        stream_name="test_stream", name="test_consumer", durable_name="test_consumer", max_deliver=10
+    )
+
+    # Update the consumer using create_or_update
+    updated_consumer = await jetstream.create_or_update_consumer(
+        stream_name="test_stream", name="test_consumer", durable_name="test_consumer", max_deliver=20
+    )
+
+    assert updated_consumer.info.config.max_deliver == 20
+
+
+@pytest.mark.asyncio
+async def test_create_or_update_consumer_via_stream(jetstream: JetStream):
+    """Test create_or_update_consumer via Stream object."""
+    # Create a stream
+    stream = await jetstream.create_stream(name="test_stream", subjects=["FOO.*"])
+
+    # Create a consumer using create_or_update
+    consumer = await stream.create_or_update_consumer(name="test_consumer", max_deliver=10)
+
+    assert consumer.info.config.name == "test_consumer"
+    assert consumer.info.config.max_deliver == 10
+
+    # Update the same consumer
+    updated_consumer = await stream.create_or_update_consumer(name="test_consumer", max_deliver=20)
+
+    assert updated_consumer.info.config.max_deliver == 20
+
+
 @pytest.mark.asyncio
 async def test_delete_nonexistent_consumer_fails(jetstream: JetStream):
     """Test that deleting a non-existent consumer fails."""