From d55117ff8aafc3b774df8f47c351eeb442e4b264 Mon Sep 17 00:00:00 2001 From: Kuang-xianxin <243476082+Kuang-xianxin@users.noreply.github.com> Date: Sun, 6 Sep 2026 22:37:36 +0800 Subject: [PATCH] fix(migration): support RESP3 key enumeration Normalize readiness and prefix metadata, preserve zero indexing progress, and share RESP2/RESP3 aggregate parsing across sync and async executors. Fixes #713 Fixes #714 Assisted-by: Codex --- redisvl/migration/async_executor.py | 14 +- redisvl/migration/executor.py | 29 +++- tests/unit/test_migration_enumeration.py | 173 +++++++++++++++++++++++ 3 files changed, 202 insertions(+), 14 deletions(-) create mode 100644 tests/unit/test_migration_enumeration.py diff --git a/redisvl/migration/async_executor.py b/redisvl/migration/async_executor.py index 149ae0e9d..1f4582795 100644 --- a/redisvl/migration/async_executor.py +++ b/redisvl/migration/async_executor.py @@ -22,6 +22,7 @@ _checkpoint_identity_matches, _delete_backup_prefix, _delete_multi_worker_backup_prefix, + _extract_aggregate_keys, _extract_prefixes_from_info, _key_prefix_map, _map_key_prefix, @@ -45,6 +46,7 @@ normalize_keys, timestamp_utc, ) +from redisvl.redis.utils import convert_bytes from redisvl.types import AsyncRedisClient from redisvl.utils.log import get_logger @@ -103,9 +105,10 @@ async def _enumerate_indexed_keys( # condition means FT.AGGREGATE would miss documents, so fall # back to SCAN for complete enumeration. try: - info = await client.ft(index_name).info() + info = convert_bytes(await client.ft(index_name).info()) failures = int(info.get("hash_indexing_failures", 0) or 0) - percent_indexed = float(info.get("percent_indexed", 1.0) or 1.0) + progress = info.get("percent_indexed") + percent_indexed = float(progress) if progress is not None else 1.0 if failures > 0: logger.warning( f"Index '{index_name}' has {failures} indexing failures. " @@ -183,11 +186,8 @@ async def _enumerate_with_aggregate( while True: results_data, cursor_id = result - # Extract keys from results - for item in results_data[1:]: - if isinstance(item, (list, tuple)) and len(item) >= 2: - key = item[1] - yield key.decode() if isinstance(key, bytes) else str(key) + for key in _extract_aggregate_keys(results_data): + yield key if cursor_id == 0: break diff --git a/redisvl/migration/executor.py b/redisvl/migration/executor.py index a4f8ae3d0..3a355ab8b 100644 --- a/redisvl/migration/executor.py +++ b/redisvl/migration/executor.py @@ -37,6 +37,7 @@ wait_for_index_ready, ) from redisvl.migration.validation import MigrationValidator +from redisvl.redis.utils import convert_bytes from redisvl.types import SyncRedisClient from redisvl.utils.log import get_logger @@ -144,6 +145,7 @@ def _map_keys_prefix( def _extract_prefixes_from_info(info: Any) -> List[str]: """Extract Redis Search index prefixes from dict or list FT.INFO shapes.""" + info = convert_bytes(info) def _prefixes_from_definition(definition: Any) -> Any: if isinstance(definition, dict): @@ -216,6 +218,22 @@ def _checkpoint_identity_matches( ) +def _extract_aggregate_keys(results_data: Any) -> Generator[str, None, None]: + """Read keys from raw RESP2 rows or a RESP3 aggregate result map.""" + results_data = convert_bytes(results_data) + if isinstance(results_data, dict): + keys = (row["extra_attributes"]["__key"] for row in results_data["results"]) + else: + # RESP2 starts with the row count, followed by field/value pairs. + keys = ( + row[1] + for row in results_data[1:] + if isinstance(row, (list, tuple)) and len(row) >= 2 + ) + for key in keys: + yield key.decode() if isinstance(key, bytes) else str(key) + + class MigrationExecutor: def __init__(self, validator: Optional[MigrationValidator] = None): self.validator = validator or MigrationValidator() @@ -250,9 +268,10 @@ def _enumerate_indexed_keys( # condition means FT.AGGREGATE would miss documents, so fall # back to SCAN for complete enumeration. try: - info = client.ft(index_name).info() + info = convert_bytes(client.ft(index_name).info()) failures = int(info.get("hash_indexing_failures", 0) or 0) - percent_indexed = float(info.get("percent_indexed", 1.0) or 1.0) + progress = info.get("percent_indexed") + percent_indexed = float(progress) if progress is not None else 1.0 if failures > 0: logger.warning( f"Index '{index_name}' has {failures} indexing failures. " @@ -331,11 +350,7 @@ def _enumerate_with_aggregate( while True: results_data, cursor_id = result - # Extract keys from results (skip first element which is count) - for item in results_data[1:]: - if isinstance(item, (list, tuple)) and len(item) >= 2: - key = item[1] - yield key.decode() if isinstance(key, bytes) else str(key) + yield from _extract_aggregate_keys(results_data) # Check if done (cursor_id == 0) if cursor_id == 0: diff --git a/tests/unit/test_migration_enumeration.py b/tests/unit/test_migration_enumeration.py new file mode 100644 index 000000000..3b2b56308 --- /dev/null +++ b/tests/unit/test_migration_enumeration.py @@ -0,0 +1,173 @@ +"""Migration key enumeration across Redis wire response formats.""" + +from unittest.mock import AsyncMock, MagicMock, call + +import pytest +from redis.exceptions import ResponseError + +from redisvl.migration import AsyncMigrationExecutor, MigrationExecutor + + +@pytest.fixture(params=[False, True], ids=["sync", "async"]) +def migration(request): + client = MagicMock() + executor = MigrationExecutor() + if request.param: + executor = AsyncMigrationExecutor() + client.ft.return_value.info = AsyncMock() + client.execute_command = AsyncMock() + client.scan = AsyncMock() + return executor, client + + +async def _collect(executor, client): + keys = executor._enumerate_indexed_keys(client, "source", batch_size=2) + if isinstance(executor, AsyncMigrationExecutor): + return [key async for key in keys] + return list(keys) + + +def _wire(value, decode_responses): + if isinstance(value, str): + return value if decode_responses else value.encode() + if isinstance(value, dict): + return { + _wire(key, decode_responses): _wire(item, decode_responses) + for key, item in value.items() + } + if isinstance(value, list): + return [_wire(item, decode_responses) for item in value] + return value + + +def _aggregate_page(keys, cursor, protocol, decode_responses): + if protocol == 2: + rows = [len(keys), *[["__key", key] for key in keys]] + else: + rows = { + "attributes": [], + "format": "STRING", + "results": [ + {"extra_attributes": {"__key": key}, "values": []} for key in keys + ], + "total_results": len(keys), + "warning": [], + } + return [_wire(rows, decode_responses), cursor] + + +@pytest.mark.parametrize("protocol", [2, 3]) +@pytest.mark.parametrize("decode_responses", [False, True]) +@pytest.mark.asyncio +async def test_enumerate_aggregate_cursor_pages(migration, protocol, decode_responses): + executor, client = migration + client.ft.return_value.info.return_value = _wire( + {"hash_indexing_failures": 0, "percent_indexed": 1.0}, decode_responses + ) + client.execute_command.side_effect = [ + _aggregate_page(["doc:1", "doc:中文"], 17, protocol, decode_responses), + _aggregate_page([], 17, protocol, decode_responses), + _aggregate_page(["doc:3"], 0, protocol, decode_responses), + ] + + assert await _collect(executor, client) == ["doc:1", "doc:中文", "doc:3"] + assert client.execute_command.call_args_list == [ + call( + "FT.AGGREGATE", + "source", + "*", + "LOAD", + "1", + "__key", + "WITHCURSOR", + "COUNT", + "2", + "MAXIDLE", + "300000", + ), + call("FT.CURSOR", "READ", "source", "17", "COUNT", "2"), + call("FT.CURSOR", "READ", "source", "17", "COUNT", "2"), + ] + client.scan.assert_not_called() + + +@pytest.mark.parametrize("decode_responses", [False, True]) +@pytest.mark.parametrize( + "readiness", + [ + {"hash_indexing_failures": 2, "percent_indexed": 1.0}, + {"hash_indexing_failures": 0, "percent_indexed": 0.5}, + {"hash_indexing_failures": 0, "percent_indexed": 0.0}, + ], + ids=["failed-documents", "partial-index", "empty-index"], +) +@pytest.mark.asyncio +async def test_incomplete_index_scans_its_prefixes( + migration, decode_responses, readiness +): + executor, client = migration + client.ft.return_value.info.return_value = _wire( + {**readiness, "index_definition": {"prefixes": ["doc:", "archive:"]}}, + decode_responses, + ) + client.scan.side_effect = [ + (5, _wire(["archive:1"], decode_responses)), + (0, _wire(["archive:1", "archive:failed"], decode_responses)), + (0, _wire(["doc:pending"], decode_responses)), + ] + # The fast path would omit failed/pending documents, even if it did not crash. + client.execute_command.return_value = [[1, [b"__key", b"archive:1"]], 0] + + assert await _collect(executor, client) == [ + "archive:1", + "archive:failed", + "doc:pending", + ] + assert client.scan.call_args_list == [ + call(cursor=0, match="archive:*", count=2), + call(cursor=5, match="archive:*", count=2), + call(cursor=0, match="doc:*", count=2), + ] + client.execute_command.assert_not_called() + + +@pytest.mark.parametrize("decode_responses", [False, True]) +@pytest.mark.asyncio +async def test_aggregate_error_preserves_scan_prefix(migration, decode_responses): + executor, client = migration + client.ft.return_value.info.return_value = _wire( + { + "hash_indexing_failures": 0, + "percent_indexed": 1.0, + "index_definition": {"prefixes": ["doc:"]}, + }, + decode_responses, + ) + client.execute_command.side_effect = ResponseError("aggregate unavailable") + client.scan.return_value = (0, [b"doc:1"]) + + assert await _collect(executor, client) == ["doc:1"] + client.scan.assert_called_once_with(cursor=0, match="doc:*", count=2) + + +@pytest.mark.parametrize("protocol", [2, 3]) +@pytest.mark.asyncio +async def test_closing_enumeration_releases_cursor(migration, protocol): + executor, client = migration + client.execute_command.return_value = _aggregate_page( + ["doc:1"], 17, protocol, False + ) + keys = executor._enumerate_with_aggregate(client, "source", batch_size=2) + if isinstance(executor, AsyncMigrationExecutor): + try: + assert await anext(keys) == "doc:1" + finally: + await keys.aclose() + else: + try: + assert next(keys) == "doc:1" + finally: + keys.close() + + assert client.execute_command.call_count == 2 + client.execute_command.assert_called_with("FT.CURSOR", "DEL", "source", "17")