KVS update

Author: vdusek

src/crawlee/storage_clients/_base/_key_value_store_client.py

@@ -6,10 +6,15 @@
 from crawlee._utils.docs import docs_group
 
 if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
     from datetime import datetime
     from pathlib import Path
 
-    from crawlee.storage_clients.models import KeyValueStoreListKeysPage, KeyValueStoreRecord
+    from crawlee.storage_clients.models import (
+        KeyValueStoreListKeysPage,
+        KeyValueStoreRecord,
+        KeyValueStoreRecordMetadata,
+    )
 
 # Properties:
 # - id
@@ -126,6 +131,23 @@ async def iterate_keys(
         *,
         exclusive_start_key: str | None = None,
         limit: int = 1000,
+    ) -> AsyncIterator[KeyValueStoreRecordMetadata]:
+        """Iterate over the existing keys in the key-value store.
+
+        The backend method for the `KeyValueStore.iterate_keys` call.
+        """
+        # This syntax is to make mypy properly work with abstract AsyncIterator.
+        # https://mypy.readthedocs.io/en/stable/more_types.html#asynchronous-iterators
+        raise NotImplementedError
+        if False:  # type: ignore[unreachable]
+            yield 0
+
+    @abstractmethod
+    async def list_keys(
+        self,
+        *,
+        exclusive_start_key: str | None = None,
+        limit: int = 1000,
     ) -> KeyValueStoreListKeysPage:
         """List the keys in the key-value store.
 

src/crawlee/storage_clients/_file_system/_key_value_store_client.py

@@ -14,7 +14,6 @@
 from crawlee._utils.crypto import crypto_random_object_id
 from crawlee.storage_clients._base import KeyValueStoreClient
 from crawlee.storage_clients.models import (
-    KeyValueStoreKeyInfo,
     KeyValueStoreListKeysPage,
     KeyValueStoreMetadata,
     KeyValueStoreRecord,
@@ -245,11 +244,15 @@ async def get_value(self, *, key: str) -> KeyValueStoreRecord | None:
         # Update the metadata to record access
         await self._update_metadata(update_accessed_at=True)
 
+        # Calculate the size of the value in bytes
+        size = len(value_bytes)
+
         return KeyValueStoreRecord(
             key=metadata.key,
             value=value,
             content_type=metadata.content_type,
             filename=filename,
+            size=size,
         )
 
     @override
@@ -271,7 +274,9 @@ async def set_value(self, *, key: str, value: Any, content_type: str | None = No
         record_path = self._path_to_kvs / filename
 
         # Get the metadata.
-        record_metadata = KeyValueStoreRecordMetadata(key=key, content_type=content_type)
+        # Calculate the size of the value in bytes
+        size = len(value_bytes)
+        record_metadata = KeyValueStoreRecordMetadata(key=key, content_type=content_type, size=size)
         record_metadata_filepath = record_path.with_name(f'{record_path.name}.{METADATA_FILENAME}')
         record_metadata_content = await json_dumps(record_metadata.model_dump())
 
@@ -330,61 +335,87 @@ async def delete_value(self, *, key: str) -> None:
             if deleted:
                 await self._update_metadata(update_accessed_at=True, update_modified_at=True)
 
+    @override
     async def iterate_keys(
         self,
         *,
         exclusive_start_key: str | None = None,
         limit: int = 1000,
-    ) -> AsyncIterator[KeyValueStoreKeyInfo]:
-        keys = []
-        has_next = False
-
+    ) -> AsyncIterator[KeyValueStoreRecordMetadata]:
         # Check if the KVS directory exists
         if not self._path_to_kvs.exists():
             return
 
+        count = 0
         async with self._lock:
             # Get all files in the KVS directory
-            files = await asyncio.to_thread(self._path_to_kvs.glob, '*')
+            files = sorted(await asyncio.to_thread(list, self._path_to_kvs.glob('*')))
 
-            # Filter out metadata files and get unique key names
-            key_files = {}
             for file_path in files:
+                # Skip the main metadata file
                 if file_path.name == METADATA_FILENAME:
                     continue
 
-                # Skip metadata files for records
-                if file_path.name.endswith(f'.{METADATA_FILENAME}'):
+                # Only process metadata files for records
+                if not file_path.name.endswith(f'.{METADATA_FILENAME}'):
                     continue
 
-                # Extract the base key name
-                key = file_path.name
-                key_files[key] = file_path
+                # Extract the base key name from the metadata filename
+                key_name = file_path.name[: -len(f'.{METADATA_FILENAME}')]
 
-            # Sort keys for consistent ordering
-            all_keys = sorted(key_files.keys())
+                # Apply exclusive_start_key filter if provided
+                if exclusive_start_key is not None and key_name <= exclusive_start_key:
+                    continue
 
-            # Apply exclusive_start_key if provided
-            if exclusive_start_key is not None:
-                start_idx = 0
-                for idx, key in enumerate(all_keys):
-                    if key > exclusive_start_key:  # exclusive start
-                        start_idx = idx
-                        break
-                all_keys = all_keys[start_idx:]
+                # Try to read and parse the metadata file
+                try:
+                    metadata_content = await asyncio.to_thread(file_path.read_text, encoding='utf-8')
+                    metadata_dict = json.loads(metadata_content)
+                    record_metadata = KeyValueStoreRecordMetadata(**metadata_dict)
 
-            # Apply limit
-            if len(all_keys) > limit:
-                keys = all_keys[:limit]
-                has_next = True
-            else:
-                keys = all_keys
-                has_next = False
+                    yield record_metadata
+
+                    count += 1
+                    if count >= limit:
+                        break
+                except (json.JSONDecodeError, ValidationError) as e:
+                    logger.warning(f'Failed to parse metadata file {file_path}: {e}')
 
         # Update accessed_at timestamp
         await self._update_metadata(update_accessed_at=True)
 
-        return KeyValueStoreListKeysPage(keys=keys, has_next=has_next)
+    @override
+    async def list_keys(
+        self,
+        *,
+        exclusive_start_key: str | None = None,
+        limit: int = 1000,
+    ) -> KeyValueStoreListKeysPage:
+        keys = []
+        had_more = False
+        next_exclusive_start_key = None
+
+        # Use the iterate_keys method to get all keys
+        async for metadata in self.iterate_keys(exclusive_start_key=exclusive_start_key, limit=limit + 1):
+            keys.append(metadata.key)
+            # If we've collected more than the limit, we know there are more keys
+            if len(keys) > limit:
+                had_more = True
+                next_exclusive_start_key = metadata.key
+                keys.pop()  # Remove the extra key
+                break
+
+        # Update the accessed_at timestamp is already handled by iterate_keys
+
+        return KeyValueStoreListKeysPage(
+            count=len(keys),
+            items=keys,
+            had_more=had_more,
+            is_truncated=had_more,
+            limit=limit,
+            exclusive_start_key=exclusive_start_key,
+            next_exclusive_start_key=next_exclusive_start_key,
+        )
 
     @override
     async def get_public_url(self, *, key: str) -> str:

src/crawlee/storage_clients/models.py

@@ -63,35 +63,35 @@ class RequestQueueMetadata(StorageMetadata):
 
 
 @docs_group('Data structures')
-class KeyValueStoreRecord(BaseModel, Generic[KvsValueType]):
-    """Model for a key-value store record."""
+class KeyValueStoreRecordMetadata(BaseModel):
+    """Model for a key-value store record metadata."""
 
     model_config = ConfigDict(populate_by_name=True)
 
     key: Annotated[str, Field(alias='key')]
-    value: Annotated[KvsValueType, Field(alias='value')]
-    content_type: Annotated[str | None, Field(alias='contentType', default=None)]
-    filename: Annotated[str | None, Field(alias='filename', default=None)]
+    """The key of the record.
 
+    A unique identifier for the record in the key-value store.
+    """
 
-@docs_group('Data structures')
-class KeyValueStoreRecordMetadata(BaseModel):
-    """Model for a key-value store record metadata."""
+    content_type: Annotated[str, Field(alias='contentType')]
+    """The MIME type of the record.
 
-    model_config = ConfigDict(populate_by_name=True)
+    Describe the format and type of data stored in the record, following the MIME specification.
+    """
 
-    key: Annotated[str, Field(alias='key')]
-    content_type: Annotated[str, Field(alias='contentType')]
+    size: Annotated[int, Field(alias='size')]
+    """The size of the record in bytes."""
 
 
 @docs_group('Data structures')
-class KeyValueStoreKeyInfo(BaseModel):
-    """Model for a key-value store key info."""
+class KeyValueStoreRecord(KeyValueStoreRecordMetadata, Generic[KvsValueType]):
+    """Model for a key-value store record."""
 
     model_config = ConfigDict(populate_by_name=True)
 
-    key: Annotated[str, Field(alias='key')]
-    size: Annotated[int, Field(alias='size')]
+    value: Annotated[KvsValueType, Field(alias='value')]
+    """The value of the record."""
 
 
 @docs_group('Data structures')
@@ -103,9 +103,9 @@ class KeyValueStoreListKeysPage(BaseModel):
     count: Annotated[int, Field(alias='count')]
     limit: Annotated[int, Field(alias='limit')]
     is_truncated: Annotated[bool, Field(alias='isTruncated')]
-    items: Annotated[list[KeyValueStoreKeyInfo], Field(alias='items', default_factory=list)]
     exclusive_start_key: Annotated[str | None, Field(alias='exclusiveStartKey', default=None)]
     next_exclusive_start_key: Annotated[str | None, Field(alias='nextExclusiveStartKey', default=None)]
+    items: Annotated[list[KeyValueStoreRecordMetadata], Field(alias='items', default_factory=list)]
 
 
 @docs_group('Data structures')

src/crawlee/storages/_key_value_store.py

@@ -1,13 +1,12 @@
 from __future__ import annotations
 
-from collections.abc import AsyncIterator
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar, TypeVar
 
 from crawlee import service_locator
 from crawlee._utils.docs import docs_group
 from crawlee.events._types import Event, EventPersistStateData
-from crawlee.storage_clients.models import KeyValueStoreKeyInfo, KeyValueStoreMetadata
+from crawlee.storage_clients.models import KeyValueStoreMetadata, KeyValueStoreRecordMetadata
 
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator
@@ -21,14 +20,27 @@
 
 # TODO:
 # - inherit from storage class
-# - caching / memoization of both datasets & dataset clients
+# - caching / memoization of both KVS & KVS clients
 
 # Suggested KVS breaking changes:
 # - from_storage_object method has been removed - Use the open method with name and/or id instead.
 # - get_info -> metadata property
 # - storage_object -> metadata property
 # - set_metadata method has been removed - Do we want to support it (e.g. for renaming)?
 
+# Properties:
+# - id
+# - name
+# - metadata
+
+# Methods:
+# - open
+# - drop
+# - get_value
+# - set_value
+# - iterate_keys
+# - list_keys (new method)
+# - get_public_url
 
 @docs_group('Classes')
 class KeyValueStore:
@@ -163,7 +175,7 @@ async def iterate_keys(
         self,
         exclusive_start_key: str | None = None,
         limit: int = 1000,
-    ) -> AsyncIterator[KeyValueStoreKeyInfo]:
+    ) -> AsyncIterator[KeyValueStoreRecordMetadata]:
         """Iterate over the existing keys in the KVS.
 
         Args: