Persist Crawler statistics to Crawler KVS

Author: Pijukatel

src/crawlee/crawlers/_basic/_basic_crawler.py

@@ -437,15 +437,23 @@ def __init__(
         self._statistics_log_format = statistics_log_format
 
         # Statistics
-        self._statistics = statistics or cast(
-            'Statistics[TStatisticsState]',
-            Statistics.with_default_state(
-                persistence_enabled=True,
-                periodic_message_logger=self._logger,
-                statistics_log_format=self._statistics_log_format,
-                log_message='Current request statistics:',
-            ),
-        )
+        if statistics:
+            self._statistics = statistics
+        else:
+
+            async def persist_state_factory() -> KeyValueStore:
+                return await self.get_key_value_store()
+
+            self._statistics = cast(
+                'Statistics[TStatisticsState]',
+                Statistics.with_default_state(
+                    persistence_enabled=True,
+                    periodic_message_logger=self._logger,
+                    statistics_log_format=self._statistics_log_format,
+                    log_message='Current request statistics:',
+                    persist_state_kvs_factory=persist_state_factory,
+                ),
+            )
 
         # Additional context managers to enter and exit
         self._additional_context_managers = _additional_context_managers or []

src/crawlee/statistics/_statistics.py

@@ -130,6 +130,7 @@ def with_default_state(
         persistence_enabled: bool = False,
         persist_state_kvs_name: str | None = None,
         persist_state_key: str | None = None,
+        persist_state_kvs_factory: Callable[[], Coroutine[None, None, KeyValueStore]] | None = None,
         log_message: str = 'Statistics',
         periodic_message_logger: Logger | None = None,
         log_interval: timedelta = timedelta(minutes=1),
@@ -141,6 +142,7 @@ def with_default_state(
             persistence_enabled=persistence_enabled,
             persist_state_kvs_name=persist_state_kvs_name,
             persist_state_key=persist_state_key,
+            persist_state_kvs_factory=persist_state_kvs_factory,
             log_message=log_message,
             periodic_message_logger=periodic_message_logger,
             log_interval=log_interval,

tests/unit/test_configuration.py

@@ -8,6 +8,7 @@
 from crawlee import service_locator
 from crawlee.configuration import Configuration
 from crawlee.crawlers import HttpCrawler, HttpCrawlingContext
+from crawlee.statistics import Statistics
 from crawlee.storage_clients import MemoryStorageClient
 from crawlee.storage_clients._file_system._storage_client import FileSystemStorageClient
 
@@ -35,26 +36,51 @@ def test_global_configuration_works_reversed() -> None:
     )
 
 
-async def test_storage_not_persisted_when_disabled(tmp_path: Path, server_url: URL) -> None:
+async def test_storage_not_persisted_when_non_persistable_storage_used(tmp_path: Path, server_url: URL) -> None:
+    """Make the Crawler use MemoryStorageClient which can't persist state."""
+    service_locator.set_configuration(
+        Configuration(
+            crawlee_storage_dir=str(tmp_path),  # type: ignore[call-arg]
+        )
+    )
+    crawler = HttpCrawler(storage_client=MemoryStorageClient())
+
+    @crawler.router.default_handler
+    async def default_handler(context: HttpCrawlingContext) -> None:
+        await context.push_data({'url': context.request.url})
+
+    await crawler.run([str(server_url)])
+
+    # Verify that no files were created in the storage directory.
+    content = list(tmp_path.iterdir())
+    assert content == [], 'Expected the storage directory to be empty, but it is not.'
+
+
+async def test_storage_persisted_with_explicit_statistics_with_persistable_storage(
+    tmp_path: Path, server_url: URL
+) -> None:
+    """Make the Crawler use MemoryStorageClient which can't persist state,
+    but pass explicit statistics to it which will use global FileSystemStorageClient() that can persist state."""
+
     configuration = Configuration(
         crawlee_storage_dir=str(tmp_path),  # type: ignore[call-arg]
     )
-    storage_client = MemoryStorageClient()
-
     service_locator.set_configuration(configuration)
-    service_locator.set_storage_client(storage_client)
+    service_locator.set_storage_client(FileSystemStorageClient())
 
-    crawler = HttpCrawler()
+    crawler = HttpCrawler(
+        storage_client=MemoryStorageClient(), statistics=Statistics.with_default_state(persistence_enabled=True)
+    )
 
     @crawler.router.default_handler
     async def default_handler(context: HttpCrawlingContext) -> None:
         await context.push_data({'url': context.request.url})
 
     await crawler.run([str(server_url)])
 
-    # Verify that no files were created in the storage directory.
+    # Verify that files were created in the storage directory.
     content = list(tmp_path.iterdir())
-    assert content == [], 'Expected the storage directory to be empty, but it is not.'
+    assert content != [], 'Expected the storage directory to contain files, but it does not.'
 
 
 async def test_storage_persisted_when_enabled(tmp_path: Path, server_url: URL) -> None: