sqlspec 0.36.0__cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl

sqlspec/adapters/bigquery/events/store.py

@@ -0,0 +1,139 @@
+"""BigQuery event queue store with clustering optimization.
+
+BigQuery uses clustered tables instead of indexes for query optimization.
+The event queue table is clustered by (channel, status, available_at) to
+optimize polling queries that filter by channel and status.
+
+Configuration:
+    extension_config={
+        "events": {
+            "queue_table": "my_events"  # Default: "sqlspec_event_queue"
+        }
+    }
+"""
+
+from sqlspec.adapters.bigquery.config import BigQueryConfig
+from sqlspec.extensions.events import BaseEventQueueStore
+
+__all__ = ("BigQueryEventQueueStore",)
+
+
+class BigQueryEventQueueStore(BaseEventQueueStore[BigQueryConfig]):
+    """BigQuery-specific event queue store with clustering optimization.
+
+    Generates DDL optimized for BigQuery. BigQuery does not support traditional
+    indexes, so the table uses CLUSTER BY for query optimization instead.
+
+    Args:
+        config: BigQueryConfig with extension_config["events"] settings.
+
+    Notes:
+        Configuration is read from config.extension_config["events"]:
+        - queue_table: Table name (default: "sqlspec_event_queue")
+
+        BigQuery-specific optimizations:
+        - Uses STRING instead of VARCHAR (BigQuery's native string type)
+        - Uses INT64 instead of INTEGER
+        - Uses CLUSTER BY instead of CREATE INDEX
+        - Supports IF NOT EXISTS / IF EXISTS in DDL
+
+    Example:
+        from sqlspec.adapters.bigquery import BigQueryConfig
+        from sqlspec.adapters.bigquery.events import BigQueryEventQueueStore
+
+        config = BigQueryConfig(
+            connection_config={"project": "my-project"},
+            extension_config={"events": {"queue_table": "my_events"}}
+        )
+        store = BigQueryEventQueueStore(config)
+        for stmt in store.create_statements():
+            driver.execute_script(stmt)
+    """
+
+    __slots__ = ()
+
+    def _column_types(self) -> "tuple[str, str, str]":
+        """Return BigQuery-specific column types.
+
+        Returns:
+            Tuple of (payload_type, metadata_type, timestamp_type).
+        """
+        return "JSON", "JSON", "TIMESTAMP"
+
+    def _string_type(self, length: int) -> str:
+        """Return BigQuery STRING type (length is ignored)."""
+        del length
+        return "STRING"
+
+    def _integer_type(self) -> str:
+        """Return BigQuery INT64 type."""
+        return "INT64"
+
+    def _timestamp_default(self) -> str:
+        """Return BigQuery timestamp default expression."""
+        return "CURRENT_TIMESTAMP()"
+
+    def _table_clause(self) -> str:
+        """Return BigQuery CLUSTER BY clause for query optimization."""
+        return " CLUSTER BY channel, status, available_at"
+
+    def _build_create_table_sql(self) -> str:
+        """Build BigQuery CREATE TABLE with CLUSTER BY optimization.
+
+        BigQuery uses CLUSTER BY for query optimization instead of indexes.
+        The clustering columns match the typical polling query pattern.
+
+        Note: BigQuery does not support column-level PRIMARY KEY, so we
+        omit it entirely. event_id uniqueness must be enforced at insert time.
+        """
+        payload_type, metadata_type, timestamp_type = self._column_types()
+        string_type = self._string_type(0)
+        integer_type = self._integer_type()
+        ts_default = self._timestamp_default()
+        table_clause = self._table_clause()
+
+        return (
+            f"CREATE TABLE IF NOT EXISTS {self.table_name} ("
+            f"event_id {string_type} NOT NULL,"
+            f" channel {string_type} NOT NULL,"
+            f" payload_json {payload_type} NOT NULL,"
+            f" metadata_json {metadata_type},"
+            f" status {string_type} NOT NULL DEFAULT 'pending',"
+            f" available_at {timestamp_type} NOT NULL DEFAULT {ts_default},"
+            f" lease_expires_at {timestamp_type},"
+            f" attempts {integer_type} NOT NULL DEFAULT 0,"
+            f" created_at {timestamp_type} NOT NULL DEFAULT {ts_default},"
+            f" acknowledged_at {timestamp_type}"
+            f"){table_clause}"
+        )
+
+    def _build_index_sql(self) -> str | None:
+        """Return None since BigQuery uses CLUSTER BY instead of indexes.
+
+        Returns:
+            None, as BigQuery does not support traditional indexes.
+        """
+        return None
+
+    def create_statements(self) -> "list[str]":
+        """Return DDL statement for table creation.
+
+        Returns:
+            List containing single CREATE TABLE statement.
+
+        Notes:
+            BigQuery uses CLUSTER BY instead of separate index creation,
+            so only one statement is returned.
+        """
+        return [self._build_create_table_sql()]
+
+    def drop_statements(self) -> "list[str]":
+        """Return DDL statement for table deletion.
+
+        Returns:
+            List containing single DROP TABLE statement.
+
+        Notes:
+            BigQuery has no index to drop, only the table.
+        """
+        return [f"DROP TABLE IF EXISTS {self.table_name}"]

sqlspec/adapters/bigquery/litestar/__init__.py

@@ -0,0 +1,5 @@
+"""BigQuery Litestar integration."""
+
+from sqlspec.adapters.bigquery.litestar.store import BigQueryStore
+
+__all__ = ("BigQueryStore",)

sqlspec/adapters/bigquery/litestar/store.py

@@ -0,0 +1,325 @@
+"""BigQuery session store for Litestar integration."""
+
+from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING
+
+from sqlspec.extensions.litestar.store import BaseSQLSpecStore
+from sqlspec.utils.sync_tools import async_
+
+if TYPE_CHECKING:
+    from sqlspec.adapters.bigquery.config import BigQueryConfig
+
+
+__all__ = ("BigQueryStore",)
+
+
+class BigQueryStore(BaseSQLSpecStore["BigQueryConfig"]):
+    """BigQuery session store using synchronous BigQuery driver.
+
+    Implements server-side session storage for Litestar using Google BigQuery.
+    Uses Litestar's sync_to_thread utility to provide an async interface
+    compatible with the Store protocol.
+
+    Provides efficient session management with:
+    - Sync operations wrapped for async compatibility
+    - MERGE for UPSERT functionality
+    - Native TIMESTAMP type support
+    - Automatic expiration handling
+    - Efficient cleanup of expired sessions
+    - Table clustering on session_id for optimized lookups
+
+    Note:
+        BigQuery is designed for analytical (OLAP) workloads and scales to petabytes.
+        For typical session store workloads, clustering by session_id provides good
+        performance. Consider partitioning by created_at if session volume exceeds
+        millions of rows per day.
+
+    Args:
+        config: BigQueryConfig instance.
+
+    Example:
+        from sqlspec.adapters.bigquery import BigQueryConfig
+        from sqlspec.adapters.bigquery.litestar.store import BigQueryStore
+
+        config = BigQueryConfig(connection_config={"project": "my-project"})
+        store = BigQueryStore(config)
+        await store.create_table()
+    """
+
+    __slots__ = ()
+
+    def __init__(self, config: "BigQueryConfig") -> None:
+        """Initialize BigQuery session store.
+
+        Args:
+            config: BigQueryConfig instance.
+
+        Notes:
+            Table name is read from config.extension_config["litestar"]["session_table"].
+        """
+        super().__init__(config)
+
+    def _get_create_table_sql(self) -> str:
+        """Get BigQuery CREATE TABLE SQL with optimized schema.
+
+        Returns:
+            SQL statement to create the sessions table with clustering.
+
+        Notes:
+            - Uses TIMESTAMP for timezone-aware expiration timestamps
+            - BYTES for binary session data storage
+            - Clustered by session_id for efficient lookups
+            - No indexes needed - BigQuery uses columnar storage
+            - Table name is internally controlled, not user input
+        """
+        return f"""
+        CREATE TABLE IF NOT EXISTS {self._table_name} (
+            session_id STRING NOT NULL,
+            data BYTES NOT NULL,
+            expires_at TIMESTAMP,
+            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP()
+        )
+        CLUSTER BY session_id
+        """
+
+    def _get_drop_table_sql(self) -> "list[str]":
+        """Get BigQuery DROP TABLE SQL statements.
+
+        Returns:
+            List containing DROP TABLE statement.
+
+        Notes:
+            BigQuery doesn't have separate indexes to drop.
+        """
+        return [f"DROP TABLE IF EXISTS {self._table_name}"]
+
+    def _datetime_to_timestamp(self, dt: "datetime | None") -> "datetime | None":
+        """Convert datetime to BigQuery TIMESTAMP.
+
+        Args:
+            dt: Datetime to convert (must be UTC-aware).
+
+        Returns:
+            UTC datetime object, or None if dt is None.
+
+        Notes:
+            BigQuery TIMESTAMP type expects UTC datetime objects.
+            The BigQuery client library handles the conversion.
+        """
+        if dt is None:
+            return None
+        if dt.tzinfo is None:
+            return dt.replace(tzinfo=timezone.utc)
+        return dt
+
+    def _timestamp_to_datetime(self, ts: "datetime | None") -> "datetime | None":
+        """Convert TIMESTAMP back to datetime.
+
+        Args:
+            ts: Datetime object from BigQuery.
+
+        Returns:
+            UTC-aware datetime, or None if ts is None.
+        """
+        if ts is None:
+            return None
+        if ts.tzinfo is None:
+            return ts.replace(tzinfo=timezone.utc)
+        return ts
+
+    def _create_table(self) -> None:
+        """Synchronous implementation of create_table."""
+        sql = self._get_create_table_sql()
+        with self._config.provide_session() as driver:
+            driver.execute_script(sql)
+        self._log_table_created()
+
+    async def create_table(self) -> None:
+        """Create the session table if it doesn't exist."""
+        await async_(self._create_table)()
+
+    def _get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | None":
+        """Synchronous implementation of get."""
+        sql = f"""
+        SELECT data, expires_at FROM {self._table_name}
+        WHERE session_id = @session_id
+        AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP())
+        """
+
+        with self._config.provide_session() as driver:
+            result = driver.select_one(sql, {"session_id": key})
+
+            if result is None:
+                return None
+
+            data = result.get("data")
+            expires_at = result.get("expires_at")
+
+            if renew_for is not None and expires_at is not None:
+                new_expires_at = self._calculate_expires_at(renew_for)
+                new_expires_at_ts = self._datetime_to_timestamp(new_expires_at)
+                if new_expires_at_ts is not None:
+                    update_sql = f"""
+                    UPDATE {self._table_name}
+                    SET expires_at = @expires_at
+                    WHERE session_id = @session_id
+                    """
+                    driver.execute(update_sql, {"expires_at": new_expires_at_ts, "session_id": key})
+
+            return bytes(data) if data is not None else None
+
+    async def get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | None":
+        """Get a session value by key.
+
+        Args:
+            key: Session ID to retrieve.
+            renew_for: If given, renew the expiry time for this duration.
+
+        Returns:
+            Session data as bytes if found and not expired, None otherwise.
+        """
+        return await async_(self._get)(key, renew_for)
+
+    def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | None" = None) -> None:
+        """Synchronous implementation of set.
+
+        Notes:
+            Uses MERGE for UPSERT functionality in BigQuery.
+            BigQuery requires source data to come from a table or inline VALUES.
+        """
+        data = self._value_to_bytes(value)
+        expires_at = self._calculate_expires_at(expires_in)
+        expires_at_ts = self._datetime_to_timestamp(expires_at)
+
+        sql = f"""
+        MERGE {self._table_name} AS target
+        USING (SELECT @session_id AS session_id, @data AS data, @expires_at AS expires_at) AS source
+        ON target.session_id = source.session_id
+        WHEN MATCHED THEN
+            UPDATE SET data = source.data, expires_at = source.expires_at
+        WHEN NOT MATCHED THEN
+            INSERT (session_id, data, expires_at, created_at)
+            VALUES (source.session_id, source.data, source.expires_at, CURRENT_TIMESTAMP())
+        """
+
+        with self._config.provide_session() as driver:
+            driver.execute(sql, {"session_id": key, "data": data, "expires_at": expires_at_ts})
+
+    async def set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | None" = None) -> None:
+        """Store a session value.
+
+        Args:
+            key: Session ID.
+            value: Session data.
+            expires_in: Time until expiration.
+        """
+        await async_(self._set)(key, value, expires_in)
+
+    def _delete(self, key: str) -> None:
+        """Synchronous implementation of delete."""
+        sql = f"DELETE FROM {self._table_name} WHERE session_id = @session_id"
+
+        with self._config.provide_session() as driver:
+            driver.execute(sql, {"session_id": key})
+
+    async def delete(self, key: str) -> None:
+        """Delete a session by key.
+
+        Args:
+            key: Session ID to delete.
+        """
+        await async_(self._delete)(key)
+
+    def _delete_all(self) -> None:
+        """Synchronous implementation of delete_all."""
+        sql = f"DELETE FROM {self._table_name} WHERE TRUE"
+
+        with self._config.provide_session() as driver:
+            driver.execute(sql)
+        self._log_delete_all()
+
+    async def delete_all(self) -> None:
+        """Delete all sessions from the store."""
+        await async_(self._delete_all)()
+
+    def _exists(self, key: str) -> bool:
+        """Synchronous implementation of exists."""
+        sql = f"""
+        SELECT 1 FROM {self._table_name}
+        WHERE session_id = @session_id
+        AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP())
+        LIMIT 1
+        """
+
+        with self._config.provide_session() as driver:
+            result = driver.select_one(sql, {"session_id": key})
+            return result is not None
+
+    async def exists(self, key: str) -> bool:
+        """Check if a session key exists and is not expired.
+
+        Args:
+            key: Session ID to check.
+
+        Returns:
+            True if the session exists and is not expired.
+        """
+        return await async_(self._exists)(key)
+
+    def _expires_in(self, key: str) -> "int | None":
+        """Synchronous implementation of expires_in."""
+        sql = f"""
+        SELECT expires_at FROM {self._table_name}
+        WHERE session_id = @session_id
+        """
+
+        with self._config.provide_session() as driver:
+            result = driver.select_one(sql, {"session_id": key})
+
+            if result is None:
+                return None
+
+            expires_at = result.get("expires_at")
+            if expires_at is None:
+                return None
+
+            expires_at_dt = self._timestamp_to_datetime(expires_at)
+            if expires_at_dt is None:
+                return None
+
+            now = datetime.now(timezone.utc)
+            if expires_at_dt <= now:
+                return 0
+
+            delta = expires_at_dt - now
+            return int(delta.total_seconds())
+
+    async def expires_in(self, key: str) -> "int | None":
+        """Get the time in seconds until the session expires.
+
+        Args:
+            key: Session ID to check.
+
+        Returns:
+            Seconds until expiration, or None if no expiry or key doesn't exist.
+        """
+        return await async_(self._expires_in)(key)
+
+    def _delete_expired(self) -> int:
+        """Synchronous implementation of delete_expired."""
+        sql = f"DELETE FROM {self._table_name} WHERE expires_at <= CURRENT_TIMESTAMP()"
+
+        with self._config.provide_session() as driver:
+            result = driver.execute(sql)
+            count = result.get_affected_count()
+            if count > 0:
+                self._log_delete_expired(count)
+            return count
+
+    async def delete_expired(self) -> int:
+        """Delete all expired sessions.
+
+        Returns:
+            Number of sessions deleted.
+        """
+        return await async_(self._delete_expired)()

sqlspec/adapters/bigquery/type_converter.py

@@ -0,0 +1,107 @@
+"""BigQuery-specific type conversion with native UUID support.
+
+Provides specialized type handling for BigQuery, including UUID support
+for the native BigQuery driver and parameter creation.
+"""
+
+from typing import Any
+from uuid import UUID
+
+from sqlspec.core.type_converter import CachedOutputConverter, convert_uuid
+
+__all__ = ("BIGQUERY_SPECIAL_CHARS", "BQ_TYPE_MAP", "BigQueryOutputConverter")
+
+BQ_TYPE_MAP: "dict[str, str]" = {
+    "str": "STRING",
+    "int": "INT64",
+    "float": "FLOAT64",
+    "bool": "BOOL",
+    "datetime": "DATETIME",
+    "date": "DATE",
+    "time": "TIME",
+    "UUID": "STRING",
+    "uuid": "STRING",
+    "Decimal": "NUMERIC",
+    "bytes": "BYTES",
+    "list": "ARRAY",
+    "dict": "STRUCT",
+}
+
+BIGQUERY_SPECIAL_CHARS: "frozenset[str]" = frozenset({"{", "[", "-", ":", "T", "."})
+
+
+class BigQueryOutputConverter(CachedOutputConverter):
+    """BigQuery-specific output conversion with native UUID support.
+
+    Extends CachedOutputConverter with BigQuery-specific functionality
+    including UUID handling and parameter creation for the native BigQuery driver.
+    """
+
+    __slots__ = ("_enable_uuid_conversion",)
+
+    def __init__(self, cache_size: int = 5000, enable_uuid_conversion: bool = True) -> None:
+        """Initialize converter with BigQuery-specific options.
+
+        Args:
+            cache_size: Maximum number of string values to cache (default: 5000)
+            enable_uuid_conversion: Enable automatic UUID string conversion (default: True)
+        """
+        super().__init__(special_chars=BIGQUERY_SPECIAL_CHARS, cache_size=cache_size)
+        self._enable_uuid_conversion = enable_uuid_conversion
+
+    def _convert_detected(self, value: str, detected_type: str) -> Any:
+        """Convert value with BigQuery-specific handling.
+
+        Args:
+            value: String value to convert.
+            detected_type: Detected type name.
+
+        Returns:
+            Converted value, respecting UUID conversion setting.
+        """
+        try:
+            return self.convert_value(value, detected_type)
+        except Exception:
+            return value
+
+    def create_parameter(self, name: str, value: Any) -> "Any | None":
+        """Create BigQuery parameter with proper type mapping.
+
+        Args:
+            name: Parameter name.
+            value: Parameter value.
+
+        Returns:
+            ScalarQueryParameter for native BigQuery driver, None if not available.
+        """
+        try:
+            from google.cloud.bigquery import ScalarQueryParameter
+        except ImportError:
+            return None
+
+        if self._enable_uuid_conversion:
+            if isinstance(value, UUID):
+                return ScalarQueryParameter(name, "STRING", str(value))
+
+            if isinstance(value, str):
+                detected_type = self.detect_type(value)
+                if detected_type == "uuid":
+                    uuid_obj = convert_uuid(value)
+                    return ScalarQueryParameter(name, "STRING", str(uuid_obj))
+
+        param_type = BQ_TYPE_MAP.get(type(value).__name__, "STRING")
+        return ScalarQueryParameter(name, param_type, value)
+
+    def convert_bigquery_value(self, value: Any, column_type: str) -> Any:
+        """Convert BigQuery value based on column type.
+
+        Args:
+            value: Value to convert.
+            column_type: BigQuery column type.
+
+        Returns:
+            Converted value appropriate for the column type.
+        """
+        if column_type == "STRING" and isinstance(value, str):
+            return self.convert(value)
+        return value

sqlspec/adapters/cockroach_asyncpg/__init__.py

@@ -0,0 +1,24 @@
+from sqlspec.adapters.cockroach_asyncpg._typing import (
+    CockroachAsyncpgConnection,
+    CockroachAsyncpgPool,
+    CockroachAsyncpgSessionContext,
+)
+from sqlspec.adapters.cockroach_asyncpg.config import (
+    CockroachAsyncpgConfig,
+    CockroachAsyncpgConnectionConfig,
+    CockroachAsyncpgDriverFeatures,
+    CockroachAsyncpgPoolConfig,
+)
+from sqlspec.adapters.cockroach_asyncpg.driver import CockroachAsyncpgDriver, CockroachAsyncpgExceptionHandler
+
+__all__ = (
+    "CockroachAsyncpgConfig",
+    "CockroachAsyncpgConnection",
+    "CockroachAsyncpgConnectionConfig",
+    "CockroachAsyncpgDriver",
+    "CockroachAsyncpgDriverFeatures",
+    "CockroachAsyncpgExceptionHandler",
+    "CockroachAsyncpgPool",
+    "CockroachAsyncpgPoolConfig",
+    "CockroachAsyncpgSessionContext",
+)

sqlspec/adapters/cockroach_asyncpg/_typing.py

@@ -0,0 +1,72 @@
+"""CockroachDB AsyncPG adapter type definitions."""
+
+from typing import TYPE_CHECKING, Any
+
+from asyncpg.pool import PoolConnectionProxy
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from typing import TypeAlias
+
+    from asyncpg import Connection, Pool, Record
+
+    from sqlspec.adapters.cockroach_asyncpg.driver import CockroachAsyncpgDriver
+    from sqlspec.core import StatementConfig
+
+    CockroachAsyncpgConnection: TypeAlias = Connection[Record] | PoolConnectionProxy[Record]
+    CockroachAsyncpgPool: TypeAlias = Pool[Record]
+else:
+    from asyncpg import Pool
+
+    CockroachAsyncpgConnection = PoolConnectionProxy
+    CockroachAsyncpgPool = Pool
+
+
+class CockroachAsyncpgSessionContext:
+    """Async context manager for CockroachDB AsyncPG sessions."""
+
+    __slots__ = (
+        "_acquire_connection",
+        "_connection",
+        "_driver",
+        "_driver_features",
+        "_prepare_driver",
+        "_release_connection",
+        "_statement_config",
+    )
+
+    def __init__(
+        self,
+        acquire_connection: "Callable[[], Any]",
+        release_connection: "Callable[[Any], Any]",
+        statement_config: "StatementConfig",
+        driver_features: "dict[str, Any]",
+        prepare_driver: "Callable[[CockroachAsyncpgDriver], CockroachAsyncpgDriver]",
+    ) -> None:
+        self._acquire_connection = acquire_connection
+        self._release_connection = release_connection
+        self._statement_config = statement_config
+        self._driver_features = driver_features
+        self._prepare_driver = prepare_driver
+        self._connection: Any = None
+        self._driver: CockroachAsyncpgDriver | None = None
+
+    async def __aenter__(self) -> "CockroachAsyncpgDriver":
+        from sqlspec.adapters.cockroach_asyncpg.driver import CockroachAsyncpgDriver
+
+        self._connection = await self._acquire_connection()
+        self._driver = CockroachAsyncpgDriver(
+            connection=self._connection, statement_config=self._statement_config, driver_features=self._driver_features
+        )
+        return self._prepare_driver(self._driver)
+
+    async def __aexit__(
+        self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: Any
+    ) -> "bool | None":
+        if self._connection is not None:
+            await self._release_connection(self._connection)
+            self._connection = None
+        return None
+
+
+__all__ = ("CockroachAsyncpgConnection", "CockroachAsyncpgPool", "CockroachAsyncpgSessionContext")

sqlspec/adapters/cockroach_asyncpg/adk/__init__.py

@@ -0,0 +1,3 @@
+from sqlspec.adapters.cockroach_asyncpg.adk.store import CockroachAsyncpgADKMemoryStore, CockroachAsyncpgADKStore
+
+__all__ = ("CockroachAsyncpgADKMemoryStore", "CockroachAsyncpgADKStore")