diff --git a/redisvl/extensions/cache/base.py b/redisvl/extensions/cache/base.py index 8ff3566db..f14abb86e 100644 --- a/redisvl/extensions/cache/base.py +++ b/redisvl/extensions/cache/base.py @@ -4,7 +4,6 @@ specific cache types such as LLM caches and embedding caches. """ -from collections.abc import Mapping from typing import Any, cast from redis import Redis # For backwards compatibility in type checking @@ -13,6 +12,10 @@ from redisvl.redis.connection import RedisConnectionFactory from redisvl.types import AsyncRedisClient, SyncRedisClient +# Keys deleted per DEL when clearing. Also the SCAN count hint, so one page of +# keys maps to one delete round-trip. +CLEAR_BATCH_SIZE = 500 + class BaseCache: """Base abstract cache interface for all RedisVL caches. @@ -182,50 +185,74 @@ async def aexpire(self, key: str, ttl: int | None = None) -> None: await client.expire(key, _ttl) def clear(self) -> None: - """Clear the cache of all keys.""" + """Clear the cache of all keys. + + Deletes every Redis key under the cache's prefix (``:``) with + ``SCAN`` + ``DEL``. The cache object itself stays usable for future + writes. + + Note: + ``SCAN`` is not a point-in-time snapshot, so this is a best-effort + sweep rather than an atomic flush: + + - Keys written by other clients while the sweep is in progress may + or may not be deleted, so the cache is not guaranteed to be empty + when this returns. Quiesce writers first if you need that. + - ``SCAN`` may return the same key on more than one page. ``DEL`` + on an already-deleted key is a no-op, so this is harmless. + - Deletion is not atomic across keys. If the call raises partway + through, some keys are already gone. The operation is idempotent, + so retrying is safe and converges. + """ client = self._get_redis_client() - prefix = self._get_prefix() - - # Scan for all keys with our prefix - cursor = 0 # Start with cursor 0 - while True: - cursor_int, keys = client.scan(cursor=cursor, match=f"{prefix}*", count=100) # type: ignore - if keys: - client.delete(*keys) - if cursor_int == 0: # Redis returns 0 when scan is complete - break - # Cluster returns a dict of cursor values. We need to stop if these all - # come back as 0. - elif isinstance(cursor_int, Mapping): - cursor_values = list(cursor_int.values()) - if all(v == 0 for v in cursor_values): - break - else: - cursor = cursor_int # Update cursor for next iteration + # scan_iter, not a hand-rolled SCAN loop: on a cluster client SCAN is + # broadcast to every primary and replies with a {node_name: cursor} + # mapping, and those cursors are node-local -- they can neither be fed + # back as a single cursor nor broadcast. redis-py's scan_iter already + # drives each primary on its own cursor via target_nodes. + batch: list[Any] = [] + for key in client.scan_iter( + match=f"{self._get_prefix()}*", count=CLEAR_BATCH_SIZE + ): + batch.append(key) + if len(batch) >= CLEAR_BATCH_SIZE: + client.delete(*batch) + batch.clear() + if batch: + client.delete(*batch) async def aclear(self) -> None: - """Async clear the cache of all keys.""" + """Asynchronously clear the cache of all keys. + + Deletes every Redis key under the cache's prefix (``:``) with + ``SCAN`` + ``DEL``. The cache object itself stays usable for future + writes. + + Note: + ``SCAN`` is not a point-in-time snapshot, so this is a best-effort + sweep rather than an atomic flush: + + - Keys written by other clients while the sweep is in progress may + or may not be deleted, so the cache is not guaranteed to be empty + when this returns. Quiesce writers first if you need that. + - ``SCAN`` may return the same key on more than one page. ``DEL`` + on an already-deleted key is a no-op, so this is harmless. + - Deletion is not atomic across keys. If the call raises partway + through, some keys are already gone. The operation is idempotent, + so retrying is safe and converges. + """ client = await self._get_async_redis_client() - prefix = self._get_prefix() - - # Scan for all keys with our prefix - cursor = 0 # Start with cursor 0 - while True: - cursor_int, keys = await client.scan( - cursor=cursor, match=f"{prefix}*", count=100 - ) # type: ignore - if keys: - await client.delete(*keys) - if cursor_int == 0: # Redis returns 0 when scan is complete - break - # Cluster returns a dict of cursor values. We need to stop if these all - # come back as 0. - elif isinstance(cursor_int, Mapping): - cursor_values = list(cursor_int.values()) - if all(v == 0 for v in cursor_values): - break - else: - cursor = cursor_int # Update cursor for next iteration + # See the note in clear() on why this delegates to scan_iter. + batch: list[Any] = [] + async for key in client.scan_iter( + match=f"{self._get_prefix()}*", count=CLEAR_BATCH_SIZE + ): + batch.append(key) + if len(batch) >= CLEAR_BATCH_SIZE: + await client.delete(*batch) + batch.clear() + if batch: + await client.delete(*batch) def disconnect(self) -> None: """Disconnect from Redis.""" diff --git a/redisvl/migration/async_executor.py b/redisvl/migration/async_executor.py index 149ae0e9d..6a7326ebc 100644 --- a/redisvl/migration/async_executor.py +++ b/redisvl/migration/async_executor.py @@ -229,21 +229,15 @@ async def _enumerate_with_scan( for match_pattern in build_scan_match_patterns( normalized_prefixes, key_separator ): - cursor: int = 0 - while True: - cursor, keys = await client.scan( - cursor=cursor, - match=match_pattern, - count=batch_size, - ) - for key in keys: - key_str = key.decode() if isinstance(key, bytes) else str(key) - if key_str not in seen_keys: - seen_keys.add(key_str) - yield key_str - - if cursor == 0: - break + # scan_iter, not a hand-rolled SCAN loop: a cluster client replies + # with a {node_name: cursor} mapping, which cannot be fed back as a + # cursor (redis-py raises DataError). scan_iter drives each primary + # on its own cursor. + async for key in client.scan_iter(match=match_pattern, count=batch_size): + key_str = key.decode() if isinstance(key, bytes) else str(key) + if key_str not in seen_keys: + seen_keys.add(key_str) + yield key_str async def _rename_keys( self, diff --git a/redisvl/migration/async_planner.py b/redisvl/migration/async_planner.py index 6c75efda2..7bd9a154b 100644 --- a/redisvl/migration/async_planner.py +++ b/redisvl/migration/async_planner.py @@ -1,6 +1,7 @@ from __future__ import annotations -from typing import Any, List, Optional +from contextlib import aclosing +from typing import Any, AsyncGenerator, List, Optional, cast from redisvl.index import AsyncSearchIndex from redisvl.migration.models import ( @@ -272,21 +273,27 @@ async def _async_sample_keys( match_pattern = f"{prefix}*" else: match_pattern = f"{prefix}{key_separator}*" - cursor: int = 0 - while True: - cursor, keys = await client.scan( - cursor=cursor, + # See the note in the sync planner's _sample_keys on why this + # delegates to scan_iter rather than driving the cursor by hand. + # aclosing because we return mid-iteration once the sample limit is + # hit: an async generator abandoned that way is only closed at loop + # shutdown, which warns that scan_iter's aclose was never awaited. + # scan_iter is annotated AsyncIterator, which does not advertise + # aclose, but every implementation of it is an async generator. + scanner = cast( + AsyncGenerator[Any, None], + client.scan_iter( match=match_pattern, count=max(self.key_sample_limit, 10), - ) - for key in keys: + ), + ) + async with aclosing(scanner) as keys: + async for key in keys: decoded_key = key.decode() if isinstance(key, bytes) else str(key) if decoded_key not in key_sample: key_sample.append(decoded_key) if len(key_sample) >= self.key_sample_limit: return key_sample - if cursor == 0: - break return key_sample def write_plan(self, plan: MigrationPlan, plan_out: str) -> None: diff --git a/redisvl/migration/async_validation.py b/redisvl/migration/async_validation.py index ce742a3d0..02146bfb6 100644 --- a/redisvl/migration/async_validation.py +++ b/redisvl/migration/async_validation.py @@ -149,14 +149,13 @@ async def _count_index_keys(self, index: AsyncSearchIndex) -> int: key_separator = index.schema.index.key_separator seen_keys: set[str] = set() for match_pattern in build_scan_match_patterns(prefix_list, key_separator): - cursor = 0 - while True: - cursor, keys = await client.scan(cursor=cursor, match=match_pattern) - for key in keys: - key_str = key.decode() if isinstance(key, bytes) else str(key) - seen_keys.add(key_str) - if cursor == 0: - break + # scan_iter, not a hand-rolled SCAN loop: a cluster client replies + # with a {node_name: cursor} mapping, which cannot be fed back as a + # cursor (redis-py raises DataError). scan_iter drives each primary + # on its own cursor. + async for key in client.scan_iter(match=match_pattern): + key_str = key.decode() if isinstance(key, bytes) else str(key) + seen_keys.add(key_str) return len(seen_keys) async def _run_query_checks( diff --git a/redisvl/migration/executor.py b/redisvl/migration/executor.py index a4f8ae3d0..06b8f3f05 100644 --- a/redisvl/migration/executor.py +++ b/redisvl/migration/executor.py @@ -386,21 +386,15 @@ def _enumerate_with_scan( for match_pattern in build_scan_match_patterns( normalized_prefixes, key_separator ): - cursor = 0 - while True: - cursor, keys = client.scan( # type: ignore[misc] - cursor=cursor, - match=match_pattern, - count=batch_size, - ) - for key in keys: - key_str = key.decode() if isinstance(key, bytes) else str(key) - if key_str not in seen_keys: - seen_keys.add(key_str) - yield key_str - - if cursor == 0: - break + # scan_iter, not a hand-rolled SCAN loop: a cluster client replies + # with a {node_name: cursor} mapping, which cannot be fed back as a + # cursor (redis-py raises DataError). scan_iter drives each primary + # on its own cursor. + for key in client.scan_iter(match=match_pattern, count=batch_size): + key_str = key.decode() if isinstance(key, bytes) else str(key) + if key_str not in seen_keys: + seen_keys.add(key_str) + yield key_str def _rename_keys( self, diff --git a/redisvl/migration/planner.py b/redisvl/migration/planner.py index 4c09fe04c..f0419aca2 100644 --- a/redisvl/migration/planner.py +++ b/redisvl/migration/planner.py @@ -664,21 +664,20 @@ def _sample_keys( # key_separator — a PREFIX of "doc" must match "doc:1", # "doca:1", etc., exactly like FT.CREATE does. match_pattern = f"{prefix}*" - cursor = 0 - while True: - cursor, keys = client.scan( - cursor=cursor, - match=match_pattern, - count=max(self.key_sample_limit, 1000), - ) - for key in keys: - decoded_key = key.decode() if isinstance(key, bytes) else str(key) - if decoded_key not in key_sample: - key_sample.append(decoded_key) - if len(key_sample) >= self.key_sample_limit: - return key_sample - if cursor == 0: - break + # scan_iter, not a hand-rolled SCAN loop: a cluster client replies + # with a {node_name: cursor} mapping, which cannot be fed back as a + # cursor (redis-py raises DataError). scan_iter drives each primary + # on its own cursor. It also lets the sample limit below stop us + # mid-page instead of draining the whole page first. + for key in client.scan_iter( + match=match_pattern, + count=max(self.key_sample_limit, 1000), + ): + decoded_key = key.decode() if isinstance(key, bytes) else str(key) + if decoded_key not in key_sample: + key_sample.append(decoded_key) + if len(key_sample) >= self.key_sample_limit: + return key_sample return key_sample def _detect_possible_field_renames( diff --git a/redisvl/migration/validation.py b/redisvl/migration/validation.py index f8735a443..9f68564b7 100644 --- a/redisvl/migration/validation.py +++ b/redisvl/migration/validation.py @@ -138,17 +138,13 @@ def _count_index_keys(self, index: SearchIndex) -> int: key_separator = index.schema.index.key_separator seen_keys: set[str] = set() for match_pattern in build_scan_match_patterns(prefix_list, key_separator): - cursor = 0 - while True: - cursor, keys = cast( - tuple[int, list[Any]], - client.scan(cursor=cursor, match=match_pattern), - ) - for key in keys: - key_str = key.decode() if isinstance(key, bytes) else str(key) - seen_keys.add(key_str) - if cursor == 0: - break + # scan_iter, not a hand-rolled SCAN loop: a cluster client replies + # with a {node_name: cursor} mapping, which cannot be fed back as a + # cursor (redis-py raises DataError). scan_iter drives each primary + # on its own cursor. + for key in client.scan_iter(match=match_pattern): + key_str = key.decode() if isinstance(key, bytes) else str(key) + seen_keys.add(key_str) return len(seen_keys) def _run_query_checks( diff --git a/redisvl/utils/utils.py b/redisvl/utils/utils.py index 85f74397c..e80165cf9 100644 --- a/redisvl/utils/utils.py +++ b/redisvl/utils/utils.py @@ -12,9 +12,10 @@ from warnings import warn from pydantic import BaseModel -from redis import Redis from ulid import ULID +from redisvl.types import SyncRedisClient + T = TypeVar("T") @@ -283,14 +284,19 @@ def norm_l2_distance(value: float) -> float: def scan_by_pattern( - redis_client: Redis, + redis_client: SyncRedisClient, pattern: str, ) -> Sequence[str]: """ Scan the Redis database for keys matching a specific pattern. + Uses scan_iter, so this is correct for both standalone and cluster clients: + on a cluster, SCAN is broadcast to every primary and each one has to be + iterated on its own node-local cursor. + Args: - redis (Redis): The Redis client instance. + redis_client (SyncRedisClient): The Redis client instance. Standalone + or cluster. pattern (str): The pattern to match keys against. Returns: diff --git a/tests/integration/test_redis_cluster_support.py b/tests/integration/test_redis_cluster_support.py index fc01922c0..fb5c1ea92 100644 --- a/tests/integration/test_redis_cluster_support.py +++ b/tests/integration/test_redis_cluster_support.py @@ -1,5 +1,7 @@ """Tests for Redis Cluster support in RedisVL.""" +import asyncio + import pytest from redis import Redis from redis.asyncio.cluster import RedisCluster as AsyncRedisCluster @@ -164,7 +166,7 @@ async def test_embeddings_cache_cluster_async(redis_cluster_url, redis_test_name try: await cache.aset( - text="hey", + content="hey", model_name="test", embedding=[1, 2, 3], ) @@ -172,6 +174,7 @@ async def test_embeddings_cache_cluster_async(redis_cluster_url, redis_test_name assert result is not None assert result["embedding"] == [1, 2, 3] await cache.aclear() + assert await cache.aget("hey", "test") is None finally: # Manually close the cluster client to prevent connection leaks await cluster_client.aclose() @@ -187,7 +190,7 @@ def test_embeddings_cache_cluster_sync(redis_cluster_url, redis_test_name): for i in range(100): cache.set( - text=f"hey_{i}", + content=f"hey_{i}", model_name="test", embedding=[1, 2, 3], ) @@ -195,11 +198,12 @@ def test_embeddings_cache_cluster_sync(redis_cluster_url, redis_test_name): assert result is not None assert result["embedding"] == [1, 2, 3] cache.clear() + assert cache.get("hey_0", "test") is None cache.mset( [ - {"text": "hey_0", "model_name": "test", "embedding": [1, 2, 3]}, - {"text": "hey_1", "model_name": "test", "embedding": [1, 2, 3]}, + {"content": "hey_0", "model_name": "test", "embedding": [1, 2, 3]}, + {"content": "hey_1", "model_name": "test", "embedding": [1, 2, 3]}, ] ) result = cache.mget(["hey_0", "hey_1"], "test") @@ -208,6 +212,7 @@ def test_embeddings_cache_cluster_sync(redis_cluster_url, redis_test_name): assert result[0]["embedding"] == [1, 2, 3] assert result[1]["embedding"] == [1, 2, 3] cache.clear() + assert cache.mget(["hey_0", "hey_1"], "test") == [None, None] @pytest.mark.requires_cluster @@ -247,3 +252,174 @@ def test_semantic_router_cluster_client( if router._index and router._index.exists(): router._index.delete(drop=True) + + +# ============================================================================= +# BaseCache.clear/aclear on a cluster keyspace larger than one SCAN page. +# +# On a cluster, SCAN is broadcast to all primaries and replies with a +# {node_name: cursor} mapping of node-local cursors. The clear loop used to +# leave its cursor at 0, so it re-issued SCAN 0 forever. It accidentally made +# progress when every key in the DB matched the cache prefix -- deleting the +# first page shrank the keyspace -- which is why small-cache tests passed. The +# genuine hang needs unrelated keys in the DB, the normal case for redisvl, +# where index docs and caches share a keyspace: then a SCAN 0 page can match +# nothing, nothing gets deleted, and the loop spins with zero progress forever. +# +# These tests therefore seed BOTH unrelated keys and a multi-page cache, and +# bound the SCAN count during clear() so a regression fails loudly instead of +# hanging the suite (pytest-timeout is not installed). +# +# Note: count keys with scan_iter, never KEYS or DBSIZE. Those are routed to a +# single node on a cluster, so they silently report roughly one shard's worth. +# ============================================================================= + +CLEAR_NOISE_KEYS = 2000 +CLEAR_CACHE_KEYS = 600 +MAX_CLEAR_SCANS = 400 + + +def _count_keys(client, pattern): + return sum(1 for _ in client.scan_iter(match=pattern, count=500)) + + +async def _acount_keys(client, pattern): + total = 0 + async for _ in client.scan_iter(match=pattern, count=500): + total += 1 + return total + + +def _drop_keys(client, pattern): + keys = list(client.scan_iter(match=pattern, count=500)) + if keys: + client.delete(*keys) + + +async def _adrop_keys(client, pattern): + keys = [k async for k in client.scan_iter(match=pattern, count=500)] + if keys: + await client.delete(*keys) + + +@pytest.mark.requires_cluster +def test_embeddings_cache_clear_multipage_cluster(redis_cluster_url, redis_test_name): + """clear() empties a multi-page cache on a cluster and spares other keys.""" + cluster_client = RedisCluster.from_url(redis_cluster_url) + name = redis_test_name("clear_multipage") + noise_prefix = redis_test_name("clear_noise") + cache = EmbeddingsCache(name=name, redis_client=cluster_client) + + try: + # Unrelated keys: without these the buggy loop accidentally terminates. + pipe = cluster_client.pipeline() + for i in range(CLEAR_NOISE_KEYS): + pipe.set(f"{noise_prefix}:{i}", "keep") + pipe.execute() + + cache.mset( + [ + { + "content": f"content-{i}", + "model_name": "test", + "embedding": [0.1, 0.2, 0.3], + } + for i in range(CLEAR_CACHE_KEYS) + ] + ) + assert _count_keys(cluster_client, f"{name}:*") == CLEAR_CACHE_KEYS + + # Bound the SCAN calls clear() itself issues: the pre-fix loop never + # terminates, and pytest-timeout is not installed. + real_scan = cluster_client.scan + calls: list = [] + + def counting_scan(*args, **kwargs): + calls.append(kwargs.get("cursor", args[0] if args else None)) + if len(calls) > MAX_CLEAR_SCANS: + raise AssertionError( + f"clear() issued {len(calls)} SCAN calls without finishing; " + f"first cursors: {calls[:8]}" + ) + return real_scan(*args, **kwargs) + + cluster_client.scan = counting_scan + try: + cache.clear() + finally: + cluster_client.scan = real_scan + + # The real-world bug: keys past the first page survived. + assert _count_keys(cluster_client, f"{name}:*") == 0 + # clear() must not touch anything outside its own prefix. + assert _count_keys(cluster_client, f"{noise_prefix}:*") == CLEAR_NOISE_KEYS + # Guard against a vacuous pass: paging must actually have happened. + assert len(calls) > 1, "one SCAN sufficed; raise CLEAR_CACHE_KEYS" + finally: + _drop_keys(cluster_client, f"{name}:*") + _drop_keys(cluster_client, f"{noise_prefix}:*") + cluster_client.close() + + +@pytest.mark.requires_cluster +@pytest.mark.asyncio +async def test_embeddings_cache_aclear_multipage_cluster( + redis_cluster_url, redis_test_name +): + """aclear() must match clear(): aclear is a separate code path.""" + cluster_client = RedisConnectionFactory.get_async_redis_cluster_connection( + redis_cluster_url + ) + name = redis_test_name("aclear_multipage") + noise_prefix = redis_test_name("aclear_noise") + cache = EmbeddingsCache(name=name, async_redis_client=cluster_client) + + try: + pipe = cluster_client.pipeline() + for i in range(CLEAR_NOISE_KEYS): + pipe.set(f"{noise_prefix}:{i}", "keep") + await pipe.execute() + + # Seeded with aset, not amset: amset silently writes nothing on an + # async cluster client (it awaits the pipeline returned by the queueing + # call, which drains the queue). That is a separate bug; this test is + # about aclear, so don't let it depend on the broken path. + await asyncio.gather( + *( + cache.aset( + content=f"content-{i}", + model_name="test", + embedding=[0.1, 0.2, 0.3], + ) + for i in range(CLEAR_CACHE_KEYS) + ) + ) + assert await _acount_keys(cluster_client, f"{name}:*") == CLEAR_CACHE_KEYS + + real_scan = cluster_client.scan + calls: list = [] + + async def counting_scan(*args, **kwargs): + calls.append(kwargs.get("cursor", args[0] if args else None)) + if len(calls) > MAX_CLEAR_SCANS: + raise AssertionError( + f"aclear() issued {len(calls)} SCAN calls without finishing; " + f"first cursors: {calls[:8]}" + ) + return await real_scan(*args, **kwargs) + + cluster_client.scan = counting_scan + try: + await cache.aclear() + finally: + cluster_client.scan = real_scan + + assert await _acount_keys(cluster_client, f"{name}:*") == 0 + assert ( + await _acount_keys(cluster_client, f"{noise_prefix}:*") == CLEAR_NOISE_KEYS + ) + assert len(calls) > 1, "one SCAN sufficed; raise CLEAR_CACHE_KEYS" + finally: + await _adrop_keys(cluster_client, f"{name}:*") + await _adrop_keys(cluster_client, f"{noise_prefix}:*") + await cluster_client.aclose() diff --git a/tests/unit/test_async_migration_planner.py b/tests/unit/test_async_migration_planner.py index 93ce3d49d..005e4f65a 100644 --- a/tests/unit/test_async_migration_planner.py +++ b/tests/unit/test_async_migration_planner.py @@ -7,6 +7,7 @@ import pytest import yaml +from redis.asyncio.client import Redis as AsyncRedis from redisvl.migration import AsyncMigrationPlanner, MigrationPlanner from redisvl.schema.schema import IndexSchema @@ -15,10 +16,14 @@ class AsyncDummyClient: """Async mock Redis client for testing.""" + # Bind redis-py's real scan_iter so key enumeration under test goes through + # the actual library loop rather than a stand-in for it. + scan_iter = AsyncRedis.scan_iter + def __init__(self, keys): self.keys = keys - async def scan(self, cursor=0, match=None, count=None): + async def scan(self, cursor=0, match=None, count=None, _type=None): matched = [] for key in self.keys: decoded_key = key.decode() if isinstance(key, bytes) else str(key) diff --git a/tests/unit/test_batch_migration.py b/tests/unit/test_batch_migration.py index 210c2cb47..3ee2dcefe 100644 --- a/tests/unit/test_batch_migration.py +++ b/tests/unit/test_batch_migration.py @@ -14,6 +14,7 @@ import pytest import yaml +from redis.client import Redis from redisvl.migration import ( BatchMigrationExecutor, @@ -33,6 +34,10 @@ class MockRedisClient: """Mock Redis client for batch migration tests.""" + # Bind redis-py's real scan_iter so key enumeration under test goes through + # the actual library loop rather than a stand-in for it. + scan_iter = Redis.scan_iter + def __init__(self, indexes: List[str] = None, keys: Dict[str, List[str]] = None): self.indexes = indexes or [] self.keys = keys or {} @@ -43,7 +48,7 @@ def execute_command(self, *args, **kwargs): return [idx.encode() for idx in self.indexes] raise NotImplementedError(f"Command not mocked: {args}") - def scan(self, cursor=0, match=None, count=None): + def scan(self, cursor=0, match=None, count=None, _type=None): matched = [] all_keys = [] for prefix_keys in self.keys.values(): diff --git a/tests/unit/test_cache_clear_cluster_cursor.py b/tests/unit/test_cache_clear_cluster_cursor.py new file mode 100644 index 000000000..1dba05fa3 --- /dev/null +++ b/tests/unit/test_cache_clear_cluster_cursor.py @@ -0,0 +1,186 @@ +"""Unit tests for BaseCache.clear/aclear enumeration on Redis Cluster. + +On a cluster client, ``scan`` is broadcast to every primary and replies with a +``{node_name: cursor}`` mapping of node-local cursors. The clear loop used to +leave its cursor at 0 in that case, so it re-issued ``SCAN 0`` forever and made +no progress once the first page stopped yielding matches. + +``clear`` now delegates enumeration to redis-py's ``scan_iter``, so the cursor +walk itself is upstream's code and not worth re-asserting here. What these tests +pin is our part: every primary gets drained, the match pattern stays scoped to +the cache prefix, and deletes are batched rather than issued per page or all at +once. The fakes bind the real upstream ``scan_iter`` so the drain runs through +the actual library loop. + +Standalone clear/aclear is covered end-to-end against a real Redis by +tests/integration/test_llmcache.py; the cluster path is covered against a real +cluster by tests/integration/test_redis_cluster_support.py, which only runs +under --run-cluster-tests. +""" + +import pytest +from redis.asyncio.cluster import RedisCluster as AsyncRedisCluster +from redis.cluster import RedisCluster + +from redisvl.extensions.cache import base as cache_base +from redisvl.extensions.cache.base import BaseCache + +PREFIX = "clear_cursor_test" +MATCH = f"{PREFIX}:*" + + +class MockClusterClient: + """Stand-in for RedisCluster SCAN/DEL, driven by a per-node key layout. + + Keys live on named primaries and each primary pages through its own keys on + its own cursor, which is what makes the cluster contract testable. + + Two built-in guards, so failures are loud rather than silent: + ``scan`` raises if a node-local cursor is ever broadcast or if the call + count runs away (a non-advancing cursor would otherwise hang the suite -- + pytest-timeout is not installed), and ``delete`` raises on a zero-argument + call, which real Redis rejects with "wrong number of arguments". + """ + + # Exercise the real upstream loop instead of imitating it. + scan_iter = RedisCluster.scan_iter + + def __init__(self, keys_by_node, page=2, max_calls=200): + self.keys_by_node = {n: list(k) for n, k in keys_by_node.items()} + self.page = page + self.max_calls = max_calls + self.scan_calls = [] + self.delete_batches = [] + + @property + def deleted(self): + return [key for batch in self.delete_batches for key in batch] + + def _page(self, node, cursor): + keys = self.keys_by_node[node] + nxt = cursor + self.page + return (0 if nxt >= len(keys) else nxt), keys[cursor : cursor + self.page] + + def _next_scan(self, cursor, match, target_nodes): + self.scan_calls.append( + {"cursor": cursor, "match": match, "target_nodes": target_nodes} + ) + if len(self.scan_calls) > self.max_calls: + raise AssertionError( + f"{len(self.scan_calls)} SCAN calls -- the cursor is not " + f"advancing: {self.scan_calls[:6]}..." + ) + if target_nodes is None: + # Broadcast to all primaries. Only cursor 0 may be broadcast; a + # node-local cursor sent here would resume other nodes mid-keyspace. + assert cursor == 0, f"node-local cursor {cursor!r} broadcast to primaries" + pages = {n: self._page(n, 0) for n in self.keys_by_node} + return ( + {n: c for n, (c, _) in pages.items()}, + [k for _, ks in pages.values() for k in ks], + ) + node = str(target_nodes).removeprefix("node-object:") + nxt, keys = self._page(node, cursor) + return {node: nxt}, keys + + def _record_delete(self, keys): + if not keys: + raise AssertionError("DEL issued with no keys; real Redis rejects this") + self.delete_batches.append(list(keys)) + return len(keys) + + def scan(self, cursor=0, match=None, count=None, target_nodes=None, _type=None): + return self._next_scan(cursor, match, target_nodes) + + def delete(self, *keys): + return self._record_delete(keys) + + def get_node(self, host=None, port=None, node_name=None): + return f"node-object:{node_name}" + + +class MockAsyncClusterClient(MockClusterClient): + """Async variant, bound to the async cluster client's ``scan_iter``.""" + + scan_iter = AsyncRedisCluster.scan_iter + + async def scan( + self, cursor=0, match=None, count=None, target_nodes=None, _type=None + ): + return self._next_scan(cursor, match, target_nodes) + + async def delete(self, *keys): + return self._record_delete(keys) + + +def _layout(): + """Three primaries with uneven scan depth, one done after the broadcast.""" + return { + "node-1": [f"{PREFIX}:a{i}" for i in range(6)], # 3 rounds + "node-2": [f"{PREFIX}:b{i}" for i in range(2)], # done on the broadcast + "node-3": [f"{PREFIX}:c{i}" for i in range(4)], # 2 rounds + } + + +def _assert_drained(client, layout): + """Every primary emptied, nothing scanned outside the cache prefix.""" + assert sorted(client.deleted) == sorted(k for ks in layout.values() for k in ks) + # Scoping the match pattern is what keeps clear() from wiping the whole DB. + assert all(c["match"] == MATCH for c in client.scan_calls) + # Anti-vacuity: a single broadcast page would prove nothing about paging. + assert any( + c["target_nodes"] is not None for c in client.scan_calls + ), "no per-node continuation happened; test proves nothing" + # node-2 finished on the broadcast and must not be revisited. + assert not any(c["target_nodes"] == "node-object:node-2" for c in client.scan_calls) + + +class TestClearOnCluster: + def test_drains_every_primary(self): + layout = _layout() + client = MockClusterClient(layout) + + BaseCache(name=PREFIX, redis_client=client).clear() + + _assert_drained(client, layout) + + def test_empty_cache_issues_no_delete(self): + client = MockClusterClient({"node-1": [], "node-2": []}) + + BaseCache(name=PREFIX, redis_client=client).clear() + + assert client.delete_batches == [] + + def test_deletes_are_batched(self, monkeypatch): + # Patch the batch size rather than asserting against its real value, so + # tuning CLEAR_BATCH_SIZE in production doesn't touch this test. + monkeypatch.setattr(cache_base, "CLEAR_BATCH_SIZE", 3) + keys = [f"{PREFIX}:{i}" for i in range(8)] + client = MockClusterClient({"node-1": keys}, page=3) + + BaseCache(name=PREFIX, redis_client=client).clear() + + # Bounded batches, and the trailing remainder is still flushed. + assert [len(b) for b in client.delete_batches] == [3, 3, 2] + assert sorted(client.deleted) == sorted(keys) + + +class TestAsyncClearOnCluster: + """aclear is written separately from clear, so it needs its own coverage.""" + + @pytest.mark.asyncio + async def test_drains_every_primary(self): + layout = _layout() + client = MockAsyncClusterClient(layout) + + await BaseCache(name=PREFIX, async_redis_client=client).aclear() + + _assert_drained(client, layout) + + @pytest.mark.asyncio + async def test_empty_cache_issues_no_delete(self): + client = MockAsyncClusterClient({"node-1": [], "node-2": []}) + + await BaseCache(name=PREFIX, async_redis_client=client).aclear() + + assert client.delete_batches == [] diff --git a/tests/unit/test_migration_cluster_scan.py b/tests/unit/test_migration_cluster_scan.py new file mode 100644 index 000000000..984f7e4a0 --- /dev/null +++ b/tests/unit/test_migration_cluster_scan.py @@ -0,0 +1,146 @@ +"""Regression tests for migration SCAN key enumeration on Redis Cluster. + +Six migration call sites used to feed the previous reply's cursor straight back +into ``client.scan(cursor=...)``. On a cluster client that reply is a +``{node_name: cursor}`` mapping, so ``cursor == 0`` is never true and the second +iteration passes a dict as the cursor -- redis-py rejects it with ``DataError``. +That is a hard crash on any cluster, not a silent under-count. + +The fake below reproduces exactly that: it replies with per-node cursors like a +real cluster and raises ``DataError`` if handed a non-integer cursor, so the old +loop fails here the same way it fails against real Redis. + +Covers the four always-reachable sites (validator key counts, planner key +sampling; sync and async). The two executor sites +(``_enumerate_with_scan``) are the same one-line ``scan_iter`` pattern but sit +behind index-info mocking, and they only run on fallback paths; they are left to +the migration integration tests. +""" + +import pytest +from redis.asyncio.cluster import RedisCluster as AsyncRedisCluster +from redis.cluster import RedisCluster +from redis.exceptions import DataError + +from redisvl.migration.async_planner import AsyncMigrationPlanner +from redisvl.migration.async_validation import AsyncMigrationValidator +from redisvl.migration.planner import MigrationPlanner +from redisvl.migration.validation import MigrationValidator + +PREFIX = "doc" +NODES = { + "node-1": [f"{PREFIX}:a{i}".encode() for i in range(6)], + "node-2": [f"{PREFIX}:b{i}".encode() for i in range(2)], + "node-3": [f"{PREFIX}:c{i}".encode() for i in range(4)], +} +ALL_KEYS = sorted(k.decode() for ks in NODES.values() for k in ks) + + +class MockClusterClient: + """Cluster-shaped SCAN: per-node cursors, and DataError on a dict cursor.""" + + scan_iter = RedisCluster.scan_iter + + def __init__(self, keys_by_node=None, page=2, max_calls=200): + self.keys_by_node = {n: list(k) for n, k in (keys_by_node or NODES).items()} + self.page = page + self.max_calls = max_calls + self.scan_calls = 0 + + def _page(self, node, cursor): + keys = self.keys_by_node[node] + nxt = cursor + self.page + return (0 if nxt >= len(keys) else nxt), keys[cursor : cursor + self.page] + + def _next_scan(self, cursor, target_nodes): + self.scan_calls += 1 + if self.scan_calls > self.max_calls: + raise AssertionError(f"{self.scan_calls} SCAN calls -- not advancing") + # What redis-py's encoder does when the old loop hands back the mapping. + if not isinstance(cursor, (int, str, bytes)): + raise DataError( + f"Invalid input of type: {type(cursor).__name__!r}. " + "Convert to a bytes, string, int or float first." + ) + cursor = int(cursor) + if target_nodes is None: + assert cursor == 0, f"node-local cursor {cursor!r} broadcast to primaries" + pages = {n: self._page(n, 0) for n in self.keys_by_node} + return ( + {n: c for n, (c, _) in pages.items()}, + [k for _, ks in pages.values() for k in ks], + ) + node = str(target_nodes).removeprefix("node-object:") + nxt, keys = self._page(node, cursor) + return {node: nxt}, keys + + def scan(self, cursor=0, match=None, count=None, target_nodes=None, _type=None): + return self._next_scan(cursor, target_nodes) + + def get_node(self, host=None, port=None, node_name=None): + return f"node-object:{node_name}" + + +class MockAsyncClusterClient(MockClusterClient): + scan_iter = AsyncRedisCluster.scan_iter + + async def scan( + self, cursor=0, match=None, count=None, target_nodes=None, _type=None + ): + return self._next_scan(cursor, target_nodes) + + +class _Index: + """Minimal SearchIndex stand-in for the validators' key-count path.""" + + def __init__(self, client): + self.client = client + self.schema = type( + "S", + (), + {"index": type("I", (), {"prefix": PREFIX, "key_separator": ":"})()}, + )() + + +class TestValidatorCountsKeysOnCluster: + """_count_index_keys runs on every validate, so this is the hottest path.""" + + def test_counts_across_all_primaries(self): + client = MockClusterClient() + validator = MigrationValidator.__new__(MigrationValidator) + + assert validator._count_index_keys(_Index(client)) == len(ALL_KEYS) + + @pytest.mark.asyncio + async def test_counts_across_all_primaries_async(self): + client = MockAsyncClusterClient() + validator = AsyncMigrationValidator.__new__(AsyncMigrationValidator) + + assert await validator._count_index_keys(_Index(client)) == len(ALL_KEYS) + + +class TestPlannerSamplesKeysOnCluster: + """_sample_keys returns early at the sample limit, mid-page.""" + + def test_samples_across_primaries(self): + client = MockClusterClient() + planner = MigrationPlanner.__new__(MigrationPlanner) + planner.key_sample_limit = len(ALL_KEYS) + + sample = planner._sample_keys( + client=client, prefixes=[PREFIX], key_separator=":" + ) + + assert sorted(sample) == ALL_KEYS + + @pytest.mark.asyncio + async def test_samples_across_primaries_async(self): + client = MockAsyncClusterClient() + planner = AsyncMigrationPlanner.__new__(AsyncMigrationPlanner) + planner.key_sample_limit = len(ALL_KEYS) + + sample = await planner._async_sample_keys( + client=client, prefixes=[PREFIX], key_separator=":" + ) + + assert sorted(sample) == ALL_KEYS diff --git a/tests/unit/test_migration_planner.py b/tests/unit/test_migration_planner.py index b07f9df93..5bd52b264 100644 --- a/tests/unit/test_migration_planner.py +++ b/tests/unit/test_migration_planner.py @@ -2,6 +2,7 @@ from unittest.mock import MagicMock import yaml +from redis.client import Redis from redisvl.migration import MigrationPlanner from redisvl.migration.executor import _extract_prefixes_from_info @@ -17,10 +18,14 @@ class DummyClient: + # Bind redis-py's real scan_iter so key enumeration under test goes through + # the actual library loop rather than a stand-in for it. + scan_iter = Redis.scan_iter + def __init__(self, keys): self.keys = keys - def scan(self, cursor=0, match=None, count=None): + def scan(self, cursor=0, match=None, count=None, _type=None): matched = [] for key in self.keys: decoded_key = key.decode() if isinstance(key, bytes) else str(key)