crawlee 1.0.3b6__py3-none-any.whl → 1.0.5b18__py3-none-any.whl

crawlee/statistics/_statistics.py

@@ -1,6 +1,7 @@
 # Inspiration: https://github.com/apify/crawlee/blob/v3.9.2/packages/core/src/crawlers/statistics.ts
 from __future__ import annotations
 
+import asyncio
 import math
 import time
 from datetime import datetime, timedelta, timezone
@@ -17,8 +18,11 @@ from crawlee.statistics import FinalStatistics, StatisticsState
 from crawlee.statistics._error_tracker import ErrorTracker
 
 if TYPE_CHECKING:
+    from collections.abc import Callable, Coroutine
     from types import TracebackType
 
+    from crawlee.storages import KeyValueStore
+
 TStatisticsState = TypeVar('TStatisticsState', bound=StatisticsState, default=StatisticsState)
 TNewStatisticsState = TypeVar('TNewStatisticsState', bound=StatisticsState, default=StatisticsState)
 logger = getLogger(__name__)
@@ -70,6 +74,7 @@ class Statistics(Generic[TStatisticsState]):
         persistence_enabled: bool | Literal['explicit_only'] = 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),
@@ -80,8 +85,6 @@ class Statistics(Generic[TStatisticsState]):
         self._id = Statistics.__next_id
         Statistics.__next_id += 1
 
-        self._instance_start: datetime | None = None
-
         self.error_tracker = ErrorTracker(
             save_error_snapshots=save_error_snapshots,
             snapshot_kvs_name=persist_state_kvs_name,
@@ -92,9 +95,10 @@ class Statistics(Generic[TStatisticsState]):
 
         self._state = RecoverableState(
             default_state=state_model(stats_id=self._id),
-            persist_state_key=persist_state_key or f'SDK_CRAWLER_STATISTICS_{self._id}',
+            persist_state_key=persist_state_key or f'__CRAWLER_STATISTICS_{self._id}',
             persistence_enabled=persistence_enabled,
             persist_state_kvs_name=persist_state_kvs_name,
+            persist_state_kvs_factory=persist_state_kvs_factory,
             logger=logger,
         )
 
@@ -106,12 +110,15 @@ class Statistics(Generic[TStatisticsState]):
         # Flag to indicate the context state.
         self._active = False
 
+        # Pre-existing runtime offset, that can be non-zero when restoring serialized state from KVS.
+        self._runtime_offset = timedelta(seconds=0)
+
     def replace_state_model(self, state_model: type[TNewStatisticsState]) -> Statistics[TNewStatisticsState]:
         """Create near copy of the `Statistics` with replaced `state_model`."""
         new_statistics: Statistics[TNewStatisticsState] = Statistics(
             persistence_enabled=self._state._persistence_enabled,  # noqa: SLF001
-            persist_state_kvs_name=self._state._persist_state_kvs_name,  # noqa: SLF001
             persist_state_key=self._state._persist_state_key,  # noqa: SLF001
+            persist_state_kvs_factory=self._state._persist_state_kvs_factory,  # noqa: SLF001
             log_message=self._log_message,
             periodic_message_logger=self._periodic_message_logger,
             state_model=state_model,
@@ -125,6 +132,7 @@ class Statistics(Generic[TStatisticsState]):
         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),
@@ -136,6 +144,7 @@ class Statistics(Generic[TStatisticsState]):
             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,
@@ -158,14 +167,17 @@ class Statistics(Generic[TStatisticsState]):
         if self._active:
             raise RuntimeError(f'The {self.__class__.__name__} is already active.')
 
-        self._active = True
-        self._instance_start = datetime.now(timezone.utc)
-
         await self._state.initialize()
-        self._after_initialize()
 
+        self._runtime_offset = self.state.crawler_runtime
+
+        # Start periodic logging and let it print initial state before activation.
         self._periodic_logger.start()
+        await asyncio.sleep(0.01)
+        self._active = True
 
+        self.state.crawler_last_started_at = datetime.now(timezone.utc)
+        self.state.crawler_started_at = self.state.crawler_started_at or self.state.crawler_last_started_at
         return self
 
     async def __aexit__(
@@ -182,13 +194,18 @@ class Statistics(Generic[TStatisticsState]):
         if not self._active:
             raise RuntimeError(f'The {self.__class__.__name__} is not active.')
 
-        self._state.current_value.crawler_finished_at = datetime.now(timezone.utc)
-
-        await self._state.teardown()
+        if not self.state.crawler_last_started_at:
+            raise RuntimeError('Statistics.state.crawler_last_started_at not set.')
 
+        # Stop logging and deactivate the statistics to prevent further changes to crawler_runtime
         await self._periodic_logger.stop()
+        self.state.crawler_finished_at = datetime.now(timezone.utc)
+        self.state.crawler_runtime = (
+            self._runtime_offset + self.state.crawler_finished_at - self.state.crawler_last_started_at
+        )
 
         self._active = False
+        await self._state.teardown()
 
     @property
     def state(self) -> TStatisticsState:
@@ -245,13 +262,21 @@ class Statistics(Generic[TStatisticsState]):
 
         del self._requests_in_progress[request_id_or_key]
 
+    def _update_crawler_runtime(self) -> None:
+        current_run_duration = (
+            (datetime.now(timezone.utc) - self.state.crawler_last_started_at)
+            if self.state.crawler_last_started_at
+            else timedelta()
+        )
+        self.state.crawler_runtime = current_run_duration + self._runtime_offset
+
     def calculate(self) -> FinalStatistics:
         """Calculate the current statistics."""
-        if self._instance_start is None:
-            raise RuntimeError('The Statistics object is not initialized')
+        if self._active:
+            # Only update state when active. If not, just report the last known runtime.
+            self._update_crawler_runtime()
 
-        crawler_runtime = datetime.now(timezone.utc) - self._instance_start
-        total_minutes = crawler_runtime.total_seconds() / 60
+        total_minutes = self.state.crawler_runtime.total_seconds() / 60
         state = self._state.current_value
         serialized_state = state.model_dump(by_alias=False)
 
@@ -262,7 +287,7 @@ class Statistics(Generic[TStatisticsState]):
             requests_failed_per_minute=math.floor(state.requests_failed / total_minutes) if total_minutes else 0,
             request_total_duration=state.request_total_finished_duration + state.request_total_failed_duration,
             requests_total=state.requests_failed + state.requests_finished,
-            crawler_runtime=crawler_runtime,
+            crawler_runtime=state.crawler_runtime,
             requests_finished=state.requests_finished,
             requests_failed=state.requests_failed,
             retry_histogram=serialized_state['request_retry_histogram'],
@@ -282,21 +307,6 @@ class Statistics(Generic[TStatisticsState]):
         else:
             self._periodic_message_logger.info(self._log_message, extra=stats.to_dict())
 
-    def _after_initialize(self) -> None:
-        state = self._state.current_value
-
-        if state.crawler_started_at is None:
-            state.crawler_started_at = datetime.now(timezone.utc)
-
-        if state.stats_persisted_at is not None and state.crawler_last_started_at:
-            self._instance_start = datetime.now(timezone.utc) - (
-                state.stats_persisted_at - state.crawler_last_started_at
-            )
-        elif state.crawler_last_started_at:
-            self._instance_start = state.crawler_last_started_at
-
-        state.crawler_last_started_at = self._instance_start
-
     def _save_retry_count_for_request(self, record: RequestProcessingRecord) -> None:
         retry_count = record.retry_count
         state = self._state.current_value

crawlee/storage_clients/__init__.py

@@ -13,9 +13,13 @@ _install_import_hook(__name__)
 with _try_import(__name__, 'SqlStorageClient'):
     from ._sql import SqlStorageClient
 
+with _try_import(__name__, 'RedisStorageClient'):
+    from ._redis import RedisStorageClient
+
 __all__ = [
     'FileSystemStorageClient',
     'MemoryStorageClient',
+    'RedisStorageClient',
     'SqlStorageClient',
     'StorageClient',
 ]

crawlee/storage_clients/_file_system/_request_queue_client.py

@@ -31,6 +31,7 @@ if TYPE_CHECKING:
     from collections.abc import Sequence
 
     from crawlee.configuration import Configuration
+    from crawlee.storages import KeyValueStore
 
 logger = getLogger(__name__)
 
@@ -92,6 +93,7 @@ class FileSystemRequestQueueClient(RequestQueueClient):
         metadata: RequestQueueMetadata,
         path_to_rq: Path,
         lock: asyncio.Lock,
+        recoverable_state: RecoverableState[RequestQueueState],
     ) -> None:
         """Initialize a new instance.
 
@@ -114,12 +116,7 @@ class FileSystemRequestQueueClient(RequestQueueClient):
         self._is_empty_cache: bool | None = None
         """Cache for is_empty result: None means unknown, True/False is cached state."""
 
-        self._state = RecoverableState[RequestQueueState](
-            default_state=RequestQueueState(),
-            persist_state_key=f'__RQ_STATE_{self._metadata.id}',
-            persistence_enabled=True,
-            logger=logger,
-        )
+        self._state = recoverable_state
         """Recoverable state to maintain request ordering, in-progress status, and handled status."""
 
     @override
@@ -136,6 +133,22 @@ class FileSystemRequestQueueClient(RequestQueueClient):
         """The full path to the request queue metadata file."""
         return self.path_to_rq / METADATA_FILENAME
 
+    @classmethod
+    async def _create_recoverable_state(cls, id: str, configuration: Configuration) -> RecoverableState:
+        async def kvs_factory() -> KeyValueStore:
+            from crawlee.storage_clients import FileSystemStorageClient  # noqa: PLC0415 avoid circular import
+            from crawlee.storages import KeyValueStore  # noqa: PLC0415 avoid circular import
+
+            return await KeyValueStore.open(storage_client=FileSystemStorageClient(), configuration=configuration)
+
+        return RecoverableState[RequestQueueState](
+            default_state=RequestQueueState(),
+            persist_state_key=f'__RQ_STATE_{id}',
+            persist_state_kvs_factory=kvs_factory,
+            persistence_enabled=True,
+            logger=logger,
+        )
+
     @classmethod
     async def open(
         cls,
@@ -194,6 +207,9 @@ class FileSystemRequestQueueClient(RequestQueueClient):
                                 metadata=metadata,
                                 path_to_rq=rq_base_path / rq_dir,
                                 lock=asyncio.Lock(),
+                                recoverable_state=await cls._create_recoverable_state(
+                                    id=id, configuration=configuration
+                                ),
                             )
                             await client._state.initialize()
                             await client._discover_existing_requests()
@@ -230,6 +246,7 @@ class FileSystemRequestQueueClient(RequestQueueClient):
                     metadata=metadata,
                     path_to_rq=path_to_rq,
                     lock=asyncio.Lock(),
+                    recoverable_state=await cls._create_recoverable_state(id=metadata.id, configuration=configuration),
                 )
 
                 await client._state.initialize()
@@ -254,6 +271,7 @@ class FileSystemRequestQueueClient(RequestQueueClient):
                     metadata=metadata,
                     path_to_rq=path_to_rq,
                     lock=asyncio.Lock(),
+                    recoverable_state=await cls._create_recoverable_state(id=metadata.id, configuration=configuration),
                 )
                 await client._state.initialize()
                 await client._update_metadata()

crawlee/storage_clients/_redis/__init__.py

@@ -0,0 +1,6 @@
+from ._dataset_client import RedisDatasetClient
+from ._key_value_store_client import RedisKeyValueStoreClient
+from ._request_queue_client import RedisRequestQueueClient
+from ._storage_client import RedisStorageClient
+
+__all__ = ['RedisDatasetClient', 'RedisKeyValueStoreClient', 'RedisRequestQueueClient', 'RedisStorageClient']

crawlee/storage_clients/_redis/_client_mixin.py

@@ -0,0 +1,295 @@
+from __future__ import annotations
+
+import asyncio
+from contextlib import asynccontextmanager
+from datetime import datetime, timezone
+from logging import getLogger
+from typing import TYPE_CHECKING, Any, ClassVar, TypedDict, overload
+
+from crawlee._utils.crypto import crypto_random_object_id
+
+from ._utils import await_redis_response, read_lua_script
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
+
+    from redis.asyncio import Redis
+    from redis.asyncio.client import Pipeline
+    from redis.commands.core import AsyncScript
+    from typing_extensions import NotRequired, Self
+
+    from crawlee.storage_clients.models import DatasetMetadata, KeyValueStoreMetadata, RequestQueueMetadata
+
+
+logger = getLogger(__name__)
+
+
+class MetadataUpdateParams(TypedDict, total=False):
+    """Parameters for updating metadata."""
+
+    update_accessed_at: NotRequired[bool]
+    update_modified_at: NotRequired[bool]
+
+
+class RedisClientMixin:
+    """Mixin class for Redis clients.
+
+    This mixin provides common Redis operations and basic methods for Redis storage clients.
+    """
+
+    _DEFAULT_NAME = 'default'
+    """Default storage name in key prefix when none provided."""
+
+    _MAIN_KEY: ClassVar[str]
+    """Main Redis key prefix for this storage type."""
+
+    _CLIENT_TYPE: ClassVar[str]
+    """Human-readable client type for error messages."""
+
+    def __init__(self, storage_name: str, storage_id: str, redis: Redis) -> None:
+        self._storage_name = storage_name
+        self._storage_id = storage_id
+        self._redis = redis
+
+        self._scripts_loaded = False
+
+    @property
+    def redis(self) -> Redis:
+        """Return the Redis client instance."""
+        return self._redis
+
+    @property
+    def metadata_key(self) -> str:
+        """Return the Redis key for the metadata of this storage."""
+        return f'{self._MAIN_KEY}:{self._storage_name}:metadata'
+
+    @classmethod
+    async def _get_metadata_by_name(cls, name: str, redis: Redis, *, with_wait: bool = False) -> dict | None:
+        """Retrieve metadata by storage name.
+
+        Args:
+            name: The name of the storage.
+            redis: The Redis client instance.
+            with_wait: Whether to wait for the storage to be created if it doesn't exist.
+        """
+        if with_wait:
+            # Wait for the creation signal (max 30 seconds)
+            await await_redis_response(redis.blpop([f'{cls._MAIN_KEY}:{name}:created_signal'], timeout=30))
+            # Signal consumed, push it back for other waiters
+            await await_redis_response(redis.lpush(f'{cls._MAIN_KEY}:{name}:created_signal', 1))
+
+        response = await await_redis_response(redis.json().get(f'{cls._MAIN_KEY}:{name}:metadata'))
+        data = response[0] if response is not None and isinstance(response, list) else response
+        if data is not None and not isinstance(data, dict):
+            raise TypeError('The metadata data was received in an incorrect format.')
+        return data
+
+    @classmethod
+    async def _get_metadata_name_by_id(cls, id: str, redis: Redis) -> str | None:
+        """Retrieve storage name by ID from id_to_name index.
+
+        Args:
+            id: The ID of the storage.
+            redis: The Redis client instance.
+        """
+        name = await await_redis_response(redis.hget(f'{cls._MAIN_KEY}:id_to_name', id))
+        if isinstance(name, str) or name is None:
+            return name
+        if isinstance(name, bytes):
+            return name.decode('utf-8')
+        return None
+
+    @classmethod
+    async def _open(
+        cls,
+        *,
+        id: str | None,
+        name: str | None,
+        alias: str | None,
+        metadata_model: type[DatasetMetadata | KeyValueStoreMetadata | RequestQueueMetadata],
+        redis: Redis,
+        extra_metadata_fields: dict[str, Any],
+        instance_kwargs: dict[str, Any],
+    ) -> Self:
+        """Open or create a new Redis storage client.
+
+        Args:
+            id: The ID of the storage. If not provided, a random ID will be generated.
+            name: The name of the storage for named (global scope) storages.
+            alias: The alias of the storage for unnamed (run scope) storages.
+            redis: Redis client instance.
+            metadata_model: Pydantic model for metadata validation.
+            extra_metadata_fields: Storage-specific metadata fields.
+            instance_kwargs: Additional arguments for the client constructor.
+
+        Returns:
+            An instance for the opened or created storage client.
+        """
+        internal_name = name or alias or cls._DEFAULT_NAME
+        storage_id: str | None = None
+        # Determine if storage exists by ID or name
+        if id:
+            storage_name = await cls._get_metadata_name_by_id(id=id, redis=redis)
+            storage_id = id
+            if storage_name is None:
+                raise ValueError(f'{cls._CLIENT_TYPE} with ID "{id}" does not exist.')
+        else:
+            metadata_data = await cls._get_metadata_by_name(name=internal_name, redis=redis)
+            storage_name = internal_name if metadata_data is not None else None
+            storage_id = metadata_data['id'] if metadata_data is not None else None
+        # If both storage_name and storage_id are found, open existing storage
+        if storage_name and storage_id:
+            client = cls(storage_name=storage_name, storage_id=storage_id, redis=redis, **instance_kwargs)
+            async with client._get_pipeline() as pipe:
+                await client._update_metadata(pipe, update_accessed_at=True)
+        # Otherwise, create a new storage
+        else:
+            now = datetime.now(timezone.utc)
+            metadata = metadata_model(
+                id=crypto_random_object_id(),
+                name=name,
+                created_at=now,
+                accessed_at=now,
+                modified_at=now,
+                **extra_metadata_fields,
+            )
+            client = cls(storage_name=internal_name, storage_id=metadata.id, redis=redis, **instance_kwargs)
+            created = await client._create_metadata_and_storage(internal_name, metadata.model_dump())
+            # The client was probably not created due to a race condition. Let's try to open it using the name.
+            if not created:
+                metadata_data = await cls._get_metadata_by_name(name=internal_name, redis=redis, with_wait=True)
+                client = cls(storage_name=internal_name, storage_id=metadata.id, redis=redis, **instance_kwargs)
+
+        # Ensure Lua scripts are loaded
+        await client._ensure_scripts_loaded()
+        return client
+
+    async def _load_scripts(self) -> None:
+        """Load Lua scripts in Redis."""
+        return
+
+    async def _ensure_scripts_loaded(self) -> None:
+        """Ensure Lua scripts are loaded in Redis."""
+        if not self._scripts_loaded:
+            await self._load_scripts()
+            self._scripts_loaded = True
+
+    @asynccontextmanager
+    async def _get_pipeline(self, *, with_execute: bool = True) -> AsyncIterator[Pipeline]:
+        """Create a new Redis pipeline."""
+        async with self._redis.pipeline() as pipe:
+            try:
+                pipe.multi()  # type: ignore[no-untyped-call]
+                yield pipe
+            finally:
+                if with_execute:
+                    await pipe.execute()
+
+    async def _create_storage(self, pipeline: Pipeline) -> None:
+        """Create the actual storage structure in Redis."""
+        _ = pipeline  # To avoid unused variable mypy error
+
+    async def _create_script(self, script_name: str) -> AsyncScript:
+        """Load a Lua script from a file and return a Script object."""
+        script_content = await asyncio.to_thread(read_lua_script, script_name)
+
+        return self._redis.register_script(script_content)
+
+    async def _create_metadata_and_storage(self, storage_name: str, metadata: dict) -> bool:
+        index_id_to_name = f'{self._MAIN_KEY}:id_to_name'
+        index_name_to_id = f'{self._MAIN_KEY}:name_to_id'
+        metadata['created_at'] = metadata['created_at'].isoformat()
+        metadata['accessed_at'] = metadata['accessed_at'].isoformat()
+        metadata['modified_at'] = metadata['modified_at'].isoformat()
+
+        # Try to create name_to_id index entry, if it already exists, return False.
+        name_to_id = await await_redis_response(self._redis.hsetnx(index_name_to_id, storage_name, metadata['id']))
+        # If name already exists, return False. Probably an attempt at parallel creation.
+        if not name_to_id:
+            return False
+
+        # Create id_to_name index entry, metadata, and storage structure in a transaction.
+        async with self._get_pipeline() as pipe:
+            await await_redis_response(pipe.hsetnx(index_id_to_name, metadata['id'], storage_name))
+            await await_redis_response(pipe.json().set(self.metadata_key, '$', metadata))
+            await await_redis_response(pipe.lpush(f'{self._MAIN_KEY}:{storage_name}:created_signal', 1))
+
+            await self._create_storage(pipe)
+
+        return True
+
+    async def _drop(self, extra_keys: list[str]) -> None:
+        async with self._get_pipeline() as pipe:
+            await pipe.delete(self.metadata_key)
+            await pipe.delete(f'{self._MAIN_KEY}:id_to_name', self._storage_id)
+            await pipe.delete(f'{self._MAIN_KEY}:name_to_id', self._storage_name)
+            await pipe.delete(f'{self._MAIN_KEY}:{self._storage_name}:created_signal')
+            for key in extra_keys:
+                await pipe.delete(key)
+
+    async def _purge(self, extra_keys: list[str], metadata_kwargs: MetadataUpdateParams) -> None:
+        async with self._get_pipeline() as pipe:
+            for key in extra_keys:
+                await pipe.delete(key)
+            await self._update_metadata(pipe, **metadata_kwargs)
+            await self._create_storage(pipe)
+
+    @overload
+    async def _get_metadata(self, metadata_model: type[DatasetMetadata]) -> DatasetMetadata: ...
+    @overload
+    async def _get_metadata(self, metadata_model: type[KeyValueStoreMetadata]) -> KeyValueStoreMetadata: ...
+    @overload
+    async def _get_metadata(self, metadata_model: type[RequestQueueMetadata]) -> RequestQueueMetadata: ...
+
+    async def _get_metadata(
+        self, metadata_model: type[DatasetMetadata | KeyValueStoreMetadata | RequestQueueMetadata]
+    ) -> DatasetMetadata | KeyValueStoreMetadata | RequestQueueMetadata:
+        """Retrieve client metadata."""
+        metadata_dict = await self._get_metadata_by_name(name=self._storage_name, redis=self._redis)
+        if metadata_dict is None:
+            raise ValueError(f'{self._CLIENT_TYPE} with name "{self._storage_name}" does not exist.')
+        async with self._get_pipeline() as pipe:
+            await self._update_metadata(pipe, update_accessed_at=True)
+
+        return metadata_model.model_validate(metadata_dict)
+
+    async def _specific_update_metadata(self, pipeline: Pipeline, **kwargs: Any) -> None:
+        """Pipeline operations storage-specific metadata updates.
+
+        Must be implemented by concrete classes.
+
+        Args:
+            pipeline: The Redis pipeline to use for the update.
+            **kwargs: Storage-specific update parameters.
+        """
+        _ = pipeline  # To avoid unused variable mypy error
+        _ = kwargs
+
+    async def _update_metadata(
+        self,
+        pipeline: Pipeline,
+        *,
+        update_accessed_at: bool = False,
+        update_modified_at: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        """Update storage metadata combining common and specific fields.
+
+        Args:
+            pipeline: The Redis pipeline to use for the update.
+            update_accessed_at: Whether to update accessed_at timestamp.
+            update_modified_at: Whether to update modified_at timestamp.
+            **kwargs: Additional arguments for _specific_update_metadata.
+        """
+        now = datetime.now(timezone.utc)
+
+        if update_accessed_at:
+            await await_redis_response(
+                pipeline.json().set(self.metadata_key, '$.accessed_at', now.isoformat(), nx=False, xx=True)
+            )
+        if update_modified_at:
+            await await_redis_response(
+                pipeline.json().set(self.metadata_key, '$.modified_at', now.isoformat(), nx=False, xx=True)
+            )
+
+        await self._specific_update_metadata(pipeline, **kwargs)