Add Apify's version of FS client that keeps the INPUT json

Author: vdusek

src/apify/storage_clients/__init__.py

@@ -1,6 +1,7 @@
-from crawlee.storage_clients import FileSystemStorageClient, MemoryStorageClient
+from crawlee.storage_clients import MemoryStorageClient
 
 from ._apify import ApifyStorageClient
+from ._file_system import FileSystemStorageClient
 
 __all__ = [
     'ApifyStorageClient',

src/apify/storage_clients/_file_system/__init__.py

@@ -0,0 +1 @@
+from ._storage_client import ApifyFileSystemStorageClient as FileSystemStorageClient

src/apify/storage_clients/_file_system/_key_value_store_client.py

@@ -0,0 +1,36 @@
+import asyncio
+
+from typing_extensions import override
+
+from crawlee._consts import METADATA_FILENAME
+from crawlee.storage_clients._file_system import FileSystemKeyValueStoreClient
+
+from apify._configuration import Configuration
+
+
+class ApifyFileSystemKeyValueStoreClient(FileSystemKeyValueStoreClient):
+    """Apify-specific implementation of the `FileSystemKeyValueStoreClient`.
+
+    The only difference is that it overrides the `purge` method to delete all files in the key-value store
+    directory, except for the metadata file and the `INPUT.json` file.
+    """
+
+    @override
+    async def purge(self) -> None:
+        """Purges the key-value store by deleting all its contents.
+
+        It deletes all files in the key-value store directory, except for the metadata file and
+        the `INPUT.json` file. It also updates the metadata to reflect that the store has been purged.
+        """
+        kvs_input_key = Configuration.get_global_configuration().input_key
+        async with self._lock:
+            for file_path in self.path_to_kvs.glob('*'):
+                if file_path.name in {METADATA_FILENAME, f'{kvs_input_key}.json'}:
+                    continue
+                if file_path.is_file():
+                    await asyncio.to_thread(file_path.unlink, missing_ok=True)
+
+            await self._update_metadata(
+                update_accessed_at=True,
+                update_modified_at=True,
+            )

src/apify/storage_clients/_file_system/_storage_client.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from typing_extensions import override
+
+from crawlee._utils.docs import docs_group
+from crawlee.configuration import Configuration
+from crawlee.storage_clients import FileSystemStorageClient
+
+from ._key_value_store_client import ApifyFileSystemKeyValueStoreClient
+
+if TYPE_CHECKING:
+    from crawlee.storage_clients._file_system import FileSystemKeyValueStoreClient
+
+
+@docs_group('Classes')
+class ApifyFileSystemStorageClient(FileSystemStorageClient):
+    """Apify-specific implementation of the file system storage client.
+
+    The only difference is that it uses `ApifyFileSystemKeyValueStoreClient` for key-value stores,
+    which overrides the `purge` method to delete all files in the key-value store directory
+    except for the metadata file and the `INPUT.json` file.
+    """
+
+    @override
+    async def create_kvs_client(
+        self,
+        *,
+        id: str | None = None,
+        name: str | None = None,
+        configuration: Configuration | None = None,
+    ) -> FileSystemKeyValueStoreClient:
+        configuration = configuration or Configuration.get_global_configuration()
+        client = await ApifyFileSystemKeyValueStoreClient.open(id=id, name=name, configuration=configuration)
+        await self._purge_if_needed(client, configuration)
+        return client

tests/unit/storage_clients/__init__.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+import asyncio
+
+from crawlee._consts import METADATA_FILENAME
+
+from apify import Configuration
+from apify.storage_clients._file_system._key_value_store_client import ApifyFileSystemKeyValueStoreClient
+
+
+async def test_purge_preserves_input_file_and_metadata() -> None:
+    """Test that purge() preserves INPUT.json and metadata files but removes other files."""
+    # Get the global configuration (storage directory is set by test fixtures)
+    config = Configuration.get_global_configuration()
+
+    # Create the key-value store client
+    kvs_client = await ApifyFileSystemKeyValueStoreClient.open(
+        id=None,
+        name='test-kvs',
+        configuration=config,
+    )
+
+    # Create some test files in the KVS directory
+    kvs_path = kvs_client.path_to_kvs
+
+    # Create various files
+    kvs_input_filename = f'{config.input_key}.json'
+    input_file = kvs_path / kvs_input_filename
+    metadata_file = kvs_path / METADATA_FILENAME
+    regular_file1 = kvs_path / 'regular_file1.json'
+    regular_file2 = kvs_path / 'another_file.txt'
+
+    # Write content to files
+    await asyncio.to_thread(input_file.write_text, '{"test": "input"}')
+    await asyncio.to_thread(regular_file1.write_text, '{"test": "data1"}')
+    await asyncio.to_thread(regular_file2.write_text, 'some text content')
+
+    # Verify all files exist before purge
+    assert input_file.exists()
+    assert metadata_file.exists()  # Should exist from client creation
+    assert regular_file1.exists()
+    assert regular_file2.exists()
+
+    # Purge the key-value store
+    await kvs_client.purge()  # Verify INPUT.json and metadata are preserved
+    assert input_file.exists(), f'{kvs_input_filename} should be preserved during purge'
+    assert metadata_file.exists(), f'{METADATA_FILENAME} should be preserved during purge'
+
+    # Verify other files are deleted
+    assert not regular_file1.exists(), 'Regular files should be deleted during purge'
+    assert not regular_file2.exists(), 'Regular files should be deleted during purge'
+
+    # Verify INPUT.json content is unchanged
+    input_content = await asyncio.to_thread(input_file.read_text)
+    assert input_content == '{"test": "input"}'