diff --git a/redisvl/migration/async_executor.py b/redisvl/migration/async_executor.py index 6a7326eb..6090ab0a 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 @@ -819,9 +819,8 @@ async def apply( source_failures = int( plan.source.stats_snapshot.get("hash_indexing_failures", 0) or 0 ) - source_percent_indexed = float( - plan.source.stats_snapshot.get("percent_indexed", 1.0) or 1.0 - ) + progress = plan.source.stats_snapshot.get("percent_indexed") + source_percent_indexed = float(progress) if progress is not None else 1.0 needs_exact_count = source_failures > 0 or source_percent_indexed < 1.0 needs_enumeration = ( needs_quantization diff --git a/redisvl/migration/executor.py b/redisvl/migration/executor.py index 06b8f3f0..84cef40d 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: + # Skip the leading RESP2 metadata; the remaining rows are 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: @@ -1092,9 +1107,8 @@ def apply( source_failures = int( plan.source.stats_snapshot.get("hash_indexing_failures", 0) or 0 ) - source_percent_indexed = float( - plan.source.stats_snapshot.get("percent_indexed", 1.0) or 1.0 - ) + progress = plan.source.stats_snapshot.get("percent_indexed") + source_percent_indexed = float(progress) if progress is not None else 1.0 needs_exact_count = source_failures > 0 or source_percent_indexed < 1.0 needs_enumeration = ( needs_quantization diff --git a/redisvl/migration/planner.py b/redisvl/migration/planner.py index e867f3d2..623bf606 100644 --- a/redisvl/migration/planner.py +++ b/redisvl/migration/planner.py @@ -171,9 +171,8 @@ def create_plan_from_patch( # falls back to SCAN automatically, but surface the condition here # so users running `rvl migrate plan` can wait for indexing to # complete before applying. - source_percent_indexed = float( - snapshot.stats_snapshot.get("percent_indexed", 1.0) or 1.0 - ) + progress = snapshot.stats_snapshot.get("percent_indexed") + source_percent_indexed = float(progress) if progress is not None else 1.0 if source_percent_indexed < 1.0: warnings.append( f"Source index is still building " diff --git a/tests/unit/test_async_migration_executor.py b/tests/unit/test_async_migration_executor.py index 53c54cb8..2d97703b 100644 --- a/tests/unit/test_async_migration_executor.py +++ b/tests/unit/test_async_migration_executor.py @@ -8,12 +8,17 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from redis.asyncio.client import Redis as AsyncRedis +from redis.client import Redis from redisvl.migration import AsyncMigrationExecutor, MigrationExecutor +from redisvl.migration import async_executor as async_executor_module +from redisvl.migration import executor as executor_module from redisvl.migration.models import ( DiffClassification, KeyspaceSnapshot, MigrationPlan, + MigrationValidation, SourceSnapshot, ValidationPolicy, ) @@ -102,6 +107,85 @@ def test_async_executor_with_validator(): assert executor.validator is custom_validator +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize( + "progress, expected_count", + [("0", 2), (0.0, 2), (1.0, None), (None, None)], + ids=["resp2-zero", "resp3-zero", "complete", "null-progress"], +) +@pytest.mark.asyncio +async def test_schema_only_migration_counts_unindexed_keys( + monkeypatch, tmp_path, is_async, progress, expected_count +): + plan = _make_basic_plan() + plan.source.stats_snapshot = { + "num_docs": 0, + "hash_indexing_failures": 0, + "percent_indexed": progress, + } + executor = AsyncMigrationExecutor() if is_async else MigrationExecutor() + module = async_executor_module if is_async else executor_module + index_class = module.AsyncSearchIndex if is_async else module.SearchIndex + mock_call = AsyncMock if is_async else MagicMock + + client = MagicMock() + client.info = mock_call(return_value={"aof_enabled": 0}) + client.ft.return_value.info = mock_call( + return_value={ + **plan.source.stats_snapshot, + "index_definition": {"prefixes": ["test:"]}, + } + ) + client.scan = mock_call(return_value=(0, [b"test:1", b"test:2"])) + redis_class = AsyncRedis if is_async else Redis + client.scan_iter = redis_class.scan_iter.__get__(client) + source_index = MagicMock() + source_index._redis_client = client + source_index._get_client = AsyncMock(return_value=client) + source_index.delete = mock_call() + target_index = MagicMock() + target_index.create = mock_call() + monkeypatch.setattr( + index_class, "from_existing", mock_call(return_value=source_index) + ) + monkeypatch.setattr(index_class, "from_dict", MagicMock(return_value=target_index)) + matches = mock_call(side_effect=[True, False]) + wait = mock_call(return_value=({"num_docs": 2}, 0.01)) + if is_async: + monkeypatch.setattr(executor, "_async_current_source_matches_snapshot", matches) + monkeypatch.setattr(executor, "_async_wait_for_index_ready", wait) + else: + monkeypatch.setattr(module, "current_source_matches_snapshot", matches) + monkeypatch.setattr(module, "wait_for_index_ready", wait) + validate = mock_call( + return_value=( + MigrationValidation( + schema_match=True, doc_count_match=True, key_sample_exists=True + ), + {"num_docs": 2}, + 0.01, + ) + ) + monkeypatch.setattr(executor.validator, "validate", validate) + + report = executor.apply( + plan, redis_client=client, backup_dir=str(tmp_path / "backups") + ) + if is_async: + report = await report + + assert report.result == "succeeded", report.validation.errors + # A schema-only change must count the source keys before dropping an + # unfinished index; num_docs == 0 is not the number of stored documents. + assert validate.call_args.kwargs["expected_source_count"] == expected_count + if expected_count is None: + client.scan.assert_not_called() + else: + client.scan.assert_called_once_with( + cursor="0", match="test:*", count=1000, _type=None + ) + + @pytest.mark.asyncio async def test_async_multi_worker_requires_redis_url_before_loading_index(tmp_path): """num_workers > 1 with redis_client only must fail before source lookup.""" diff --git a/tests/unit/test_migration_enumeration.py b/tests/unit/test_migration_enumeration.py new file mode 100644 index 00000000..1eb3404b --- /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.asyncio.client import Redis as AsyncRedis +from redis.client import Redis +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() + redis_class = AsyncRedis if request.param else Redis + client.scan_iter = redis_class.scan_iter.__get__(client) + 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: + # The leading metadata value need not equal the number of rows. + rows = [1, *[["__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, readiness", + [ + (False, {"hash_indexing_failures": 2, "percent_indexed": 1.0}), + (True, {"hash_indexing_failures": 2, "percent_indexed": 1.0}), + (True, {"hash_indexing_failures": 0, "percent_indexed": 0.5}), + (True, {"hash_indexing_failures": 0, "percent_indexed": 0.0}), + ], + ids=["byte-keys", "failed-documents", "partial-index", "zero-progress"], +) +@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)), + ] + assert await _collect(executor, client) == [ + "archive:1", + "archive:failed", + "doc:pending", + ] + assert client.scan.call_args_list == [ + call(cursor="0", match="archive:*", count=2, _type=None), + call(cursor=5, match="archive:*", count=2, _type=None), + call(cursor="0", match="doc:*", count=2, _type=None), + ] + # The fast path would omit failed/pending documents, even if it did not crash. + 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, _type=None) + + +@pytest.mark.asyncio +async def test_closing_enumeration_releases_cursor(migration): + executor, client = migration + client.execute_command.return_value = _aggregate_page(["doc:1"], 17, 3, 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") diff --git a/tests/unit/test_migration_planner.py b/tests/unit/test_migration_planner.py index 5bd52b26..7f3d19e2 100644 --- a/tests/unit/test_migration_planner.py +++ b/tests/unit/test_migration_planner.py @@ -1,6 +1,7 @@ from fnmatch import fnmatch from unittest.mock import MagicMock +import pytest import yaml from redis.client import Redis @@ -1023,12 +1024,13 @@ def test_plan_no_warning_when_stats_missing_failures_key(monkeypatch, tmp_path): assert len(failure_warnings) == 0 -def test_plan_warns_when_source_is_still_indexing(monkeypatch, tmp_path): +@pytest.mark.parametrize("progress", ["0.42", 0.42, "0", 0.0]) +def test_plan_warns_when_source_is_still_indexing(monkeypatch, tmp_path, progress): """Plan should warn when the source index has percent_indexed < 1.0.""" source_schema = _make_source_schema() dummy_index = DummyIndex( source_schema, - {"num_docs": 100, "hash_indexing_failures": 0, "percent_indexed": "0.42"}, + {"num_docs": 100, "hash_indexing_failures": 0, "percent_indexed": progress}, [b"docs:1"], ) monkeypatch.setattr( @@ -1060,7 +1062,7 @@ def test_plan_warns_when_source_is_still_indexing(monkeypatch, tmp_path): indexing_warnings = [w for w in plan.warnings if "still building" in w] assert len(indexing_warnings) == 1 - assert "0.4200" in indexing_warnings[0] + assert f"{float(progress):.4f}" in indexing_warnings[0] def test_plan_no_warning_when_source_fully_indexed(monkeypatch, tmp_path):