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

sqlspec/adapters/asyncpg/events/backend.py

@@ -0,0 +1,286 @@
+# pyright: reportPrivateUsage=false
+"""Native and hybrid PostgreSQL backends for EventChannel."""
+
+import asyncio
+import contextlib
+import logging
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any, cast
+
+from sqlspec.core import SQL
+from sqlspec.exceptions import EventChannelError, ImproperConfigurationError
+from sqlspec.extensions.events import (
+    AsyncTableEventQueue,
+    EventMessage,
+    build_queue_backend,
+    decode_notify_payload,
+    encode_notify_payload,
+    normalize_event_channel_name,
+)
+from sqlspec.utils.logging import get_logger, log_with_context
+from sqlspec.utils.serializers import to_json
+from sqlspec.utils.type_guards import has_add_listener, has_notifies, is_notification
+from sqlspec.utils.uuids import uuid4
+
+if TYPE_CHECKING:
+    from sqlspec.adapters.asyncpg.config import AsyncpgConfig
+
+logger = get_logger("sqlspec.events.postgres")
+
+__all__ = ("AsyncpgEventsBackend", "AsyncpgHybridEventsBackend", "create_event_backend")
+
+
+class _AsyncpgNotificationListener:
+    __slots__ = ("_channel", "_future", "_loop")
+
+    def __init__(self, channel: str, future: "asyncio.Future[str]", loop: "asyncio.AbstractEventLoop") -> None:
+        self._channel = channel
+        self._future = future
+        self._loop = loop
+
+    def __call__(self, _conn: Any, _pid: int, notified_channel: str, payload: str) -> None:
+        if notified_channel != self._channel or self._future.done():
+            return
+        self._loop.call_soon_threadsafe(self._future.set_result, payload)
+
+
+class AsyncpgHybridEventsBackend:
+    """Hybrid backend combining durable queue with LISTEN/NOTIFY wakeups."""
+
+    __slots__ = ("_config", "_listen_connection", "_listen_connection_cm", "_queue", "_runtime")
+
+    supports_sync = False
+    supports_async = True
+    backend_name = "listen_notify_durable"
+
+    def __init__(self, config: "AsyncpgConfig", queue: "AsyncTableEventQueue") -> None:
+        if not config.is_async:
+            msg = "Asyncpg hybrid backend requires an async adapter"
+            raise ImproperConfigurationError(msg)
+        self._config = config
+        self._runtime = config.get_observability_runtime()
+        self._queue = queue
+        self._listen_connection: Any | None = None
+        self._listen_connection_cm: Any | None = None
+        log_with_context(
+            logger,
+            logging.DEBUG,
+            "event.listen",
+            adapter_name="asyncpg",
+            backend_name=self.backend_name,
+            mode="async",
+            status="backend_ready",
+        )
+
+    async def publish(self, channel: str, payload: "dict[str, Any]", metadata: "dict[str, Any] | None" = None) -> str:
+        event_id = uuid4().hex
+        await self._publish_durable(channel, event_id, payload, metadata)
+        self._runtime.increment_metric("events.publish.native")
+        return event_id
+
+    async def dequeue(self, channel: str, poll_interval: float) -> EventMessage | None:
+        connection = await self._ensure_listener(channel)
+        if has_notifies(connection):
+            message = await self._dequeue_with_notifies(connection, channel, poll_interval)
+        else:
+            message = await self._queue.dequeue(channel, poll_interval)
+        return message
+
+    async def ack(self, event_id: str) -> None:
+        await self._queue.ack(event_id)
+        self._runtime.increment_metric("events.ack")
+
+    async def nack(self, event_id: str) -> None:
+        await self._queue.nack(event_id)
+        self._runtime.increment_metric("events.nack")
+
+    async def shutdown(self) -> None:
+        """Shutdown the listener connection and release resources."""
+        if self._listen_connection_cm is not None:
+            with contextlib.suppress(Exception):
+                await self._listen_connection_cm.__aexit__(None, None, None)
+            self._listen_connection = None
+            self._listen_connection_cm = None
+
+    async def _ensure_listener(self, channel: str) -> Any:
+        """Ensure a dedicated connection is listening on the given channel."""
+        if self._listen_connection is None:
+            validated_channel = normalize_event_channel_name(channel)
+            self._listen_connection_cm = self._config.provide_connection()
+            self._listen_connection = await self._listen_connection_cm.__aenter__()
+            if self._listen_connection is not None:
+                await self._listen_connection.execute(f"LISTEN {validated_channel}")
+        return self._listen_connection
+
+    async def _publish_durable(
+        self, channel: str, event_id: str, payload: "dict[str, Any]", metadata: "dict[str, Any] | None"
+    ) -> None:
+        """Insert event into durable queue and send NOTIFY wakeup signal."""
+        now = datetime.now(timezone.utc)
+        async with self._config.provide_session() as driver:
+            await driver.execute(
+                SQL(
+                    self._queue._upsert_sql,
+                    {
+                        "event_id": event_id,
+                        "channel": channel,
+                        "payload_json": to_json(payload),
+                        "metadata_json": to_json(metadata) if metadata else None,
+                        "status": "pending",
+                        "available_at": now,
+                        "lease_expires_at": None,
+                        "attempts": 0,
+                        "created_at": now,
+                    },
+                    statement_config=self._queue._statement_config,
+                )
+            )
+            await driver.execute(SQL("SELECT pg_notify($1, $2)", channel, to_json({"event_id": event_id})))
+            await driver.commit()
+
+    async def _dequeue_with_notifies(self, connection: Any, channel: str, poll_interval: float) -> EventMessage | None:
+        """Wait for a NOTIFY wakeup then dequeue from the durable table."""
+        try:
+            notify = await asyncio.wait_for(connection.notifies.get(), timeout=poll_interval)
+        except asyncio.TimeoutError:
+            return None
+        notify_payload = notify.payload if is_notification(notify) else None
+        if notify_payload:
+            return await self._queue.dequeue(channel)
+        return None
+
+
+class AsyncpgEventsBackend:
+    """Async backend that relies on PostgreSQL LISTEN/NOTIFY primitives.
+
+    This backend uses asyncpg's native LISTEN/NOTIFY support for real-time
+    event delivery. Messages are ephemeral and not persisted.
+    """
+
+    __slots__ = ("_config", "_listen_connection", "_listen_connection_cm", "_notify_mode", "_runtime")
+
+    supports_sync = False
+    supports_async = True
+    backend_name = "listen_notify"
+
+    def __init__(self, config: "AsyncpgConfig") -> None:
+        if not config.is_async:
+            msg = "AsyncpgEventsBackend requires an async adapter"
+            raise ImproperConfigurationError(msg)
+        self._config = config
+        self._runtime = config.get_observability_runtime()
+        self._listen_connection: Any | None = None
+        self._listen_connection_cm: Any | None = None
+        self._notify_mode: str | None = None
+        log_with_context(
+            logger,
+            logging.DEBUG,
+            "event.listen",
+            adapter_name="asyncpg",
+            backend_name=self.backend_name,
+            mode="async",
+            status="backend_ready",
+        )
+
+    async def publish(self, channel: str, payload: "dict[str, Any]", metadata: "dict[str, Any] | None" = None) -> str:
+        event_id = uuid4().hex
+        async with self._config.provide_session() as driver:
+            await driver.execute(
+                SQL("SELECT pg_notify($1, $2)", channel, encode_notify_payload(event_id, payload, metadata))
+            )
+            await driver.commit()
+        self._runtime.increment_metric("events.publish.native")
+        return event_id
+
+    async def dequeue(self, channel: str, poll_interval: float) -> EventMessage | None:
+        connection = await self._ensure_listener(channel)
+        match self._notify_mode:
+            case "add_listener":
+                return await self._dequeue_with_listener(connection, channel, poll_interval)
+            case "notifies":
+                return await self._dequeue_with_notifies(connection, channel, poll_interval)
+            case _:
+                msg = "PostgreSQL connection does not support LISTEN/NOTIFY callbacks"
+                raise EventChannelError(msg)
+
+    async def ack(self, _event_id: str) -> None:
+        """Acknowledge an event. Native notifications are fire-and-forget."""
+        self._runtime.increment_metric("events.ack")
+
+    async def nack(self, _event_id: str) -> None:
+        """Return an event to the queue (no-op for native LISTEN/NOTIFY)."""
+
+    async def shutdown(self) -> None:
+        """Shutdown the listener connection and release resources."""
+        if self._listen_connection_cm is not None:
+            with contextlib.suppress(Exception):
+                await self._listen_connection_cm.__aexit__(None, None, None)
+            self._listen_connection = None
+            self._listen_connection_cm = None
+            self._notify_mode = None
+
+    async def _ensure_listener(self, channel: str) -> Any:
+        """Ensure a dedicated connection is listening and detect notify mode."""
+        if self._listen_connection is None:
+            validated_channel = normalize_event_channel_name(channel)
+            self._listen_connection_cm = self._config.provide_connection()
+            self._listen_connection = await self._listen_connection_cm.__aenter__()
+            if self._listen_connection is not None and has_add_listener(self._listen_connection):
+                self._notify_mode = "add_listener"
+            elif self._listen_connection is not None and has_notifies(self._listen_connection):
+                self._notify_mode = "notifies"
+                if self._listen_connection is not None:
+                    await self._listen_connection.execute(f"LISTEN {validated_channel}")
+            else:
+                msg = "PostgreSQL connection does not support LISTEN/NOTIFY callbacks"
+                raise EventChannelError(msg)
+        return self._listen_connection
+
+    async def _dequeue_with_listener(self, connection: Any, channel: str, poll_interval: float) -> EventMessage | None:
+        """Wait for notification using add_listener callback API."""
+        loop = asyncio.get_running_loop()
+        future: asyncio.Future[str] = loop.create_future()
+        listener = _AsyncpgNotificationListener(channel, future, loop)
+
+        await connection.add_listener(channel, listener)
+        try:
+            payload_str = await asyncio.wait_for(future, timeout=poll_interval)
+        except asyncio.TimeoutError:
+            return None
+        finally:
+            with contextlib.suppress(Exception):
+                await connection.remove_listener(channel, listener)
+        return decode_notify_payload(channel, payload_str)
+
+    async def _dequeue_with_notifies(self, connection: Any, channel: str, poll_interval: float) -> EventMessage | None:
+        """Wait for notification using notifies queue API."""
+        try:
+            notify = await asyncio.wait_for(connection.notifies.get(), timeout=poll_interval)
+        except asyncio.TimeoutError:
+            return None
+        notify_channel = notify.channel if is_notification(notify) else None
+        if notify_channel != channel:
+            return None
+        return decode_notify_payload(channel, notify.payload)
+
+
+def create_event_backend(
+    config: "AsyncpgConfig", backend_name: str, extension_settings: "dict[str, Any]"
+) -> AsyncpgEventsBackend | AsyncpgHybridEventsBackend | None:
+    """Factory used by EventChannel to create the native backend."""
+    match backend_name:
+        case "listen_notify":
+            try:
+                return AsyncpgEventsBackend(config)
+            except ImproperConfigurationError:
+                return None
+        case "listen_notify_durable":
+            queue_backend = cast(
+                "AsyncTableEventQueue", build_queue_backend(config, extension_settings, adapter_name="asyncpg")
+            )
+            try:
+                return AsyncpgHybridEventsBackend(config, queue_backend)
+            except ImproperConfigurationError:
+                return None
+        case _:
+            return None

sqlspec/adapters/asyncpg/events/store.py

@@ -0,0 +1,40 @@
+"""AsyncPG event queue store for PostgreSQL JSONB storage."""
+
+from sqlspec.adapters.asyncpg.config import AsyncpgConfig
+from sqlspec.extensions.events import BaseEventQueueStore
+
+__all__ = ("AsyncpgEventQueueStore",)
+
+
+class AsyncpgEventQueueStore(BaseEventQueueStore[AsyncpgConfig]):
+    """PostgreSQL event queue store with JSONB columns.
+
+    Uses PostgreSQL-native JSONB for efficient JSON storage and querying.
+    TIMESTAMPTZ ensures proper timezone handling.
+
+    Args:
+        config: AsyncpgConfig with extension_config["events"] settings.
+
+    Notes:
+        Configuration is read from config.extension_config["events"]:
+        - queue_table: Table name (default: "sqlspec_event_queue")
+
+    Example:
+        from sqlspec.adapters.asyncpg import AsyncpgConfig
+        from sqlspec.adapters.asyncpg.events import AsyncpgEventQueueStore
+
+        config = AsyncpgConfig(connection_config={"dsn": "postgresql://..."})
+        store = AsyncpgEventQueueStore(config)
+        for stmt in store.create_statements():
+            await driver.execute_script(stmt)
+    """
+
+    __slots__ = ()
+
+    def _column_types(self) -> "tuple[str, str, str]":
+        """Return PostgreSQL-native column types.
+
+        Returns:
+            Tuple of (payload_type, metadata_type, timestamp_type).
+        """
+        return "JSONB", "JSONB", "TIMESTAMPTZ"

sqlspec/adapters/asyncpg/litestar/__init__.py

@@ -0,0 +1,5 @@
+"""Litestar integration for AsyncPG adapter."""
+
+from sqlspec.adapters.asyncpg.litestar.store import AsyncpgStore
+
+__all__ = ("AsyncpgStore",)

sqlspec/adapters/asyncpg/litestar/store.py

@@ -0,0 +1,251 @@
+"""AsyncPG session store for Litestar integration."""
+
+from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING
+
+from sqlspec.extensions.litestar.store import BaseSQLSpecStore
+
+if TYPE_CHECKING:
+    from sqlspec.adapters.asyncpg.config import AsyncpgConfig
+
+
+__all__ = ("AsyncpgStore",)
+
+
+class AsyncpgStore(BaseSQLSpecStore["AsyncpgConfig"]):
+    """PostgreSQL session store using AsyncPG driver.
+
+    Implements server-side session storage for Litestar using PostgreSQL
+    via the AsyncPG driver. Provides efficient session management with:
+    - Native async PostgreSQL operations
+    - UPSERT support using ON CONFLICT
+    - Automatic expiration handling
+    - Efficient cleanup of expired sessions
+
+    Args:
+        config: AsyncpgConfig instance with extension_config["litestar"] settings.
+
+    Example:
+        from sqlspec.adapters.asyncpg import AsyncpgConfig
+        from sqlspec.adapters.asyncpg.litestar.store import AsyncpgStore
+
+        config = AsyncpgConfig(
+            connection_config={"dsn": "postgresql://..."},
+            extension_config={"litestar": {"session_table": "my_sessions"}}
+        )
+        store = AsyncpgStore(config)
+        await store.create_table()
+    """
+
+    __slots__ = ()
+
+    def __init__(self, config: "AsyncpgConfig") -> None:
+        """Initialize AsyncPG session store.
+
+        Args:
+            config: AsyncpgConfig instance.
+
+        Notes:
+            Table name is read from config.extension_config["litestar"]["session_table"].
+        """
+        super().__init__(config)
+
+    def _get_create_table_sql(self) -> str:
+        """Get PostgreSQL CREATE TABLE SQL with optimized schema.
+
+        Returns:
+            SQL statement to create the sessions table with proper indexes.
+
+        Notes:
+            - Uses TIMESTAMPTZ for timezone-aware expiration timestamps
+            - Partial index WHERE expires_at IS NOT NULL reduces index size/maintenance
+            - FILLFACTOR 80 leaves space for HOT updates, reducing table bloat
+            - Audit columns (created_at, updated_at) help with debugging
+            - Table name is internally controlled, not user input (S608 suppressed)
+        """
+        return f"""
+        CREATE TABLE IF NOT EXISTS {self._table_name} (
+            session_id TEXT PRIMARY KEY,
+            data BYTEA NOT NULL,
+            expires_at TIMESTAMPTZ,
+            created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
+            updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP
+        ) WITH (fillfactor = 80);
+
+        CREATE INDEX IF NOT EXISTS idx_{self._table_name}_expires_at
+        ON {self._table_name}(expires_at) WHERE expires_at IS NOT NULL;
+
+        ALTER TABLE {self._table_name} SET (
+            autovacuum_vacuum_scale_factor = 0.05,
+            autovacuum_analyze_scale_factor = 0.02
+        );
+        """
+
+    def _get_drop_table_sql(self) -> "list[str]":
+        """Get PostgreSQL DROP TABLE SQL statements.
+
+        Returns:
+            List of SQL statements to drop indexes and table.
+        """
+        return [f"DROP INDEX IF EXISTS idx_{self._table_name}_expires_at", f"DROP TABLE IF EXISTS {self._table_name}"]
+
+    async def create_table(self) -> None:
+        """Create the session table if it doesn't exist."""
+        sql = self._get_create_table_sql()
+        async with self._config.provide_session() as driver:
+            await driver.execute_script(sql)
+        self._log_table_created()
+
+    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.
+
+        Notes:
+            Uses CURRENT_TIMESTAMP instead of NOW() for SQL standard compliance.
+            The query planner can use the partial index for expires_at > CURRENT_TIMESTAMP.
+        """
+        sql = f"""
+        SELECT data, expires_at FROM {self._table_name}
+        WHERE session_id = $1
+        AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP)
+        """
+
+        async with self._config.provide_connection() as conn:
+            row = await conn.fetchrow(sql, key)
+
+            if row is None:
+                return None
+
+            if renew_for is not None and row["expires_at"] is not None:
+                new_expires_at = self._calculate_expires_at(renew_for)
+                if new_expires_at is not None:
+                    update_sql = f"""
+                    UPDATE {self._table_name}
+                    SET expires_at = $1, updated_at = CURRENT_TIMESTAMP
+                    WHERE session_id = $2
+                    """
+                    await conn.execute(update_sql, new_expires_at, key)
+
+            return bytes(row["data"])
+
+    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.
+
+        Notes:
+            Uses EXCLUDED to reference the proposed insert values in ON CONFLICT.
+            Updates updated_at timestamp on every write for audit trail.
+        """
+        data = self._value_to_bytes(value)
+        expires_at = self._calculate_expires_at(expires_in)
+
+        sql = f"""
+        INSERT INTO {self._table_name} (session_id, data, expires_at)
+        VALUES ($1, $2, $3)
+        ON CONFLICT (session_id)
+        DO UPDATE SET
+            data = EXCLUDED.data,
+            expires_at = EXCLUDED.expires_at,
+            updated_at = CURRENT_TIMESTAMP
+        """
+
+        async with self._config.provide_connection() as conn:
+            await conn.execute(sql, key, data, expires_at)
+
+    async def delete(self, key: str) -> None:
+        """Delete a session by key.
+
+        Args:
+            key: Session ID to delete.
+        """
+        sql = f"DELETE FROM {self._table_name} WHERE session_id = $1"
+
+        async with self._config.provide_connection() as conn:
+            await conn.execute(sql, key)
+
+    async def delete_all(self) -> None:
+        """Delete all sessions from the store."""
+        sql = f"DELETE FROM {self._table_name}"
+
+        async with self._config.provide_connection() as conn:
+            await conn.execute(sql)
+        self._log_delete_all()
+
+    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.
+
+        Notes:
+            Uses CURRENT_TIMESTAMP for consistency with get() method.
+        """
+        sql = f"""
+        SELECT 1 FROM {self._table_name}
+        WHERE session_id = $1
+        AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP)
+        """
+
+        async with self._config.provide_connection() as conn:
+            result = await conn.fetchval(sql, key)
+            return result is not None
+
+    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.
+        """
+        sql = f"""
+        SELECT expires_at FROM {self._table_name}
+        WHERE session_id = $1
+        """
+
+        async with self._config.provide_connection() as conn:
+            expires_at = await conn.fetchval(sql, key)
+
+            if expires_at is None:
+                return None
+
+            now = datetime.now(timezone.utc)
+            if expires_at <= now:
+                return 0
+
+            delta = expires_at - now
+            return int(delta.total_seconds())
+
+    async def delete_expired(self) -> int:
+        """Delete all expired sessions.
+
+        Returns:
+            Number of sessions deleted.
+
+        Notes:
+            Uses CURRENT_TIMESTAMP for consistency.
+            For very large tables (10M+ rows), consider batching deletes
+            to avoid holding locks too long.
+        """
+        sql = f"DELETE FROM {self._table_name} WHERE expires_at <= CURRENT_TIMESTAMP"
+
+        async with self._config.provide_connection() as conn:
+            result = await conn.execute(sql)
+            count = int(result.split()[-1])
+            if count > 0:
+                self._log_delete_expired(count)
+            return count

sqlspec/adapters/bigquery/__init__.py

@@ -0,0 +1,14 @@
+from sqlspec.adapters.bigquery._typing import BigQueryConnection
+from sqlspec.adapters.bigquery.config import BigQueryConfig, BigQueryConnectionParams
+from sqlspec.adapters.bigquery.core import default_statement_config
+from sqlspec.adapters.bigquery.driver import BigQueryCursor, BigQueryDriver, BigQueryExceptionHandler
+
+__all__ = (
+    "BigQueryConfig",
+    "BigQueryConnection",
+    "BigQueryConnectionParams",
+    "BigQueryCursor",
+    "BigQueryDriver",
+    "BigQueryExceptionHandler",
+    "default_statement_config",
+)

sqlspec/adapters/bigquery/_typing.py

@@ -0,0 +1,86 @@
+"""BigQuery adapter type definitions.
+
+This module contains type aliases and classes that are excluded from mypyc
+compilation to avoid ABI boundary issues.
+"""
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from typing import TypeAlias
+
+    from google.cloud.bigquery import ArrayQueryParameter, Client, ScalarQueryParameter
+
+    from sqlspec.adapters.bigquery.driver import BigQueryDriver
+    from sqlspec.core import StatementConfig
+
+    BigQueryConnection: TypeAlias = Client
+    BigQueryParam: TypeAlias = ArrayQueryParameter | ScalarQueryParameter
+else:
+    try:
+        from google.cloud.bigquery import ArrayQueryParameter, Client, ScalarQueryParameter
+    except Exception:
+        BigQueryConnection = Any
+        BigQueryParam = Any
+    else:
+        BigQueryConnection = Client
+        BigQueryParam = ArrayQueryParameter | ScalarQueryParameter
+
+
+class BigQuerySessionContext:
+    """Sync context manager for BigQuery sessions.
+
+    This class is intentionally excluded from mypyc compilation to avoid ABI
+    boundary issues. It receives callables from uncompiled config classes and
+    instantiates compiled Driver objects, acting as a bridge between compiled
+    and uncompiled code.
+
+    Uses callable-based connection management to decouple from config implementation.
+    """
+
+    __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[[BigQueryDriver], BigQueryDriver]",
+    ) -> 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: BigQueryDriver | None = None
+
+    def __enter__(self) -> "BigQueryDriver":
+        from sqlspec.adapters.bigquery.driver import BigQueryDriver
+
+        self._connection = self._acquire_connection()
+        self._driver = BigQueryDriver(
+            connection=self._connection, statement_config=self._statement_config, driver_features=self._driver_features
+        )
+        return self._prepare_driver(self._driver)
+
+    def __exit__(
+        self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: Any
+    ) -> "bool | None":
+        if self._connection is not None:
+            self._release_connection(self._connection)
+            self._connection = None
+        return None
+
+
+__all__ = ("BigQueryConnection", "BigQueryParam", "BigQuerySessionContext")

sqlspec/adapters/bigquery/adk/__init__.py

@@ -0,0 +1,5 @@
+"""BigQuery ADK store for Google Agent Development Kit session/event storage."""
+
+from sqlspec.adapters.bigquery.adk.store import BigQueryADKMemoryStore, BigQueryADKStore
+
+__all__ = ("BigQueryADKMemoryStore", "BigQueryADKStore")