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

crawlee/storage_clients/_sql/_dataset_client.py

@@ -0,0 +1,310 @@
+from __future__ import annotations
+
+from logging import getLogger
+from typing import TYPE_CHECKING, Any
+
+from sqlalchemy import Select, insert, select
+from typing_extensions import override
+
+from crawlee.storage_clients._base import DatasetClient
+from crawlee.storage_clients.models import DatasetItemsListPage, DatasetMetadata
+
+from ._client_mixin import MetadataUpdateParams, SqlClientMixin
+from ._db_models import DatasetItemDb, DatasetMetadataDb
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
+
+    from sqlalchemy import Select
+    from typing_extensions import NotRequired
+
+    from ._storage_client import SqlStorageClient
+
+
+logger = getLogger(__name__)
+
+
+class _DatasetMetadataUpdateParams(MetadataUpdateParams):
+    """Parameters for updating dataset metadata."""
+
+    new_item_count: NotRequired[int]
+    delta_item_count: NotRequired[int]
+
+
+class SqlDatasetClient(DatasetClient, SqlClientMixin):
+    """SQL implementation of the dataset client.
+
+    This client persists dataset items to a SQL database using two tables for storage
+    and retrieval. Items are stored as JSON with automatic ordering preservation.
+
+    The dataset data is stored in SQL database tables following the pattern:
+    - `datasets` table: Contains dataset metadata (id, name, timestamps, item_count)
+    - `dataset_records` table: Contains individual items with JSON data and auto-increment ordering
+
+    Items are stored as a JSON object in SQLite and as JSONB in PostgreSQL. These objects must be JSON-serializable.
+    The `item_id` auto-increment primary key ensures insertion order is preserved.
+    All operations are wrapped in database transactions with CASCADE deletion support.
+    """
+
+    _DEFAULT_NAME = 'default'
+    """Default dataset name used when no name is provided."""
+
+    _METADATA_TABLE = DatasetMetadataDb
+    """SQLAlchemy model for dataset metadata."""
+
+    _ITEM_TABLE = DatasetItemDb
+    """SQLAlchemy model for dataset items."""
+
+    _CLIENT_TYPE = 'Dataset'
+    """Human-readable client type for error messages."""
+
+    def __init__(
+        self,
+        *,
+        id: str,
+        storage_client: SqlStorageClient,
+    ) -> None:
+        """Initialize a new instance.
+
+        Preferably use the `SqlDatasetClient.open` class method to create a new instance.
+        """
+        super().__init__(id=id, storage_client=storage_client)
+
+    @classmethod
+    async def open(
+        cls,
+        *,
+        id: str | None,
+        name: str | None,
+        alias: str | None,
+        storage_client: SqlStorageClient,
+    ) -> SqlDatasetClient:
+        """Open an existing dataset or create a new one.
+
+        Args:
+            id: The ID of the dataset to open. If provided, searches for existing dataset by ID.
+            name: The name of the dataset for named (global scope) storages.
+            alias: The alias of the dataset for unnamed (run scope) storages.
+            storage_client: The SQL storage client instance.
+
+        Returns:
+            An instance for the opened or created storage client.
+
+        Raises:
+            ValueError: If a dataset with the specified ID is not found.
+        """
+        return await cls._safely_open(
+            id=id,
+            name=name,
+            alias=alias,
+            storage_client=storage_client,
+            metadata_model=DatasetMetadata,
+            extra_metadata_fields={'item_count': 0},
+        )
+
+    @override
+    async def get_metadata(self) -> DatasetMetadata:
+        # The database is a single place of truth
+        return await self._get_metadata(DatasetMetadata)
+
+    @override
+    async def drop(self) -> None:
+        """Delete this dataset and all its items from the database.
+
+        This operation is irreversible. Uses CASCADE deletion to remove all related items.
+        """
+        await self._drop()
+
+    @override
+    async def purge(self) -> None:
+        """Remove all items from this dataset while keeping the dataset structure.
+
+        Resets item_count to 0 and deletes all records from dataset_records table.
+        """
+        await self._purge(
+            metadata_kwargs=_DatasetMetadataUpdateParams(
+                new_item_count=0,
+                update_accessed_at=True,
+                update_modified_at=True,
+                force=True,
+            )
+        )
+
+    @override
+    async def push_data(self, data: list[dict[str, Any]] | dict[str, Any]) -> None:
+        if not isinstance(data, list):
+            data = [data]
+
+        db_items: list[dict[str, Any]] = []
+        db_items = [{'dataset_id': self._id, 'data': item} for item in data]
+        stmt = insert(self._ITEM_TABLE).values(db_items)
+
+        async with self.get_session(with_simple_commit=True) as session:
+            await session.execute(stmt)
+
+            await self._update_metadata(
+                session,
+                **_DatasetMetadataUpdateParams(
+                    update_accessed_at=True,
+                    update_modified_at=True,
+                    delta_item_count=len(data),
+                    new_item_count=len(data),
+                    force=True,
+                ),
+            )
+
+    @override
+    async def get_data(
+        self,
+        *,
+        offset: int = 0,
+        limit: int | None = 999_999_999_999,
+        clean: bool = False,
+        desc: bool = False,
+        fields: list[str] | None = None,
+        omit: list[str] | None = None,
+        unwind: list[str] | None = None,
+        skip_empty: bool = False,
+        skip_hidden: bool = False,
+        flatten: list[str] | None = None,
+        view: str | None = None,
+    ) -> DatasetItemsListPage:
+        stmt = self._prepare_get_stmt(
+            offset=offset,
+            limit=limit,
+            clean=clean,
+            desc=desc,
+            fields=fields,
+            omit=omit,
+            unwind=unwind,
+            skip_empty=skip_empty,
+            skip_hidden=skip_hidden,
+            flatten=flatten,
+            view=view,
+        )
+
+        async with self.get_session() as session:
+            result = await session.execute(stmt)
+            db_items = result.scalars().all()
+
+            updated = await self._update_metadata(session, **_DatasetMetadataUpdateParams(update_accessed_at=True))
+
+            # Commit updates to the metadata
+            if updated:
+                await session.commit()
+
+        items = [db_item.data for db_item in db_items]
+        metadata = await self.get_metadata()
+        return DatasetItemsListPage(
+            items=items,
+            count=len(items),
+            desc=desc,
+            limit=limit or 0,
+            offset=offset or 0,
+            total=metadata.item_count,
+        )
+
+    @override
+    async def iterate_items(
+        self,
+        *,
+        offset: int = 0,
+        limit: int | None = None,
+        clean: bool = False,
+        desc: bool = False,
+        fields: list[str] | None = None,
+        omit: list[str] | None = None,
+        unwind: list[str] | None = None,
+        skip_empty: bool = False,
+        skip_hidden: bool = False,
+    ) -> AsyncIterator[dict[str, Any]]:
+        stmt = self._prepare_get_stmt(
+            offset=offset,
+            limit=limit,
+            clean=clean,
+            desc=desc,
+            fields=fields,
+            omit=omit,
+            unwind=unwind,
+            skip_empty=skip_empty,
+            skip_hidden=skip_hidden,
+        )
+
+        async with self.get_session() as session:
+            db_items = await session.stream_scalars(stmt)
+
+            async for db_item in db_items:
+                yield db_item.data
+
+            updated = await self._update_metadata(session, **_DatasetMetadataUpdateParams(update_accessed_at=True))
+
+            # Commit updates to the metadata
+            if updated:
+                await session.commit()
+
+    def _prepare_get_stmt(
+        self,
+        *,
+        offset: int = 0,
+        limit: int | None = 999_999_999_999,
+        clean: bool = False,
+        desc: bool = False,
+        fields: list[str] | None = None,
+        omit: list[str] | None = None,
+        unwind: list[str] | None = None,
+        skip_empty: bool = False,
+        skip_hidden: bool = False,
+        flatten: list[str] | None = None,
+        view: str | None = None,
+    ) -> Select:
+        # Check for unsupported arguments and log a warning if found.
+        unsupported_args: dict[str, Any] = {
+            'clean': clean,
+            'fields': fields,
+            'omit': omit,
+            'unwind': unwind,
+            'skip_hidden': skip_hidden,
+            'flatten': flatten,
+            'view': view,
+        }
+        unsupported = {k: v for k, v in unsupported_args.items() if v not in (False, None)}
+
+        if unsupported:
+            logger.warning(
+                f'The arguments {list(unsupported.keys())} of get_data are not supported by the '
+                f'{self.__class__.__name__} client.'
+            )
+
+        stmt = select(self._ITEM_TABLE).where(self._ITEM_TABLE.dataset_id == self._id)
+
+        if skip_empty:
+            # Skip items that are empty JSON objects
+            stmt = stmt.where(self._ITEM_TABLE.data != {})
+
+        # Apply ordering by insertion order (item_id)
+        stmt = stmt.order_by(self._ITEM_TABLE.item_id.desc()) if desc else stmt.order_by(self._ITEM_TABLE.item_id.asc())
+
+        return stmt.offset(offset).limit(limit)
+
+    def _specific_update_metadata(
+        self,
+        new_item_count: int | None = None,
+        delta_item_count: int | None = None,
+        **_kwargs: dict[str, Any],
+    ) -> dict[str, Any]:
+        """Update the dataset metadata in the database.
+
+        Args:
+            session: The SQLAlchemy AsyncSession to use for the update.
+            new_item_count: If provided, set item count to this value.
+            delta_item_count: If provided, add this value to the current item count.
+        """
+        values_to_set: dict[str, Any] = {}
+
+        if new_item_count is not None:
+            values_to_set['item_count'] = new_item_count
+        elif delta_item_count:
+            # Use database-level for atomic updates
+            values_to_set['item_count'] = self._METADATA_TABLE.item_count + delta_item_count
+
+        return values_to_set

crawlee/storage_clients/_sql/_db_models.py

@@ -0,0 +1,269 @@
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any
+
+from sqlalchemy import JSON, BigInteger, Boolean, ForeignKey, Index, Integer, LargeBinary, String, text
+from sqlalchemy.dialects.postgresql import JSONB
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship, synonym
+from sqlalchemy.types import DateTime, TypeDecorator
+from typing_extensions import override
+
+if TYPE_CHECKING:
+    from sqlalchemy.engine import Dialect
+    from sqlalchemy.types import TypeEngine
+
+
+class AwareDateTime(TypeDecorator):
+    """Custom SQLAlchemy type for timezone-aware datetime handling.
+
+    Ensures all datetime values are timezone-aware by adding UTC timezone to
+    naive datetime values from databases that don't store timezone information.
+    """
+
+    impl = DateTime(timezone=True)
+    cache_ok = True
+
+    @override
+    def process_result_value(self, value: datetime | None, dialect: Dialect) -> datetime | None:
+        """Add UTC timezone to naive datetime values."""
+        if value is not None and value.tzinfo is None:
+            return value.replace(tzinfo=timezone.utc)
+        return value
+
+
+class JsonField(TypeDecorator):
+    """Uses JSONB for PostgreSQL and JSON for other databases."""
+
+    impl = JSON
+    cache_ok = True
+
+    def load_dialect_impl(self, dialect: Dialect) -> TypeEngine[JSON | JSONB]:
+        """Load the appropriate dialect implementation for the JSON type."""
+        if dialect.name == 'postgresql':
+            return dialect.type_descriptor(JSONB())
+        return dialect.type_descriptor(JSON())
+
+
+class Base(DeclarativeBase):
+    """Base class for all database models for correct type annotations."""
+
+
+class StorageMetadataDb:
+    """Base database model for storage metadata."""
+
+    internal_name: Mapped[str] = mapped_column(String, nullable=False, index=True, unique=True)
+    """Internal unique name for a storage instance based on a name or alias."""
+
+    name: Mapped[str | None] = mapped_column(String, nullable=True, unique=True)
+    """Human-readable name. None becomes 'default' in database to enforce uniqueness."""
+
+    accessed_at: Mapped[datetime] = mapped_column(AwareDateTime, nullable=False)
+    """Last access datetime for usage tracking."""
+
+    created_at: Mapped[datetime] = mapped_column(AwareDateTime, nullable=False)
+    """Creation datetime."""
+
+    modified_at: Mapped[datetime] = mapped_column(AwareDateTime, nullable=False)
+    """Last modification datetime."""
+
+
+class DatasetMetadataDb(StorageMetadataDb, Base):
+    """Metadata table for datasets."""
+
+    __tablename__ = 'datasets'
+
+    dataset_id: Mapped[str] = mapped_column(String(20), nullable=False, primary_key=True)
+    """Unique identifier for the dataset."""
+
+    item_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
+    """Number of items in the dataset."""
+
+    # Relationship to dataset items with cascade deletion
+    items: Mapped[list[DatasetItemDb]] = relationship(
+        back_populates='dataset', cascade='all, delete-orphan', lazy='noload'
+    )
+
+    id = synonym('dataset_id')
+    """Alias for dataset_id to match Pydantic expectations."""
+
+
+class RequestQueueMetadataDb(StorageMetadataDb, Base):
+    """Metadata table for request queues."""
+
+    __tablename__ = 'request_queues'
+
+    request_queue_id: Mapped[str] = mapped_column(String(20), nullable=False, primary_key=True)
+    """Unique identifier for the request queue."""
+
+    had_multiple_clients: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)
+    """Flag indicating if multiple clients have accessed this queue."""
+
+    handled_request_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
+    """Number of requests processed."""
+
+    pending_request_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
+    """Number of requests waiting to be processed."""
+
+    total_request_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
+    """Total number of requests ever added to this queue."""
+
+    # Relationship to queue requests with cascade deletion
+    requests: Mapped[list[RequestDb]] = relationship(
+        back_populates='queue', cascade='all, delete-orphan', lazy='noload'
+    )
+    # Relationship to queue state
+    state: Mapped[RequestQueueStateDb] = relationship(
+        back_populates='queue', cascade='all, delete-orphan', lazy='noload'
+    )
+
+    id = synonym('request_queue_id')
+    """Alias for request_queue_id to match Pydantic expectations."""
+
+
+class KeyValueStoreMetadataDb(StorageMetadataDb, Base):
+    """Metadata table for key-value stores."""
+
+    __tablename__ = 'key_value_stores'
+
+    key_value_store_id: Mapped[str] = mapped_column(String(20), nullable=False, primary_key=True)
+    """Unique identifier for the key-value store."""
+
+    # Relationship to store records with cascade deletion
+    records: Mapped[list[KeyValueStoreRecordDb]] = relationship(
+        back_populates='kvs', cascade='all, delete-orphan', lazy='noload'
+    )
+
+    id = synonym('key_value_store_id')
+    """Alias for key_value_store_id to match Pydantic expectations."""
+
+
+class KeyValueStoreRecordDb(Base):
+    """Records table for key-value stores."""
+
+    __tablename__ = 'key_value_store_records'
+
+    key_value_store_id: Mapped[str] = mapped_column(
+        String(20),
+        ForeignKey('key_value_stores.key_value_store_id', ondelete='CASCADE'),
+        primary_key=True,
+        index=True,
+        nullable=False,
+    )
+    """Foreign key to metadata key-value store record."""
+
+    key: Mapped[str] = mapped_column(String(255), primary_key=True)
+    """The key part of the key-value pair."""
+
+    value: Mapped[bytes] = mapped_column(LargeBinary, nullable=False)
+    """Value stored as binary data to support any content type."""
+
+    content_type: Mapped[str] = mapped_column(String(50), nullable=False)
+    """MIME type for proper value deserialization."""
+
+    size: Mapped[int | None] = mapped_column(Integer, nullable=False, default=0)
+    """Size of stored value in bytes."""
+
+    # Relationship back to parent store
+    kvs: Mapped[KeyValueStoreMetadataDb] = relationship(back_populates='records')
+
+    storage_id = synonym('key_value_store_id')
+    """Alias for key_value_store_id to match SqlClientMixin expectations."""
+
+
+class DatasetItemDb(Base):
+    """Items table for datasets."""
+
+    __tablename__ = 'dataset_records'
+
+    item_id: Mapped[int] = mapped_column(Integer, primary_key=True)
+    """Auto-increment primary key preserving insertion order."""
+
+    dataset_id: Mapped[str] = mapped_column(
+        String(20),
+        ForeignKey('datasets.dataset_id', ondelete='CASCADE'),
+        index=True,
+    )
+    """Foreign key to metadata dataset record."""
+
+    data: Mapped[list[dict[str, Any]] | dict[str, Any]] = mapped_column(JsonField, nullable=False)
+    """JSON serializable item data."""
+
+    # Relationship back to parent dataset
+    dataset: Mapped[DatasetMetadataDb] = relationship(back_populates='items')
+
+    storage_id = synonym('dataset_id')
+    """Alias for dataset_id to match SqlClientMixin expectations."""
+
+
+class RequestDb(Base):
+    """Requests table for request queues."""
+
+    __tablename__ = 'request_queue_records'
+    __table_args__ = (
+        Index(
+            'idx_fetch_available',
+            'request_queue_id',
+            'is_handled',
+            'time_blocked_until',
+            'sequence_number',
+            postgresql_where=text('is_handled = false'),
+        ),
+    )
+
+    request_id: Mapped[int] = mapped_column(BigInteger, primary_key=True)
+    """Unique identifier for the request representing the unique_key."""
+
+    request_queue_id: Mapped[str] = mapped_column(
+        String(20), ForeignKey('request_queues.request_queue_id', ondelete='CASCADE'), primary_key=True
+    )
+    """Foreign key to metadata request queue record."""
+
+    data: Mapped[str] = mapped_column(String, nullable=False)
+    """JSON-serialized Request object."""
+
+    sequence_number: Mapped[int] = mapped_column(Integer, nullable=False)
+    """Ordering sequence: negative for forefront, positive for regular."""
+
+    is_handled: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)
+    """Processing status flag."""
+
+    time_blocked_until: Mapped[datetime | None] = mapped_column(AwareDateTime, nullable=True)
+    """Timestamp until which this request is considered blocked for processing by other clients."""
+
+    client_key: Mapped[str | None] = mapped_column(String(32), nullable=True)
+    """Identifier of the client that has currently locked this request for processing."""
+
+    # Relationship back to metadata table
+    queue: Mapped[RequestQueueMetadataDb] = relationship(back_populates='requests')
+
+    storage_id = synonym('request_queue_id')
+    """Alias for request_queue_id to match SqlClientMixin expectations."""
+
+
+class RequestQueueStateDb(Base):
+    """State table for request queues."""
+
+    __tablename__ = 'request_queue_state'
+
+    request_queue_id: Mapped[str] = mapped_column(
+        String(20), ForeignKey('request_queues.request_queue_id', ondelete='CASCADE'), primary_key=True
+    )
+    """Foreign key to metadata request queue record."""
+
+    sequence_counter: Mapped[int] = mapped_column(Integer, nullable=False, default=1)
+    """Counter for regular request ordering (positive)."""
+
+    forefront_sequence_counter: Mapped[int] = mapped_column(Integer, nullable=False, default=-1)
+    """Counter for forefront request ordering (negative)."""
+
+    # Relationship back to metadata table
+    queue: Mapped[RequestQueueMetadataDb] = relationship(back_populates='state')
+
+
+class VersionDb(Base):
+    """Table for storing the database schema version."""
+
+    __tablename__ = 'version'
+
+    version: Mapped[str] = mapped_column(String(10), nullable=False, primary_key=True)