crawlee 1.0.0rc1__py3-none-any.whl → 1.0.1__py3-none-any.whl

crawlee/request_loaders/_sitemap_request_loader.py

@@ -1,19 +1,25 @@
 from __future__ import annotations
 
 import asyncio
+from collections import deque
 from contextlib import suppress
 from logging import getLogger
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Annotated, Any
+
+from pydantic import BaseModel, ConfigDict, Field
+from typing_extensions import override
 
 from crawlee import Request
 from crawlee._utils.docs import docs_group
 from crawlee._utils.globs import Glob
-from crawlee._utils.sitemap import ParseSitemapOptions, SitemapSource, SitemapUrl, parse_sitemap
+from crawlee._utils.recoverable_state import RecoverableState
+from crawlee._utils.sitemap import NestedSitemap, ParseSitemapOptions, SitemapSource, SitemapUrl, parse_sitemap
 from crawlee.request_loaders._request_loader import RequestLoader
 
 if TYPE_CHECKING:
     import re
     from collections.abc import Sequence
+    from types import TracebackType
 
     from crawlee.http_clients import HttpClient
     from crawlee.proxy_configuration import ProxyInfo
@@ -23,12 +29,72 @@ if TYPE_CHECKING:
 logger = getLogger(__name__)
 
 
+class SitemapRequestLoaderState(BaseModel):
+    """State model for persisting sitemap request loader data.
+
+    The crawler processes one sitemap at a time. The current sitemap is stored in `in_progress_sitemap_url`.
+    The `parse_sitemap` function parses the sitemap and returns elements as an async iterator. Each element retrieved
+    from the iterator is processed based on its type. If the element is a `NestedSitemap`, its URL is added to
+    `pending_sitemap_urls` if it hasn't been processed yet (not in `processed_sitemap_urls`). If the element is a
+    `SitemapUrl`, the system checks whether it already exists in `current_sitemap_processed_urls`. If it exists,
+    the loader was restarted from a saved state and the URL is skipped.
+
+    If the URL is new, it is first added to `url_queue`, then to `current_sitemap_processed_urls`, and `total_count` is
+    incremented by 1. When all elements from the current sitemap iterator have been processed, `in_progress_sitemap_url`
+    is set to `None`, the sitemap URL is added to `processed_sitemap_urls`, and `current_sitemap_processed_urls` is
+    cleared. The next sitemap is retrieved from `pending_sitemap_urls`, skipping any URLs that already exist in
+    `processed_sitemap_urls`. If `pending_sitemap_urls` is empty, `completed` is set to `True`.
+
+    When `fetch_next_request` is called, a URL is extracted from `url_queue` and placed in `in_progress`.
+    When `mark_request_as_handled` is called for the extracted URL, it is removed from `in_progress` and
+    `handled_count` is incremented by 1.
+
+    During initial startup or restart after persistence, state validation occurs in `_get_state`. If both
+    `pending_sitemap_urls` and `in_progress_sitemap_url` are empty and `completed` is False, this indicates a
+    fresh start. In this case, `self._sitemap_urls` are moved to `pending_sitemap_urls`. Otherwise, the system is
+    restarting from a persisted state. If `in_progress` contains any URLs, they are moved back to `url_queue` and
+    `in_progress` is cleared.
+    """
+
+    model_config = ConfigDict(validate_by_name=True, validate_by_alias=True)
+
+    url_queue: Annotated[deque[str], Field(alias='urlQueue')]
+    """Queue of URLs extracted from sitemaps and ready for processing."""
+
+    in_progress: Annotated[set[str], Field(alias='inProgress')] = set()
+    """Set of request URLs currently being processed."""
+
+    pending_sitemap_urls: Annotated[deque[str], Field(alias='pendingSitemapUrls')]
+    """Queue of sitemap URLs that need to be fetched and processed."""
+
+    in_progress_sitemap_url: Annotated[str | None, Field(alias='inProgressSitemapUrl')] = None
+    """The sitemap URL currently being processed."""
+
+    current_sitemap_processed_urls: Annotated[set[str], Field(alias='currentSitemapProcessedUrls')] = set()
+    """URLs from the current sitemap that have been added to the queue."""
+
+    processed_sitemap_urls: Annotated[set[str], Field(alias='processedSitemapUrls')] = set()
+    """Set of processed sitemap URLs."""
+
+    completed: Annotated[bool, Field(alias='sitemapCompleted')] = False
+    """Whether all sitemaps have been fully processed."""
+
+    total_count: Annotated[int, Field(alias='totalCount')] = 0
+    """Total number of URLs found and added to the queue from all processed sitemaps."""
+
+    handled_count: Annotated[int, Field(alias='handledCount')] = 0
+    """Number of URLs that have been successfully handled."""
+
+
 @docs_group('Request loaders')
 class SitemapRequestLoader(RequestLoader):
     """A request loader that reads URLs from sitemap(s).
 
     The loader fetches and parses sitemaps in the background, allowing crawling to start
     before all URLs are loaded. It supports filtering URLs using glob and regex patterns.
+
+    The loader supports state persistence, allowing it to resume from where it left off
+    after interruption when a `persist_state_key` is provided during initialization.
     """
 
     def __init__(
@@ -40,7 +106,7 @@ class SitemapRequestLoader(RequestLoader):
         include: list[re.Pattern[Any] | Glob] | None = None,
         exclude: list[re.Pattern[Any] | Glob] | None = None,
         max_buffer_size: int = 200,
-        parse_sitemap_options: ParseSitemapOptions | None = None,
+        persist_state_key: str | None = None,
     ) -> None:
         """Initialize the sitemap request loader.
 
@@ -50,27 +116,64 @@ class SitemapRequestLoader(RequestLoader):
             include: List of glob or regex patterns to include URLs.
             exclude: List of glob or regex patterns to exclude URLs.
             max_buffer_size: Maximum number of URLs to buffer in memory.
-            parse_sitemap_options: Options for parsing sitemaps, such as `SitemapSource` and `max_urls`.
             http_client: the instance of `HttpClient` to use for fetching sitemaps.
+            persist_state_key: A key for persisting the loader's state in the KeyValueStore.
+                When provided, allows resuming from where it left off after interruption.
+                If None, no state persistence occurs.
         """
         self._http_client = http_client
-
         self._sitemap_urls = sitemap_urls
         self._include = include
         self._exclude = exclude
         self._proxy_info = proxy_info
-        self._parse_sitemap_options = parse_sitemap_options or ParseSitemapOptions()
+        self._max_buffer_size = max_buffer_size
+
+        # Synchronization for queue operations
+        self._queue_has_capacity = asyncio.Event()
+        self._queue_has_capacity.set()
+        self._queue_lock = asyncio.Lock()
+
+        # Initialize recoverable state
+        self._state = RecoverableState(
+            default_state=SitemapRequestLoaderState(
+                url_queue=deque(),
+                pending_sitemap_urls=deque(),
+            ),
+            persistence_enabled=bool(persist_state_key),
+            persist_state_key=persist_state_key or '',
+            logger=logger,
+        )
+
+        # Start background loading
+        self._loading_task = asyncio.create_task(self._load_sitemaps())
 
-        self._handled_count = 0
-        self._total_count = 0
+    async def _get_state(self) -> SitemapRequestLoaderState:
+        """Initialize and return the current state."""
+        async with self._queue_lock:
+            if self._state.is_initialized:
+                return self._state.current_value
 
-        # URL queue and tracking
-        self._url_queue: asyncio.Queue[str] = asyncio.Queue(maxsize=max_buffer_size)
-        self._in_progress: set[str] = set()
-        self._processed_urls: set[str] = set()
+            await self._state.initialize()
 
-        # Loading state
-        self._loading_task = asyncio.create_task(self._load_sitemaps())
+            # Initialize pending sitemaps on first run
+            has_sitemap_for_processing = (
+                self._state.current_value.pending_sitemap_urls or self._state.current_value.in_progress_sitemap_url
+            )
+            if not has_sitemap_for_processing and not self._state.current_value.completed:
+                self._state.current_value.pending_sitemap_urls.extend(self._sitemap_urls)
+
+            if self._state.current_value.in_progress:
+                self._state.current_value.url_queue.extendleft(self._state.current_value.in_progress)
+                self._state.current_value.in_progress.clear()
+
+            if (
+                self._state.current_value.url_queue
+                and len(self._state.current_value.url_queue) >= self._max_buffer_size
+            ):
+                # Notify that the queue is full
+                self._queue_has_capacity.clear()
+
+            return self._state.current_value
 
     def _check_url_patterns(
         self,
@@ -105,73 +208,150 @@ class SitemapRequestLoader(RequestLoader):
     async def _load_sitemaps(self) -> None:
         """Load URLs from sitemaps in the background."""
         try:
-            async for item in parse_sitemap(
-                [SitemapSource(type='url', url=url) for url in self._sitemap_urls],
-                self._http_client,
-                proxy_info=self._proxy_info,
-                options=self._parse_sitemap_options,
-            ):
-                # Only process URL items (not nested sitemaps)
-                if isinstance(item, SitemapUrl):
-                    url = item.loc
-
-                    # Skip if already processed
-                    if url in self._processed_urls:
+            # Get actual state
+            while (state := await self._get_state()) and (state.pending_sitemap_urls or state.in_progress_sitemap_url):
+                # Get sitemap URL for parsing
+                sitemap_url = state.in_progress_sitemap_url
+                if not sitemap_url:
+                    sitemap_url = state.pending_sitemap_urls.popleft()
+                    # Skip processed urls
+                    if sitemap_url in state.processed_sitemap_urls:
                         continue
-
-                    # Check if URL should be included
-                    if not self._check_url_patterns(url, self._include, self._exclude):
+                    state.in_progress_sitemap_url = sitemap_url
+
+                parse_options = ParseSitemapOptions(max_depth=0, emit_nested_sitemaps=True)
+
+                async for item in parse_sitemap(
+                    [SitemapSource(type='url', url=sitemap_url)],
+                    self._http_client,
+                    proxy_info=self._proxy_info,
+                    options=parse_options,
+                ):
+                    if isinstance(item, NestedSitemap):
+                        # Add nested sitemap to queue
+                        if item.loc not in state.pending_sitemap_urls and item.loc not in state.processed_sitemap_urls:
+                            state.pending_sitemap_urls.append(item.loc)
                         continue
 
-                    await self._url_queue.put(url)
-                    self._processed_urls.add(url)
-                    self._total_count += 1
+                    if isinstance(item, SitemapUrl):
+                        url = item.loc
+
+                        state = await self._get_state()
+
+                        # Skip if already processed
+                        if url in state.current_sitemap_processed_urls:
+                            continue
+
+                        # Check if URL should be included
+                        if not self._check_url_patterns(url, self._include, self._exclude):
+                            continue
+
+                        # Check if we have capacity in the queue
+                        await self._queue_has_capacity.wait()
+
+                        state = await self._get_state()
+                        async with self._queue_lock:
+                            state.url_queue.append(url)
+                            state.current_sitemap_processed_urls.add(url)
+                            state.total_count += 1
+                            if len(state.url_queue) >= self._max_buffer_size:
+                                # Notify that the queue is full
+                                self._queue_has_capacity.clear()
+
+                # Clear current sitemap after processing
+                state = await self._get_state()
+                current_sitemap_url = state.in_progress_sitemap_url
+                state.in_progress_sitemap_url = None
+                if current_sitemap_url:
+                    state.processed_sitemap_urls.add(current_sitemap_url)
+                state.current_sitemap_processed_urls.clear()
+
+            # Mark as completed after processing all sitemap urls
+            state.completed = True
 
         except Exception:
             logger.exception('Error loading sitemaps')
             raise
 
+    @override
     async def get_total_count(self) -> int:
         """Return the total number of URLs found so far."""
-        return self._total_count
+        state = await self._get_state()
+        return state.total_count
 
+    @override
+    async def get_handled_count(self) -> int:
+        """Return the number of URLs that have been handled."""
+        state = await self._get_state()
+        return state.handled_count
+
+    @override
     async def is_empty(self) -> bool:
         """Check if there are no more URLs to process."""
-        return self._url_queue.empty() and self._loading_task.done()
+        state = await self._get_state()
+        return not state.url_queue
 
+    @override
     async def is_finished(self) -> bool:
         """Check if all URLs have been processed."""
-        return self._url_queue.empty() and len(self._in_progress) == 0 and self._loading_task.done()
+        state = await self._get_state()
+        return not state.url_queue and len(state.in_progress) == 0 and self._loading_task.done()
 
+    @override
     async def fetch_next_request(self) -> Request | None:
         """Fetch the next request to process."""
-        while not (self._loading_task.done() and self._url_queue.empty()):
-            if self._url_queue.empty():
-                await asyncio.sleep(0.5)
+        while not (await self.is_finished()):
+            state = await self._get_state()
+            if not state.url_queue:
+                await asyncio.sleep(0.1)
                 continue
 
-            url = await self._url_queue.get()
+            async with self._queue_lock:
+                url = state.url_queue.popleft()
+
+                request = Request.from_url(url)
+                state.in_progress.add(request.url)
+                if len(state.url_queue) < self._max_buffer_size:
+                    self._queue_has_capacity.set()
 
-            request = Request.from_url(url)
-            self._in_progress.add(request.unique_key)
             return request
 
         return None
 
+    @override
     async def mark_request_as_handled(self, request: Request) -> ProcessedRequest | None:
         """Mark a request as successfully handled."""
-        if request.unique_key in self._in_progress:
-            self._in_progress.remove(request.unique_key)
-            self._handled_count += 1
+        state = await self._get_state()
+        if request.url in state.in_progress:
+            state.in_progress.remove(request.url)
+            state.handled_count += 1
         return None
 
-    async def get_handled_count(self) -> int:
-        """Return the number of handled requests."""
-        return self._handled_count
-
     async def abort_loading(self) -> None:
         """Abort the sitemap loading process."""
         if self._loading_task and not self._loading_task.done():
             self._loading_task.cancel()
             with suppress(asyncio.CancelledError):
                 await self._loading_task
+
+    async def start(self) -> None:
+        """Start the sitemap loading process."""
+        if self._loading_task and not self._loading_task.done():
+            return
+        self._loading_task = asyncio.create_task(self._load_sitemaps())
+
+    async def close(self) -> None:
+        """Close the request loader."""
+        await self.abort_loading()
+        await self._state.teardown()
+
+    async def __aenter__(self) -> SitemapRequestLoader:
+        """Enter the context manager."""
+        await self.start()
+        return self
+
+    async def __aexit__(
+        self, exc_type: type[BaseException] | None, exc_value: BaseException | None, exc_traceback: TracebackType | None
+    ) -> None:
+        """Exit the context manager."""
+        await self.close()

crawlee/sessions/_models.py

@@ -20,7 +20,7 @@ from ._session import Session
 class SessionModel(BaseModel):
     """Model for a Session object."""
 
-    model_config = ConfigDict(populate_by_name=True)
+    model_config = ConfigDict(validate_by_name=True, validate_by_alias=True)
 
     id: Annotated[str, Field(alias='id')]
     max_age: Annotated[timedelta, Field(alias='maxAge')]
@@ -38,7 +38,7 @@ class SessionModel(BaseModel):
 class SessionPoolModel(BaseModel):
     """Model for a SessionPool object."""
 
-    model_config = ConfigDict(populate_by_name=True)
+    model_config = ConfigDict(validate_by_name=True, validate_by_alias=True)
 
     max_pool_size: Annotated[int, Field(alias='maxPoolSize')]
 

crawlee/statistics/_models.py

@@ -57,7 +57,7 @@ class FinalStatistics:
 class StatisticsState(BaseModel):
     """Statistic data about a crawler run."""
 
-    model_config = ConfigDict(populate_by_name=True, ser_json_inf_nan='constants')
+    model_config = ConfigDict(validate_by_name=True, validate_by_alias=True, ser_json_inf_nan='constants')
     stats_id: Annotated[int | None, Field(alias='statsId')] = None
 
     requests_finished: Annotated[int, Field(alias='requestsFinished')] = 0

crawlee/storage_clients/__init__.py

@@ -1,9 +1,21 @@
+from crawlee._utils.try_import import install_import_hook as _install_import_hook
+from crawlee._utils.try_import import try_import as _try_import
+
+# These imports have only mandatory dependencies, so they are imported directly.
 from ._base import StorageClient
 from ._file_system import FileSystemStorageClient
 from ._memory import MemoryStorageClient
 
+_install_import_hook(__name__)
+
+# The following imports are wrapped in try_import to handle optional dependencies,
+# ensuring the module can still function even if these dependencies are missing.
+with _try_import(__name__, 'SqlStorageClient'):
+    from ._sql import SqlStorageClient
+
 __all__ = [
     'FileSystemStorageClient',
     'MemoryStorageClient',
+    'SqlStorageClient',
     'StorageClient',
 ]

crawlee/storage_clients/_base/_storage_client.py

@@ -6,6 +6,8 @@ from typing import TYPE_CHECKING
 from crawlee._utils.docs import docs_group
 
 if TYPE_CHECKING:
+    from collections.abc import Hashable
+
     from crawlee.configuration import Configuration
 
     from ._dataset_client import DatasetClient
@@ -28,12 +30,21 @@ class StorageClient(ABC):
     (where applicable), and consistent access patterns across all storage types it supports.
     """
 
+    def get_storage_client_cache_key(self, configuration: Configuration) -> Hashable:  # noqa: ARG002
+        """Return a cache key that can differentiate between different storages of this and other clients.
+
+        Can be based on configuration or on the client itself. By default, returns a module and name of the client
+        class.
+        """
+        return f'{self.__class__.__module__}.{self.__class__.__name__}'
+
     @abstractmethod
     async def create_dataset_client(
         self,
         *,
         id: str | None = None,
         name: str | None = None,
+        alias: str | None = None,
         configuration: Configuration | None = None,
     ) -> DatasetClient:
         """Create a dataset client."""
@@ -44,6 +55,7 @@ class StorageClient(ABC):
         *,
         id: str | None = None,
         name: str | None = None,
+        alias: str | None = None,
         configuration: Configuration | None = None,
     ) -> KeyValueStoreClient:
         """Create a key-value store client."""
@@ -54,6 +66,7 @@ class StorageClient(ABC):
         *,
         id: str | None = None,
         name: str | None = None,
+        alias: str | None = None,
         configuration: Configuration | None = None,
     ) -> RequestQueueClient:
         """Create a request queue client."""

crawlee/storage_clients/_file_system/_dataset_client.py

@@ -14,6 +14,7 @@ from typing_extensions import override
 from crawlee._consts import METADATA_FILENAME
 from crawlee._utils.crypto import crypto_random_object_id
 from crawlee._utils.file import atomic_write, json_dumps
+from crawlee._utils.raise_if_too_many_kwargs import raise_if_too_many_kwargs
 from crawlee.storage_clients._base import DatasetClient
 from crawlee.storage_clients.models import DatasetItemsListPage, DatasetMetadata
 
@@ -56,7 +57,7 @@ class FileSystemDatasetClient(DatasetClient):
         self,
         *,
         metadata: DatasetMetadata,
-        storage_dir: Path,
+        path_to_dataset: Path,
         lock: asyncio.Lock,
     ) -> None:
         """Initialize a new instance.
@@ -65,8 +66,8 @@ class FileSystemDatasetClient(DatasetClient):
         """
         self._metadata = metadata
 
-        self._storage_dir = storage_dir
-        """The base directory where the storage data are being persisted."""
+        self._path_to_dataset = path_to_dataset
+        """The full path to the dataset directory."""
 
         self._lock = lock
         """A lock to ensure that only one operation is performed at a time."""
@@ -78,10 +79,7 @@ class FileSystemDatasetClient(DatasetClient):
     @property
     def path_to_dataset(self) -> Path:
         """The full path to the dataset directory."""
-        if self._metadata.name is None:
-            return self._storage_dir / self._STORAGE_SUBDIR / self._STORAGE_SUBSUBDIR_DEFAULT
-
-        return self._storage_dir / self._STORAGE_SUBDIR / self._metadata.name
+        return self._path_to_dataset
 
     @property
     def path_to_metadata(self) -> Path:
@@ -94,6 +92,7 @@ class FileSystemDatasetClient(DatasetClient):
         *,
         id: str | None,
         name: str | None,
+        alias: str | None,
         configuration: Configuration,
     ) -> FileSystemDatasetClient:
         """Open or create a file system dataset client.
@@ -104,17 +103,21 @@ class FileSystemDatasetClient(DatasetClient):
 
         Args:
             id: The ID of the dataset to open. If provided, searches for existing dataset by ID.
-            name: The name of the dataset to open. If not provided, uses the default dataset.
+            name: The name of the dataset for named (global scope) storages.
+            alias: The alias of the dataset for unnamed (run scope) storages.
             configuration: The configuration object containing storage directory settings.
 
         Returns:
             An instance for the opened or created storage client.
 
         Raises:
-            ValueError: If a dataset with the specified ID is not found, or if metadata is invalid.
+            ValueError: If a dataset with the specified ID is not found, if metadata is invalid,
+                or if both name and alias are provided.
         """
-        storage_dir = Path(configuration.storage_dir)
-        dataset_base_path = storage_dir / cls._STORAGE_SUBDIR
+        # Validate input parameters.
+        raise_if_too_many_kwargs(id=id, name=name, alias=alias)
+
+        dataset_base_path = Path(configuration.storage_dir) / cls._STORAGE_SUBDIR
 
         if not dataset_base_path.exists():
             await asyncio.to_thread(dataset_base_path.mkdir, parents=True, exist_ok=True)
@@ -126,19 +129,19 @@ class FileSystemDatasetClient(DatasetClient):
                 if not dataset_dir.is_dir():
                     continue
 
-                metadata_path = dataset_dir / METADATA_FILENAME
-                if not metadata_path.exists():
+                path_to_metadata = dataset_dir / METADATA_FILENAME
+                if not path_to_metadata.exists():
                     continue
 
                 try:
-                    file = await asyncio.to_thread(metadata_path.open)
+                    file = await asyncio.to_thread(path_to_metadata.open)
                     try:
                         file_content = json.load(file)
                         metadata = DatasetMetadata(**file_content)
                         if metadata.id == id:
                             client = cls(
                                 metadata=metadata,
-                                storage_dir=storage_dir,
+                                path_to_dataset=dataset_base_path / dataset_dir,
                                 lock=asyncio.Lock(),
                             )
                             await client._update_metadata(update_accessed_at=True)
@@ -152,16 +155,15 @@ class FileSystemDatasetClient(DatasetClient):
             if not found:
                 raise ValueError(f'Dataset with ID "{id}" not found')
 
-        # Get a new instance by name.
+        # Get a new instance by name or alias.
         else:
-            dataset_path = (
-                dataset_base_path / cls._STORAGE_SUBSUBDIR_DEFAULT if name is None else dataset_base_path / name
-            )
-            metadata_path = dataset_path / METADATA_FILENAME
+            dataset_dir = Path(name) if name else Path(alias) if alias else Path('default')
+            path_to_dataset = dataset_base_path / dataset_dir
+            path_to_metadata = path_to_dataset / METADATA_FILENAME
 
             # If the dataset directory exists, reconstruct the client from the metadata file.
-            if dataset_path.exists() and metadata_path.exists():
-                file = await asyncio.to_thread(open, metadata_path)
+            if path_to_dataset.exists() and path_to_metadata.exists():
+                file = await asyncio.to_thread(open, path_to_metadata)
                 try:
                     file_content = json.load(file)
                 finally:
@@ -169,11 +171,11 @@ class FileSystemDatasetClient(DatasetClient):
                 try:
                     metadata = DatasetMetadata(**file_content)
                 except ValidationError as exc:
-                    raise ValueError(f'Invalid metadata file for dataset "{name}"') from exc
+                    raise ValueError(f'Invalid metadata file for dataset "{name or alias}"') from exc
 
                 client = cls(
                     metadata=metadata,
-                    storage_dir=storage_dir,
+                    path_to_dataset=path_to_dataset,
                     lock=asyncio.Lock(),
                 )
 
@@ -192,7 +194,7 @@ class FileSystemDatasetClient(DatasetClient):
                 )
                 client = cls(
                     metadata=metadata,
-                    storage_dir=storage_dir,
+                    path_to_dataset=path_to_dataset,
                     lock=asyncio.Lock(),
                 )
                 await client._update_metadata()