Caching of Dataset and KVS

Author: vdusek

src/crawlee/storage_clients/_file_system/_dataset_client.py

@@ -6,7 +6,7 @@
 from datetime import datetime, timezone
 from logging import getLogger
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, ClassVar
 
 from pydantic import ValidationError
 from typing_extensions import override
@@ -23,9 +23,6 @@
 
 logger = getLogger(__name__)
 
-_cache_by_name = dict[str, 'FileSystemDatasetClient']()
-"""A dictionary to cache clients by their names."""
-
 
 class FileSystemDatasetClient(DatasetClient):
     """A file system implementation of the dataset client.
@@ -44,6 +41,9 @@ class FileSystemDatasetClient(DatasetClient):
     _LOCAL_ENTRY_NAME_DIGITS = 9
     """Number of digits used for the file names (e.g., 000000019.json)."""
 
+    _cache_by_name: ClassVar[dict[str, FileSystemDatasetClient]] = {}
+    """A dictionary to cache clients by their names."""
+
     def __init__(
         self,
         *,
@@ -131,8 +131,8 @@ async def open(
         name = name or cls._DEFAULT_NAME
 
         # Check if the client is already cached by name.
-        if name in _cache_by_name:
-            client = _cache_by_name[name]
+        if name in cls._cache_by_name:
+            client = cls._cache_by_name[name]
             await client._update_metadata(update_accessed_at=True)  # noqa: SLF001
             return client
 
@@ -182,7 +182,7 @@ async def open(
             await client._update_metadata()
 
         # Cache the client by name.
-        _cache_by_name[name] = client
+        cls._cache_by_name[name] = client
 
         return client
 
@@ -194,8 +194,8 @@ async def drop(self) -> None:
                 await asyncio.to_thread(shutil.rmtree, self.path_to_dataset)
 
         # Remove the client from the cache.
-        if self.name in _cache_by_name:
-            del _cache_by_name[self.name]
+        if self.name in self.__class__._cache_by_name:  # noqa: SLF001
+            del self.__class__._cache_by_name[self.name]  # noqa: SLF001
 
     @override
     async def push_data(self, data: list[Any] | dict[str, Any]) -> None:

src/crawlee/storage_clients/_file_system/_key_value_store_client.py

@@ -6,7 +6,7 @@
 from datetime import datetime, timezone
 from logging import getLogger
 from pathlib import Path
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, ClassVar
 
 from pydantic import ValidationError
 from typing_extensions import override
@@ -24,9 +24,6 @@
 
 logger = getLogger(__name__)
 
-_cache_by_name = dict[str, 'FileSystemKeyValueStoreClient']()
-"""A dictionary to cache clients by their names."""
-
 
 class FileSystemKeyValueStoreClient(KeyValueStoreClient):
     """A file system implementation of the key-value store client.
@@ -42,6 +39,9 @@ class FileSystemKeyValueStoreClient(KeyValueStoreClient):
     _STORAGE_SUBDIR = 'key_value_stores'
     """The name of the subdirectory where key-value stores are stored."""
 
+    _cache_by_name: ClassVar[dict[str, FileSystemKeyValueStoreClient]] = {}
+    """A dictionary to cache clients by their names."""
+
     def __init__(
         self,
         *,
@@ -122,8 +122,8 @@ async def open(
         name = name or cls._DEFAULT_NAME
 
         # Check if the client is already cached by name.
-        if name in _cache_by_name:
-            return _cache_by_name[name]
+        if name in cls._cache_by_name:
+            return cls._cache_by_name[name]
 
         storage_dir = storage_dir or Path.cwd()
         kvs_path = storage_dir / cls._STORAGE_SUBDIR / name
@@ -169,7 +169,7 @@ async def open(
             await client._update_metadata()
 
         # Cache the client by name.
-        _cache_by_name[name] = client
+        cls._cache_by_name[name] = client
 
         return client
 
@@ -181,8 +181,8 @@ async def drop(self) -> None:
                 await asyncio.to_thread(shutil.rmtree, self.path_to_kvs)
 
         # Remove the client from the cache.
-        if self.name in _cache_by_name:
-            del _cache_by_name[self.name]
+        if self.name in self.__class__._cache_by_name:  # noqa: SLF001
+            del self.__class__._cache_by_name[self.name]  # noqa: SLF001
 
     @override
     async def get_value(self, *, key: str) -> KeyValueStoreRecord | None:

src/crawlee/storage_clients/_memory/_dataset_client.py

@@ -2,7 +2,7 @@
 
 from datetime import datetime, timezone
 from logging import getLogger
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, ClassVar
 
 from typing_extensions import override
 
@@ -16,9 +16,6 @@
 
 logger = getLogger(__name__)
 
-_cache_by_name = dict[str, 'MemoryDatasetClient']()
-"""A dictionary to cache clients by their names."""
-
 
 class MemoryDatasetClient(DatasetClient):
     """A memory implementation of the dataset client.
@@ -31,6 +28,9 @@ class MemoryDatasetClient(DatasetClient):
     _DEFAULT_NAME = 'default'
     """The default name for the dataset when no name is provided."""
 
+    _cache_by_name: ClassVar[dict[str, MemoryDatasetClient]] = {}
+    """A dictionary to cache clients by their names."""
+
     def __init__(
         self,
         *,
@@ -102,8 +102,8 @@ async def open(
         name = name or cls._DEFAULT_NAME
 
         # Check if the client is already cached by name.
-        if name in _cache_by_name:
-            client = _cache_by_name[name]
+        if name in cls._cache_by_name:
+            client = cls._cache_by_name[name]
             await client._update_metadata(update_accessed_at=True)  # noqa: SLF001
             return client
 
@@ -120,7 +120,7 @@ async def open(
         )
 
         # Cache the client by name
-        _cache_by_name[name] = client
+        cls._cache_by_name[name] = client
 
         return client
 
@@ -130,8 +130,8 @@ async def drop(self) -> None:
         self._metadata.item_count = 0
 
         # Remove the client from the cache
-        if self.name in _cache_by_name:
-            del _cache_by_name[self.name]
+        if self.name in self.__class__._cache_by_name:
+            del self.__class__._cache_by_name[self.name]
 
     @override
     async def push_data(self, data: list[Any] | dict[str, Any]) -> None:

src/crawlee/storage_clients/_memory/_key_value_store_client.py

@@ -3,7 +3,7 @@
 import sys
 from datetime import datetime, timezone
 from logging import getLogger
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, ClassVar
 
 from typing_extensions import override
 
@@ -18,9 +18,6 @@
 
 logger = getLogger(__name__)
 
-_cache_by_name = dict[str, 'MemoryKeyValueStoreClient']()
-"""A dictionary to cache clients by their names."""
-
 
 class MemoryKeyValueStoreClient(KeyValueStoreClient):
     """A memory implementation of the key-value store client.
@@ -33,6 +30,9 @@ class MemoryKeyValueStoreClient(KeyValueStoreClient):
     _DEFAULT_NAME = 'default'
     """The default name for the key-value store when no name is provided."""
 
+    _cache_by_name: ClassVar[dict[str, MemoryKeyValueStoreClient]] = {}
+    """A dictionary to cache clients by their names."""
+
     def __init__(
         self,
         *,
@@ -97,8 +97,8 @@ async def open(
         name = name or cls._DEFAULT_NAME
 
         # Check if the client is already cached by name
-        if name in _cache_by_name:
-            client = _cache_by_name[name]
+        if name in cls._cache_by_name:
+            client = cls._cache_by_name[name]
             await client._update_metadata(update_accessed_at=True)  # noqa: SLF001
             return client
 
@@ -115,7 +115,7 @@ async def open(
         )
 
         # Cache the client by name
-        _cache_by_name[name] = client
+        cls._cache_by_name[name] = client
 
         return client
 
@@ -125,8 +125,8 @@ async def drop(self) -> None:
         self._store.clear()
 
         # Remove from cache
-        if self.name in _cache_by_name:
-            del _cache_by_name[self.name]
+        if self.name in self.__class__._cache_by_name:
+            del self.__class__._cache_by_name[self.name]
 
     @override
     async def get_value(self, *, key: str) -> KeyValueStoreRecord | None:

src/crawlee/storages/_dataset.py

@@ -3,7 +3,7 @@
 import logging
 from io import StringIO
 from pathlib import Path
-from typing import TYPE_CHECKING, Literal
+from typing import TYPE_CHECKING, ClassVar, Literal
 
 from typing_extensions import override
 
@@ -30,32 +30,6 @@
 
 logger = logging.getLogger(__name__)
 
-# TODO:
-# - caching / memoization of Dataset
-
-# Properties:
-# - id
-# - name
-# - metadata
-
-# Methods:
-# - open
-# - drop
-# - push_data
-# - get_data
-# - iterate_items
-# - export_to
-# - export_to_json
-# - export_to_csv
-
-# 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)?
-# - write_to_json -> export_to_json
-# - write_to_csv -> export_to_csv
-
 
 @docs_group('Classes')
 class Dataset(Storage):
@@ -90,6 +64,12 @@ class Dataset(Storage):
     ```
     """
 
+    _cache_by_id: ClassVar[dict[str, Dataset]] = {}
+    """A dictionary to cache datasets by their IDs."""
+
+    _cache_by_name: ClassVar[dict[str, Dataset]] = {}
+    """A dictionary to cache datasets by their names."""
+
     def __init__(self, client: DatasetClient) -> None:
         """Initialize a new instance.
 
@@ -137,6 +117,12 @@ async def open(
         if id and name:
             raise ValueError('Only one of "id" or "name" can be specified, not both.')
 
+        # Check if dataset is already cached by id or name
+        if id and id in cls._cache_by_id:
+            return cls._cache_by_id[id]
+        if name and name in cls._cache_by_name:
+            return cls._cache_by_name[name]
+
         configuration = service_locator.get_configuration() if configuration is None else configuration
         storage_client = service_locator.get_storage_client() if storage_client is None else storage_client
         purge_on_start = configuration.purge_on_start if purge_on_start is None else purge_on_start
@@ -149,10 +135,24 @@ async def open(
             storage_dir=storage_dir,
         )
 
-        return cls(client)
+        dataset = cls(client)
+
+        # Cache the dataset by id and name if available
+        if dataset.id:
+            cls._cache_by_id[dataset.id] = dataset
+        if dataset.name:
+            cls._cache_by_name[dataset.name] = dataset
+
+        return dataset
 
     @override
     async def drop(self) -> None:
+        # Remove from cache before dropping
+        if self.id in self._cache_by_id:
+            del self._cache_by_id[self.id]
+        if self.name and self.name in self._cache_by_name:
+            del self._cache_by_name[self.name]
+
         await self._client.drop()
 
     async def push_data(self, data: list[Any] | dict[str, Any]) -> None: