Add per-message TTL support for KV operations (#783)

* Add per-message TTL support for KV operations

---------

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

Author: caspervonb

nats/src/nats/js/client.py

@@ -1403,6 +1403,7 @@ async def create_key_value(
             subjects=[f"$KV.{config.bucket}.>"],
             allow_direct=config.direct,
             allow_rollup_hdrs=True,
+            allow_msg_ttl=True,
             deny_delete=True,
             discard=api.DiscardPolicy.NEW,
             duplicate_window=duplicate_window,

nats/src/nats/js/kv.py

@@ -193,23 +193,34 @@ async def put(self, key: str, value: bytes, validate_keys: bool = True) -> int:
         """
         put will place the new value for the key into the store
         and return the revision number.
+
+        Note: This method does not support TTL. Use create() if you need TTL support.
+
+        :param key: The key to put
+        :param value: The value to store
+        :param validate_keys: Whether to validate the key format
         """
         if validate_keys and not _is_key_valid(key):
             raise nats.js.errors.InvalidKeyError(key)
 
         pa = await self._js.publish(f"{self._pre}{key}", value)
         return pa.seq
 
-    async def create(self, key: str, value: bytes, validate_keys: bool = True) -> int:
+    async def create(self, key: str, value: bytes, validate_keys: bool = True, msg_ttl: Optional[float] = None) -> int:
         """
         create will add the key/value pair iff it does not exist.
+
+        :param key: The key to create
+        :param value: The value to store
+        :param validate_keys: Whether to validate the key format
+        :param msg_ttl: Optional TTL (time-to-live) in seconds for this specific message
         """
         if validate_keys and not _is_key_valid(key):
             raise nats.js.errors.InvalidKeyError(key)
 
         pa = None
         try:
-            pa = await self.update(key, value, last=0, validate_keys=validate_keys)
+            pa = await self.update(key, value, last=0, validate_keys=validate_keys, msg_ttl=msg_ttl)
         except nats.js.errors.KeyWrongLastSequenceError as err:
             # In case of attempting to recreate an already deleted key,
             # the client would get a KeyWrongLastSequenceError.  When this happens,
@@ -229,13 +240,25 @@ async def create(self, key: str, value: bytes, validate_keys: bool = True) -> in
                 # to recreate using the last revision.
                 raise err
             except nats.js.errors.KeyDeletedError as err:
-                pa = await self.update(key, value, last=err.entry.revision, validate_keys=validate_keys)
+                pa = await self.update(
+                    key, value, last=err.entry.revision, validate_keys=validate_keys, msg_ttl=msg_ttl
+                )
 
         return pa
 
-    async def update(self, key: str, value: bytes, last: Optional[int] = None, validate_keys: bool = True) -> int:
+    async def update(
+        self,
+        key: str,
+        value: bytes,
+        last: Optional[int] = None,
+        validate_keys: bool = True,
+        msg_ttl: Optional[float] = None,
+    ) -> int:
         """
         update will update the value if the latest revision matches.
+
+        Note: TTL parameter is accepted for internal use by create(), but should not be
+        used directly on update operations per NATS KV semantics.
         """
         if validate_keys and not _is_key_valid(key):
             raise nats.js.errors.InvalidKeyError(key)
@@ -247,7 +270,7 @@ async def update(self, key: str, value: bytes, last: Optional[int] = None, valid
 
         pa = None
         try:
-            pa = await self._js.publish(f"{self._pre}{key}", value, headers=hdrs)
+            pa = await self._js.publish(f"{self._pre}{key}", value, headers=hdrs, msg_ttl=msg_ttl)
         except nats.js.errors.APIError as err:
             # Check for a BadRequest::KeyWrongLastSequenceError error code.
             if err.err_code == 10071:
@@ -256,9 +279,16 @@ async def update(self, key: str, value: bytes, last: Optional[int] = None, valid
                 raise err
         return pa.seq
 
-    async def delete(self, key: str, last: Optional[int] = None, validate_keys: bool = True) -> bool:
+    async def delete(
+        self, key: str, last: Optional[int] = None, validate_keys: bool = True, msg_ttl: Optional[float] = None
+    ) -> bool:
         """
         delete will place a delete marker and remove all previous revisions.
+
+        :param key: The key to delete
+        :param last: Expected last revision number (for optimistic concurrency)
+        :param validate_keys: Whether to validate the key format
+        :param msg_ttl: Optional TTL (time-to-live) in seconds for the delete marker
         """
         if validate_keys and not _is_key_valid(key):
             raise nats.js.errors.InvalidKeyError(key)
@@ -269,17 +299,20 @@ async def delete(self, key: str, last: Optional[int] = None, validate_keys: bool
         if last and last > 0:
             hdrs[api.Header.EXPECTED_LAST_SUBJECT_SEQUENCE] = str(last)
 
-        await self._js.publish(f"{self._pre}{key}", headers=hdrs)
+        await self._js.publish(f"{self._pre}{key}", headers=hdrs, msg_ttl=msg_ttl)
         return True
 
-    async def purge(self, key: str) -> bool:
+    async def purge(self, key: str, msg_ttl: Optional[float] = None) -> bool:
         """
         purge will remove the key and all revisions.
+
+        :param key: The key to purge
+        :param msg_ttl: Optional TTL (time-to-live) in seconds for the purge marker
         """
         hdrs = {}
         hdrs[KV_OP] = KV_PURGE
         hdrs[api.Header.ROLLUP] = MSG_ROLLUP_SUBJECT
-        await self._js.publish(f"{self._pre}{key}", headers=hdrs)
+        await self._js.publish(f"{self._pre}{key}", headers=hdrs, msg_ttl=msg_ttl)
         return True
 
     async def purge_deletes(self, olderthan: int = 30 * 60) -> bool:

nats/tests/test_js.py

@@ -3830,6 +3830,185 @@ async def error_handler(e):
         # Clean up
         await nc.close()
 
+    @async_test
+    async def test_kv_create_with_ttl(self):
+        """Test that create() supports msg_ttl parameter"""
+        errors = []
+
+        async def error_handler(e):
+            print("Error:", e, type(e))
+            errors.append(e)
+
+        nc = await nats.connect(error_cb=error_handler)
+
+        server_version = nc.connected_server_version
+        if server_version.major == 2 and server_version.minor < 11:
+            pytest.skip("per-message TTL requires nats-server v2.11.0 or later")
+
+        js = nc.jetstream()
+
+        # Create a KV bucket
+        kv = await js.create_key_value(bucket="TEST_TTL_CREATE", history=5)
+
+        # Create a key with TTL of 2 seconds
+        seq = await kv.create("age", b"30", msg_ttl=2.0)
+        assert seq == 1
+
+        # Key should exist immediately
+        entry = await kv.get("age")
+        assert entry.key == "age"
+        assert entry.value == b"30"
+        assert entry.revision == 1
+
+        # Wait for TTL to expire (2 seconds + buffer)
+        await asyncio.sleep(2.5)
+
+        # Key should be gone after TTL expires
+        with pytest.raises(KeyNotFoundError):
+            await kv.get("age")
+
+        await nc.close()
+
+    @async_test
+    async def test_kv_purge_with_ttl(self):
+        """Test that purge() supports msg_ttl parameter for the purge marker"""
+        errors = []
+
+        async def error_handler(e):
+            print("Error:", e, type(e))
+            errors.append(e)
+
+        nc = await nats.connect(error_cb=error_handler)
+
+        server_version = nc.connected_server_version
+        if server_version.major == 2 and server_version.minor < 11:
+            pytest.skip("per-message TTL requires nats-server v2.11.0 or later")
+
+        js = nc.jetstream()
+
+        # Create a KV bucket
+        kv = await js.create_key_value(bucket="TEST_TTL_PURGE", history=10)
+
+        # Put a key
+        seq = await kv.put("name", b"alice")
+        assert seq == 1
+
+        # Purge with TTL of 2 seconds on the purge marker
+        await kv.purge("name", msg_ttl=2.0)
+
+        # Key should be purged immediately
+        with pytest.raises(KeyNotFoundError):
+            await kv.get("name")
+
+        # The purge marker should exist in the stream
+        # We can verify by checking stream info - there should be a message
+        status = await kv.status()
+        # After purge, there should still be a marker message
+        assert status.values >= 1
+
+        # Wait for the purge marker TTL to expire (2 seconds + buffer)
+        await asyncio.sleep(2.5)
+
+        # The marker itself should now be removed from the stream
+        # Note: This behavior depends on server version and configuration
+
+        await nc.close()
+
+    @async_test
+    async def test_kv_delete_with_ttl(self):
+        """Test that delete() supports msg_ttl parameter for the delete marker"""
+        errors = []
+
+        async def error_handler(e):
+            print("Error:", e, type(e))
+            errors.append(e)
+
+        nc = await nats.connect(error_cb=error_handler)
+
+        server_version = nc.connected_server_version
+        if server_version.major == 2 and server_version.minor < 11:
+            pytest.skip("per-message TTL requires nats-server v2.11.0 or later")
+
+        js = nc.jetstream()
+
+        # Create a KV bucket
+        kv = await js.create_key_value(bucket="TEST_TTL_DELETE", history=10)
+
+        # Put a key
+        seq = await kv.put("city", b"paris")
+        assert seq == 1
+
+        # Verify the key exists
+        entry = await kv.get("city")
+        assert entry.value == b"paris"
+
+        # Delete with TTL of 2 seconds on the delete marker
+        await kv.delete("city", msg_ttl=2.0)
+
+        # Key should be deleted immediately
+        with pytest.raises(KeyNotFoundError):
+            await kv.get("city")
+
+        # The delete marker should exist in the stream
+        status = await kv.status()
+        # After delete, there should be both the original message and delete marker
+        assert status.values >= 1
+
+        # Wait for the delete marker TTL to expire (2 seconds + buffer)
+        await asyncio.sleep(2.5)
+
+        # The marker itself should now be removed from the stream
+        # Note: This behavior depends on server version and configuration
+
+        await nc.close()
+
+    @async_test
+    async def test_kv_put_no_ttl(self):
+        """Test that put() does NOT support TTL (should not have msg_ttl parameter)"""
+        nc = await nats.connect()
+        js = nc.jetstream()
+
+        # Create a KV bucket
+        kv = await js.create_key_value(bucket="TEST_NO_TTL_PUT", history=5)
+
+        # Put should work normally without TTL
+        seq = await kv.put("key1", b"value1")
+        assert seq == 1
+
+        # Verify put() method signature doesn't accept msg_ttl
+        # This is a compile-time check - if this test compiles, the signature is correct
+        import inspect
+
+        sig = inspect.signature(kv.put)
+        params = list(sig.parameters.keys())
+        assert "msg_ttl" not in params, "put() should not accept msg_ttl parameter"
+
+        await nc.close()
+
+    @async_test
+    async def test_kv_update_no_direct_ttl(self):
+        """Test that update() does not expose TTL for direct use"""
+        nc = await nats.connect()
+        js = nc.jetstream()
+
+        # Create a KV bucket
+        kv = await js.create_key_value(bucket="TEST_NO_TTL_UPDATE", history=5)
+
+        # Put initial value
+        seq = await kv.put("counter", b"1")
+        assert seq == 1
+
+        # Update should work normally
+        seq = await kv.update("counter", b"2", last=1)
+        assert seq == 2
+
+        # While update() technically has msg_ttl parameter for internal use by create(),
+        # it's documented as not for direct use
+        entry = await kv.get("counter")
+        assert entry.value == b"2"
+
+        await nc.close()
+
 
 class ObjectStoreTest(SingleJetStreamServerTestCase):
     @async_test