fix(migration): support RESP3 key enumeration

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

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:

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")