From ae19de12abbacedac2ff98c17c4b91ff0912ac02 Mon Sep 17 00:00:00 2001 From: peterolkhov Date: Mon, 9 Nov 2026 18:53:24 -0500 Subject: [PATCH] fix: paginate filesystem key-value stores by logical key --- .../_file_system/_key_value_store_client.py | 18 +++++----- .../_file_system/test_fs_kvs_client.py | 35 +++++++++++++++++++ tests/unit/storages/test_key_value_store.py | 20 +++++++++++ 3 files changed, 65 insertions(+), 8 deletions(-) diff --git a/src/crawlee/storage_clients/_file_system/_key_value_store_client.py b/src/crawlee/storage_clients/_file_system/_key_value_store_client.py index 984a88651f..ee67037dba 100644 --- a/src/crawlee/storage_clients/_file_system/_key_value_store_client.py +++ b/src/crawlee/storage_clients/_file_system/_key_value_store_client.py @@ -373,11 +373,11 @@ async def iterate_keys( if not self.path_to_kvs.exists(): return - # List and sort all files *inside* a brief lock, then release it immediately: + # List all files *inside* a brief lock, then release it immediately: async with self._lock: - files = sorted(await asyncio.to_thread(lambda: list(self.path_to_kvs.glob('*')))) + files = await asyncio.to_thread(lambda: list(self.path_to_kvs.glob('*'))) - count = 0 + records: list[KeyValueStoreRecordMetadata] = [] for file_path in files: # Skip the main metadata file @@ -391,10 +391,6 @@ async def iterate_keys( # Extract the base key name from the metadata filename key_name = self._decode_key(file_path.name[: -len(f'.{METADATA_FILENAME}')]) - # Apply exclusive_start_key filter if provided - if exclusive_start_key is not None and key_name <= exclusive_start_key: - continue - # Try to read and parse the metadata file try: metadata_content = await asyncio.to_thread(file_path.read_text, encoding='utf-8') @@ -412,10 +408,16 @@ async def iterate_keys( record_metadata = KeyValueStoreRecordMetadata(**metadata_dict) except ValidationError: logger.warning(f'Invalid metadata schema for key "{key_name}", skipping it.') + continue + + if exclusive_start_key is not None and record_metadata.key <= exclusive_start_key: + continue + + records.append(record_metadata) + for count, record_metadata in enumerate(sorted(records, key=lambda record: record.key), start=1): yield record_metadata - count += 1 if limit and count >= limit: break diff --git a/tests/unit/storage_clients/_file_system/test_fs_kvs_client.py b/tests/unit/storage_clients/_file_system/test_fs_kvs_client.py index 6866f59d05..e9caf9055e 100644 --- a/tests/unit/storage_clients/_file_system/test_fs_kvs_client.py +++ b/tests/unit/storage_clients/_file_system/test_fs_kvs_client.py @@ -226,3 +226,38 @@ async def test_data_persistence_across_reopens(configuration: Configuration) -> assert record.value == test_value await reopened_client.drop() + + +async def test_iterate_keys_uses_metadata_key(kvs_client: FileSystemKeyValueStoreClient) -> None: + """Record metadata supplies the logical key for ordering and cursor exclusion.""" + await kvs_client.set_value(key='INPUT.json', value={'record': 'input'}) + await kvs_client.set_value(key='INPUT-1', value='another record') + metadata_path = kvs_client.path_to_kvs / f'INPUT.json.{METADATA_FILENAME}' + metadata = json.loads(metadata_path.read_text(encoding='utf-8')) + metadata['key'] = 'INPUT' + metadata_path.write_text(json.dumps(metadata), encoding='utf-8') + + first_page = [item async for item in kvs_client.iterate_keys(limit=1)] + assert [item.key for item in first_page] == ['INPUT'] + second_page = [item async for item in kvs_client.iterate_keys(exclusive_start_key=first_page[-1].key, limit=1)] + assert [item.key for item in second_page] == ['INPUT-1'] + assert [item async for item in kvs_client.iterate_keys(exclusive_start_key=second_page[-1].key, limit=1)] == [] + + +@pytest.mark.parametrize( + 'metadata_content', + [ + pytest.param('invalid json', id='invalid-json'), + pytest.param('{}', id='invalid-schema'), + ], +) +async def test_iterate_keys_skips_invalid_metadata( + kvs_client: FileSystemKeyValueStoreClient, metadata_content: str, caplog: pytest.LogCaptureFixture +) -> None: + """Invalid record metadata is logged and omitted from iteration.""" + metadata_path = kvs_client.path_to_kvs / f'invalid.{METADATA_FILENAME}' + metadata_path.write_text(metadata_content, encoding='utf-8') + await kvs_client.set_value(key='valid', value='record') + + assert [item.key async for item in kvs_client.iterate_keys()] == ['valid'] + assert 'for key "invalid", skipping it.' in caplog.text diff --git a/tests/unit/storages/test_key_value_store.py b/tests/unit/storages/test_key_value_store.py index 302c8eb30e..970ba84ea5 100644 --- a/tests/unit/storages/test_key_value_store.py +++ b/tests/unit/storages/test_key_value_store.py @@ -240,6 +240,26 @@ async def test_list_keys_with_exclusive_start_key(kvs: KeyValueStore) -> None: assert all_key_names.index(first_returned_key) > all_key_names.index(start_key) +@pytest.mark.parametrize( + 'keys', + [ + pytest.param(['item', 'item-1'], id='prefix'), + pytest.param(['zebra', '\u00e9clair'], id='unicode'), + pytest.param(['a.b.json', 'a/b.json'], id='path'), + pytest.param(['alpha', 'zebra'], id='ascii'), + ], +) +async def test_list_keys_cursor_pages(kvs: KeyValueStore, keys: list[str]) -> None: + """Cursor pages visit every logical key once in sorted order.""" + for key in keys: + await kvs.set_value(key, {'key': key}) + + first_page = await kvs.list_keys(limit=1) + second_page = await kvs.list_keys(exclusive_start_key=first_page[-1].key, limit=1) + assert [item.key for item in first_page + second_page] == sorted(keys) + assert await kvs.list_keys(exclusive_start_key=second_page[-1].key, limit=1) == [] + + async def test_iterate_keys(kvs: KeyValueStore) -> None: """Test iterating over keys in the key-value store.""" # Add some items