Merge remote-tracking branch 'dfroger/fix-pipeline-ttl' into fix/dfroger/fix-pipeline-ttl

Author: vishal-bala

redisvl/index/storage.py

@@ -175,6 +175,20 @@ async def _aget(
         """Asynchronously get data from Redis using the provided client or pipeline."""
         raise NotImplementedError
 
+    @staticmethod
+    async def _aexpire(client: AsyncRedisClientOrPipeline, key: str, ttl: int):
+        """Asynchronously set TTL on a key using the provided client or pipeline
+
+        Args:
+            client (AsyncRedisClientOrPipeline): The async Redis client or pipeline instance.
+            key (str): The key for which to set the TTL.
+            ttl (int): Time-to-live in seconds for each key.
+        """
+        if isinstance(client, (AsyncPipeline, AsyncClusterPipeline)):
+            client.expire(key, ttl)
+        else:
+            await client.expire(key, ttl)
+
     def _validate(self, obj: Dict[str, Any]) -> Dict[str, Any]:
         """
         Validate an object against the schema using Pydantic-based validation.
@@ -490,7 +504,7 @@ async def awrite(
 
                 # Set TTL if provided
                 if ttl:
-                    await pipe.expire(key, ttl)
+                    await self._aexpire(pipe, key, ttl)
 
                 added_keys.append(key)
 
@@ -615,7 +629,7 @@ async def _aset(client: AsyncRedisClientOrPipeline, key: str, obj: Dict[str, Any
         """Asynchronously set a hash value in Redis for the given key.
 
         Args:
-            client (AsyncClientOrPipeline): The async Redis client or pipeline instance.
+            client (AsyncRedisClientOrPipeline): The async Redis client or pipeline instance.
             key (str): The key under which to store the hash.
             obj (Dict[str, Any]): The hash to store in Redis.
         """
@@ -644,7 +658,7 @@ async def _aget(
         """Asynchronously retrieve a hash value from Redis for the given key.
 
         Args:
-            client (AsyncRedisClient): The async Redis client or pipeline instance.
+            client (AsyncRedisClientOrPipeline): The async Redis client or pipeline instance.
             key (str): The key for which to retrieve the hash.
 
         Returns:

tests/integration/test_async_search_index.py

@@ -707,3 +707,26 @@ async def test_async_search_index_connect(index_schema, redis_url):
         await async_index.connect(redis_url=redis_url)
     assert async_index.client is not None
     await async_index.disconnect()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("ttl", [None, 30])
+async def test_search_index_load_with_ttl(async_index, ttl):
+    """Test that TTL is correctly set on keys when using load() with ttl parameter."""
+    await async_index.create(overwrite=True, drop=True)
+
+    # Load test data with TTL parameter
+    data = [{"id": "1", "test": "foo"}]
+    keys = await async_index.load(data, id_field="id", ttl=ttl)
+
+    # Check TTL on the loaded key
+    client = await async_index._get_client()
+    key_ttl = await client.ttl(keys[0])
+
+    if ttl is None:
+        # No TTL set, should return -1
+        assert key_ttl == -1
+    else:
+        # TTL should be set and close to the expected value
+        assert key_ttl > 0
+        assert abs(key_ttl - ttl) <= 5

tests/integration/test_cluster_pipelining.py

@@ -152,3 +152,42 @@ def test_batch_search_with_real_cluster(redis_cluster_url):
 
     finally:
         index.delete()
+
+
+@pytest.mark.requires_cluster
+@pytest.mark.parametrize("ttl", [None, 30])
+def test_cluster_load_with_ttl(redis_cluster_url, ttl):
+    """
+    Test that TTL is correctly set on keys when using load() with ttl parameter on cluster.
+    """
+    schema_dict = {
+        "index": {"name": "test-ttl-cluster", "prefix": "ttl", "storage_type": "hash"},
+        "fields": [
+            {"name": "id", "type": "tag"},
+            {"name": "text", "type": "text"},
+        ],
+    }
+
+    schema = IndexSchema.from_dict(schema_dict)
+    index = SearchIndex(schema, redis_url=redis_cluster_url)
+
+    index.create(overwrite=True)
+
+    try:
+        # Load test data with TTL parameter
+        data = [{"id": "1", "text": "foo"}]
+        keys = index.load(data, id_field="id", ttl=ttl)
+
+        # Check TTL on the loaded key
+        key_ttl = index.client.ttl(keys[0])
+
+        if ttl is None:
+            # No TTL set, should return -1
+            assert key_ttl == -1
+        else:
+            # TTL should be set and close to the expected value
+            assert key_ttl > 0
+            assert abs(key_ttl - ttl) <= 5
+
+    finally:
+        index.delete()

tests/integration/test_search_index.py

@@ -697,3 +697,24 @@ def test_search_index_validates_query_with_hnsw_algorithm(hnsw_index, sample_dat
     )
     # Should not raise
     hnsw_index.query(query)
+
+
+@pytest.mark.parametrize("ttl", [None, 30])
+def test_search_index_load_with_ttl(index, ttl):
+    """Test that TTL is correctly set on keys when using load() with ttl parameter."""
+    index.create(overwrite=True, drop=True)
+
+    # Load test data with TTL parameter
+    data = [{"id": "1", "test": "foo"}]
+    keys = index.load(data, id_field="id", ttl=ttl)
+
+    # Check TTL on the loaded key
+    key_ttl = index.client.ttl(keys[0])
+
+    if ttl is None:
+        # No TTL set, should return -1
+        assert key_ttl == -1
+    else:
+        # TTL should be set and close to the expected value
+        assert key_ttl > 0
+        assert abs(key_ttl - ttl) <= 5