scrapely-python 0.1.0__py3-none-any.whl

scrapely/storage/patch_storage.py

@@ -0,0 +1,257 @@
+"""Register Scrapely storage with Crawlee's service locator (TS setApiStorage equivalent)."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, MutableMapping, Optional, TypedDict, cast
+
+from crawlee import service_locator
+from crawlee.storage_clients._base import StorageClient
+
+from ._file_system import ScrapelyFileSystemStorageClient
+from .scrapely_storage_client import ScrapelyStorageClient
+from .smart_scrapely_storage_client import SmartScrapelyStorageClient
+
+logger = logging.getLogger(__name__)
+
+_api_factory: Optional[ScrapelyStorageClient] = None
+_smart_client: Optional[SmartScrapelyStorageClient] = None
+
+# Cache for auto-saved values: store_id -> key -> value
+# This allows get_auto_saved_value to return the same object reference,
+# and persist_state listeners can iterate and save all cached values.
+_auto_save_cache: Dict[str, Dict[str, Any]] = {}
+
+# Track which stores have registered persistState listeners
+_persist_state_registered: set = set()
+
+
+class ApiStorageConfig(TypedDict, total=False):
+    api_url: str
+    token: str
+    run_id: str
+    transport: Any
+
+
+def set_api_storage(
+    api_config: Optional[ApiStorageConfig] = None,
+    *,
+    cloud_storage_client: Optional[ScrapelyStorageClient] = None,
+    local_storage_client: Optional[StorageClient] = None,
+) -> SmartScrapelyStorageClient:
+    """Configure Crawlee to use Scrapely storage (Smart wrapper + ``/v2`` cloud).
+
+    Equivalent to TypeScript ``setApiStorage()`` plus Apify's ``SmartApifyStorageClient``
+    registration on the service locator.
+    """
+    global _api_factory, _smart_client
+
+    if cloud_storage_client is not None:
+        cloud = cloud_storage_client
+    elif api_config is not None:
+        cloud = ScrapelyStorageClient(
+            api_url=api_config.get('api_url'),
+            token=api_config.get('token'),
+            run_id=api_config.get('run_id'),
+            transport=api_config.get('transport'),
+        )
+    else:
+        cloud = ScrapelyStorageClient()
+
+    _api_factory = cloud
+    local = local_storage_client or ScrapelyFileSystemStorageClient()
+    smart = SmartScrapelyStorageClient(
+        cloud_storage_client=cloud,
+        local_storage_client=local,
+    )
+    _smart_client = smart
+    service_locator.set_storage_client(smart)
+    return smart
+
+
+def set_storage_client(client: StorageClient) -> SmartScrapelyStorageClient:
+    """Register a storage client; wrap non-Smart clients like Apify Actor expects."""
+    global _api_factory, _smart_client
+
+    if isinstance(client, SmartScrapelyStorageClient):
+        smart = client
+    elif isinstance(client, ScrapelyStorageClient):
+        _api_factory = client
+        smart = SmartScrapelyStorageClient(
+            cloud_storage_client=client,
+            local_storage_client=ScrapelyFileSystemStorageClient(),
+        )
+    else:
+        smart = SmartScrapelyStorageClient(
+            cloud_storage_client=_api_factory or ScrapelyStorageClient(),
+            local_storage_client=client,
+        )
+
+    _smart_client = smart
+    service_locator.set_storage_client(smart)
+    return smart
+
+
+def get_api_storage_factory() -> Optional[ScrapelyStorageClient]:
+    return _api_factory
+
+
+def get_smart_storage_client() -> Optional[SmartScrapelyStorageClient]:
+    return _smart_client
+
+
+def is_api_storage_configured() -> bool:
+    return _api_factory is not None
+
+
+def disconnect_storage() -> None:
+    global _api_factory, _smart_client, _auto_save_cache, _persist_state_registered
+    _api_factory = None
+    _smart_client = None
+    _auto_save_cache.clear()
+    _persist_state_registered.clear()
+
+
+def _ensure_persist_state_listener(store_id: str, config: Any) -> None:
+    """Register a persistState listener for a store if not already registered.
+    
+    This listener saves all cached values for that store when persistState is emitted.
+    """
+    if store_id in _persist_state_registered:
+        return
+    
+    event_manager = config.get_event_manager() if hasattr(config, 'get_event_manager') else None
+    if event_manager is None:
+        logger.warning('No event manager available for persistState listener')
+        return
+    
+    async def on_persist_state(_: Any) -> None:
+        store_cache = _auto_save_cache.get(store_id)
+        if not store_cache:
+            return
+        
+        factory = _api_factory
+        if not factory:
+            return
+        
+        try:
+            api = factory._ensure_api()
+            for key, value in store_cache.items():
+                try:
+                    await api.kvs_set_record(store_id, key, value)
+                except Exception as e:
+                    logger.warning('Failed to persist state value to %s: %s', key, e)
+        except Exception as e:
+            logger.warning('Failed to persist auto-saved values: %s', e)
+    
+    event_manager.on('persistState', on_persist_state)
+    _persist_state_registered.add(store_id)
+
+
+async def get_auto_saved_value(
+    key: str,
+    default_value: MutableMapping[str, Any],
+    *,
+    store_id: Optional[str] = None,
+) -> MutableMapping[str, Any]:
+    """Get a value from KVS that will be automatically saved on changes.
+    
+    This implementation matches Crawlee/Apify's approach:
+    1. Check cache first (return same object reference for subsequent calls)
+    2. Fetch value from storage if not cached
+    3. Cache the value (so mutations are tracked)
+    4. Register persistState listener to save cached values
+    
+    Args:
+        key: The key of the record to get
+        default_value: The default value to return and store if the key doesn't exist
+        store_id: Optional store ID (uses default KVS if not provided)
+        
+    Returns:
+        The stored value or the default_value
+    """
+    from scrapely.configuration import Configuration as ScrapelyConfiguration
+    
+    config = ScrapelyConfiguration.get_global_config()
+    
+    # Resolve store ID from environment if not provided
+    if store_id is None:
+        store_id = (
+            config.get('default_key_value_store_id')
+            or config.actor_run_id
+        )
+    
+    if store_id is None:
+        raise ValueError('No key-value store ID available for get_auto_saved_value')
+    
+    # Check if we have a cached value for this store/key
+    store_cache = _auto_save_cache.get(store_id)
+    if store_cache is not None and key in store_cache:
+        return cast(MutableMapping[str, Any], store_cache[key])
+    
+    # Fetch value from storage
+    factory = _api_factory
+    if factory is None:
+        raise ValueError('API storage client not initialized. Call set_api_storage() first.')
+    
+    api = factory._ensure_api()
+    
+    # Try to get existing value
+    value, content_type = await api.kvs_get_record(store_id, key)
+    
+    if value is None:
+        # Store the default value and return it
+        await api.kvs_set_record(store_id, key, default_value)
+        value = default_value
+    elif not isinstance(value, MutableMapping):
+        # If the stored value is not a mutable mapping, we can't track mutations
+        # Return it as-is but don't cache for auto-save
+        logger.warning(
+            'get_auto_saved_value: stored value for key "%s" is not a mutable mapping, '
+            'auto-save will not work',
+            key,
+        )
+        return cast(MutableMapping[str, Any], value)
+    
+    # Initialize store cache if needed
+    if store_cache is None:
+        store_cache = {}
+        _auto_save_cache[store_id] = store_cache
+    
+    # Check cache again (race condition: another call might have cached it)
+    if key in store_cache:
+        return cast(MutableMapping[str, Any], store_cache[key])
+    
+    # Cache the value - mutations to this object will be persisted
+    store_cache[key] = value
+    
+    # Register persistState listener for this store
+    _ensure_persist_state_listener(store_id, config)
+    
+    return cast(MutableMapping[str, Any], value)
+
+
+async def persist_autosaved_values(*, store_id: Optional[str] = None) -> None:
+    """Force autosaved values to be saved without waiting for an event.
+    
+    Args:
+        store_id: Optional store ID to persist. If None, persists all stores.
+    """
+    factory = _api_factory
+    if factory is None:
+        return
+    
+    api = factory._ensure_api()
+    
+    stores_to_persist = [store_id] if store_id else list(_auto_save_cache.keys())
+    
+    for sid in stores_to_persist:
+        store_cache = _auto_save_cache.get(sid)
+        if not store_cache:
+            continue
+        
+        for key, value in store_cache.items():
+            try:
+                await api.kvs_set_record(sid, key, value)
+            except Exception as e:
+                logger.warning('Failed to persist state value to %s: %s', key, e)

scrapely/storage/scrapely_dataset_client.py

@@ -0,0 +1,305 @@
+"""Scrapely cloud DatasetClient for Crawlee."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, List, Mapping, Optional, Sequence, Union
+
+from typing_extensions import Self, override
+
+from crawlee.storage_clients._base import DatasetClient
+from crawlee.storage_clients.models import DatasetItemsListPage, DatasetMetadata
+
+from .client import ApiStorageClient
+from ._metadata import resolve_storage_id, to_dataset_metadata
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
+
+    from crawlee._types import JsonSerializable
+    from crawlee.configuration import Configuration
+    from crawlee.storage_clients import StorageClient
+
+
+def _old_dataset_to_crawlee(meta: Any) -> DatasetMetadata:
+    return to_dataset_metadata(
+        {
+            'id': meta.id,
+            'name': meta.name,
+            'createdAt': meta.created_at.isoformat() if getattr(meta, 'created_at', None) else None,
+            'accessedAt': meta.accessed_at.isoformat() if getattr(meta, 'accessed_at', None) else None,
+            'modifiedAt': meta.modified_at.isoformat() if getattr(meta, 'modified_at', None) else None,
+            'itemCount': getattr(meta, 'item_count', 0),
+        },
+        meta.id,
+    )
+
+
+class ScrapelyDatasetClient(DatasetClient):
+    """Dataset client backed by Scrapely `/v2/datasets`."""
+
+    def __init__(self, *, api: ApiStorageClient, metadata: DatasetMetadata) -> None:
+        self._api = api
+        self._metadata = metadata
+
+    @property
+    def id(self) -> str:
+        """Get the storage ID."""
+        return self._metadata.id
+
+    @property
+    def name(self) -> Optional[str]:
+        """Get the storage name."""
+        return self._metadata.name
+
+    @classmethod
+    async def open(
+        cls,
+        *,
+        api: ApiStorageClient,
+        id: Optional[str] = None,
+        name: Optional[str] = None,
+        alias: Optional[str] = None,
+        default_id: Optional[str] = None,
+    ) -> Self:
+        open_id, open_name = resolve_storage_id(
+            id=id, name=name, alias=alias, default_id=default_id,
+        )
+        meta = await api.get_or_create_dataset(id=open_id, name=open_name)
+        return cls(api=api, metadata=_old_dataset_to_crawlee(meta))
+
+    @override
+    async def get_metadata(self) -> DatasetMetadata:
+        meta = await self._api.dataset_get(self._metadata.id)
+        self._metadata = _old_dataset_to_crawlee(meta)
+        return self._metadata
+
+    @override
+    async def drop(self) -> None:
+        await self._api.dataset_delete(self._metadata.id)
+
+    @override
+    async def purge(self) -> None:
+        if not self._metadata.name:
+            raise NotImplementedError(
+                'Dataset.purge() is not supported for unnamed/default datasets on Scrapely '
+                '(no clear endpoint). Use a named dataset or drop() and reopen.'
+            )
+        name = self._metadata.name
+        await self.drop()
+        meta = await self._api.get_or_create_dataset(name=name)
+        self._metadata = _old_dataset_to_crawlee(meta)
+
+    @override
+    async def push_data(
+        self,
+        data: Union[Sequence[Mapping[str, 'JsonSerializable']], Mapping[str, 'JsonSerializable']],
+    ) -> None:
+        items: List[Mapping[str, Any]]
+        if isinstance(data, Mapping):
+            items = [data]
+        else:
+            items = list(data)
+        await self._api.dataset_push_items(self._metadata.id, items)
+
+    @override
+    async def get_data(
+        self,
+        *,
+        offset: int = 0,
+        limit: Optional[int] = 999_999_999_999,
+        clean: bool = False,
+        desc: bool = False,
+        fields: Optional[List[str]] = None,
+        omit: Optional[List[str]] = None,
+        unwind: Optional[List[str]] = None,
+        skip_empty: bool = False,
+        skip_hidden: bool = False,
+        flatten: Optional[List[str]] = None,
+        view: Optional[str] = None,
+    ) -> DatasetItemsListPage:
+        _, page = await self._api.dataset_list_items(
+            self._metadata.id,
+            offset=offset,
+            limit=None if limit == 999_999_999_999 else limit,
+            clean=clean,
+            desc=desc,
+            fields=fields,
+            omit=omit,
+            unwind=unwind,
+            skip_empty=skip_empty,
+            skip_hidden=skip_hidden,
+            flatten=flatten,
+            view=view,
+        )
+        return DatasetItemsListPage(
+            count=page.count,
+            offset=page.offset,
+            limit=page.limit,
+            total=page.total,
+            desc=page.desc,
+            items=list(page.items),
+        )
+
+    @override
+    async def iterate_items(
+        self,
+        *,
+        offset: int = 0,
+        limit: Optional[int] = None,
+        clean: bool = False,
+        desc: bool = False,
+        fields: Optional[List[str]] = None,
+        omit: Optional[List[str]] = None,
+        unwind: Optional[List[str]] = None,
+        skip_empty: bool = False,
+        skip_hidden: bool = False,
+    ) -> AsyncIterator[Mapping[str, 'JsonSerializable']]:
+        current = offset
+        remaining = limit
+        chunk = 1000
+        while True:
+            fetch = chunk if remaining is None else min(chunk, remaining)
+            if fetch <= 0:
+                break
+            page = await self.get_data(
+                offset=current,
+                limit=fetch,
+                clean=clean,
+                desc=desc,
+                fields=fields,
+                omit=omit,
+                unwind=unwind,
+                skip_empty=skip_empty,
+                skip_hidden=skip_hidden,
+            )
+            if not page.items:
+                break
+            for item in page.items:
+                yield item
+                if remaining is not None:
+                    remaining -= 1
+                    if remaining <= 0:
+                        return
+            current += len(page.items)
+            if len(page.items) < fetch:
+                break
+
+    async def list_items(
+        self,
+        *,
+        offset: int = 0,
+        limit: Optional[int] = None,
+        clean: bool = False,
+        desc: bool = False,
+        fields: Optional[List[str]] = None,
+        omit: Optional[List[str]] = None,
+        unwind: Optional[List[str]] = None,
+        skip_empty: bool = False,
+        skip_hidden: bool = False,
+    ) -> list[Mapping[str, 'JsonSerializable']]:
+        """List all items in the dataset.
+
+        Args:
+            offset: Number of items to skip at the start.
+            limit: Maximum number of items to return.
+            clean: Return only non-empty items, exclude hidden fields.
+            desc: Sort in descending order.
+            fields: Fields to include in each item.
+            omit: Fields to exclude from each item.
+            unwind: Unwind items by a specified array field.
+            skip_empty: Exclude empty items.
+            skip_hidden: Exclude fields starting with '#'.
+            flatten: Fields to flatten in returned items.
+
+        Returns:
+            List of items as dictionaries.
+        """
+        result: list[Mapping[str, 'JsonSerializable']] = []
+        async for item in self.iterate_items(
+            offset=offset,
+            limit=limit,
+            clean=clean,
+            desc=desc,
+            fields=fields,
+            omit=omit,
+            unwind=unwind,
+            skip_empty=skip_empty,
+            skip_hidden=skip_hidden,
+        ):
+            result.append(item)
+        return result
+
+    async def export_to(
+        self,
+        key: str,
+        *,
+        content_type: str = 'json',
+        to_kvs_id: Optional[str] = None,
+        to_kvs_name: Optional[str] = None,
+        to_kvs_storage_client: Optional['StorageClient'] = None,
+        to_kvs_configuration: Optional['Configuration'] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Export the entire dataset to a key-value store.
+
+        Args:
+            key: Key under which to save the data in the KVS.
+            content_type: Format for export ('json' or 'csv').
+            to_kvs_id: ID of the target key-value store.
+            to_kvs_name: Name of the target key-value store.
+            to_kvs_storage_client: Storage client for the target KVS.
+            to_kvs_configuration: Configuration for the target KVS.
+            **kwargs: Additional parameters for the export (e.g., JSON/CSV options).
+        """
+        import json
+        import csv
+        import io
+
+        # Get all items from the dataset
+        items = await self.list_items(
+            clean=kwargs.get('clean', False),
+            skip_empty=kwargs.get('skip_empty', False),
+            skip_hidden=kwargs.get('skip_hidden', False),
+        )
+
+        # Determine target KVS - use this API client's KVS by default
+        if to_kvs_id or to_kvs_name:
+            target_kvs_meta = await self._api.get_or_create_key_value_store(
+                id=to_kvs_id,
+                name=to_kvs_name,
+            )
+            target_kvs_id = target_kvs_meta.id
+        else:
+            # Use default KVS
+            target_kvs_meta = await self._api.get_or_create_key_value_store()
+            target_kvs_id = target_kvs_meta.id
+
+        # Format the data
+        if content_type == 'csv':
+            output = io.StringIO()
+            if items:
+                writer = csv.DictWriter(
+                    output,
+                    fieldnames=items[0].keys() if items else [],
+                    **{k: v for k, v in kwargs.items() if k in [
+                        'dialect', 'delimiter', 'doublequote', 'escapechar',
+                        'lineterminator', 'quotechar', 'quoting', 'skipinitialspace', 'strict'
+                    ]},
+                )
+                writer.writeheader()
+                writer.writerows(items)
+            value = output.getvalue()
+            ct = 'text/csv'
+        else:
+            # JSON format
+            value = json.dumps(
+                items,
+                **{k: v for k, v in kwargs.items() if k in [
+                    'skipkeys', 'ensure_ascii', 'check_circular', 'allow_nan',
+                    'cls', 'indent', 'separators', 'default', 'sort_keys'
+                ]},
+            )
+            ct = 'application/json'
+
+        # Store in KVS
+        await self._api.kvs_set_record(target_kvs_id, key, value, ct)