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

sqlspec/adapters/cockroach_asyncpg/core.py

@@ -0,0 +1,55 @@
+"""CockroachDB AsyncPG adapter helpers."""
+
+import secrets
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Final
+
+from sqlspec.utils.type_guards import has_sqlstate
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+__all__ = ("CockroachAsyncpgRetryConfig", "calculate_backoff_seconds", "is_retryable_error")
+
+# Retry configuration defaults (module-level for mypyc compatibility)
+_DEFAULT_MAX_RETRIES: Final[int] = 10
+_DEFAULT_BASE_DELAY_MS: Final[float] = 50.0
+_DEFAULT_MAX_DELAY_MS: Final[float] = 5000.0
+_DEFAULT_ENABLE_LOGGING: Final[bool] = True
+
+
+@dataclass(frozen=True)
+class CockroachAsyncpgRetryConfig:
+    """CockroachDB asyncpg transaction retry configuration."""
+
+    max_retries: int = _DEFAULT_MAX_RETRIES
+    base_delay_ms: float = _DEFAULT_BASE_DELAY_MS
+    max_delay_ms: float = _DEFAULT_MAX_DELAY_MS
+    enable_logging: bool = _DEFAULT_ENABLE_LOGGING
+
+    @classmethod
+    def from_features(cls, driver_features: "Mapping[str, Any]") -> "CockroachAsyncpgRetryConfig":
+        """Build retry config from driver feature mappings."""
+        return cls(
+            max_retries=int(driver_features.get("max_retries", _DEFAULT_MAX_RETRIES)),
+            base_delay_ms=float(driver_features.get("retry_delay_base_ms", _DEFAULT_BASE_DELAY_MS)),
+            max_delay_ms=float(driver_features.get("retry_delay_max_ms", _DEFAULT_MAX_DELAY_MS)),
+            enable_logging=bool(driver_features.get("enable_retry_logging", _DEFAULT_ENABLE_LOGGING)),
+        )
+
+
+def is_retryable_error(error: BaseException) -> bool:
+    """Return True when the error should trigger a CockroachDB retry."""
+    if has_sqlstate(error):
+        return str(error.sqlstate) == "40001"
+    return False
+
+
+def calculate_backoff_seconds(attempt: int, config: "CockroachAsyncpgRetryConfig") -> float:
+    """Calculate exponential backoff delay in seconds."""
+    base: float = config.base_delay_ms * (2**attempt)
+    scale: int = 1000
+    max_jitter: int = max(int(base * scale), 0)
+    jitter: float = secrets.randbelow(max_jitter + 1) / scale if max_jitter else 0.0
+    delay_ms: float = min(base + jitter, config.max_delay_ms)
+    return delay_ms / 1000.0

sqlspec/adapters/cockroach_asyncpg/data_dictionary.py

@@ -0,0 +1,107 @@
+"""CockroachDB AsyncPG data dictionary."""
+
+from typing import TYPE_CHECKING, ClassVar
+
+from mypy_extensions import mypyc_attr
+
+from sqlspec.driver import AsyncDataDictionaryBase
+from sqlspec.typing import ColumnMetadata, ForeignKeyMetadata, IndexMetadata, TableMetadata, VersionInfo
+
+if TYPE_CHECKING:
+    from sqlspec.adapters.cockroach_asyncpg.driver import CockroachAsyncpgDriver
+
+__all__ = ("CockroachAsyncpgDataDictionary",)
+
+
+@mypyc_attr(allow_interpreted_subclasses=True, native_class=False)
+class CockroachAsyncpgDataDictionary(AsyncDataDictionaryBase):
+    """CockroachDB async data dictionary (AsyncPG)."""
+
+    dialect: ClassVar[str] = "cockroachdb"
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    async def get_version(self, driver: "CockroachAsyncpgDriver") -> "VersionInfo | None":
+        driver_id = id(driver)
+        if driver_id in self._version_fetch_attempted:
+            return self._version_cache.get(driver_id)
+
+        version_value = await driver.select_value_or_none(self.get_query("version"))
+        if not version_value:
+            self._log_version_unavailable(type(self).dialect, "missing")
+            self.cache_version(driver_id, None)
+            return None
+
+        version_info = self.parse_version_with_pattern(self.get_dialect_config().version_pattern, str(version_value))
+        if version_info is None:
+            self._log_version_unavailable(type(self).dialect, "parse_failed")
+            self.cache_version(driver_id, None)
+            return None
+
+        self._log_version_detected(type(self).dialect, version_info)
+        self.cache_version(driver_id, version_info)
+        return version_info
+
+    async def get_feature_flag(self, driver: "CockroachAsyncpgDriver", feature: str) -> bool:
+        version_info = await self.get_version(driver)
+        return self.resolve_feature_flag(feature, version_info)
+
+    async def get_optimal_type(self, driver: "CockroachAsyncpgDriver", type_category: str) -> str:
+        config = self.get_dialect_config()
+        return config.get_optimal_type(type_category)
+
+    async def get_tables(self, driver: "CockroachAsyncpgDriver", schema: "str | None" = None) -> "list[TableMetadata]":
+        schema_name = self.resolve_schema(schema)
+        self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="tables")
+        return await driver.select(
+            self.get_query("tables_by_schema"), schema_name=schema_name, schema_type=TableMetadata
+        )
+
+    async def get_columns(
+        self, driver: "CockroachAsyncpgDriver", table: "str | None" = None, schema: "str | None" = None
+    ) -> "list[ColumnMetadata]":
+        schema_name = self.resolve_schema(schema)
+        if table is None:
+            self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="columns")
+            return await driver.select(
+                self.get_query("columns_by_schema"), schema_name=schema_name, schema_type=ColumnMetadata
+            )
+
+        self._log_table_describe(driver, schema_name=schema_name, table_name=table, operation="columns")
+        return await driver.select(
+            self.get_query("columns_by_table"), schema_name=schema_name, table_name=table, schema_type=ColumnMetadata
+        )
+
+    async def get_indexes(
+        self, driver: "CockroachAsyncpgDriver", table: "str | None" = None, schema: "str | None" = None
+    ) -> "list[IndexMetadata]":
+        schema_name = self.resolve_schema(schema)
+        if table is None:
+            self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="indexes")
+            return await driver.select(
+                self.get_query("indexes_by_schema"), schema_name=schema_name, schema_type=IndexMetadata
+            )
+
+        self._log_table_describe(driver, schema_name=schema_name, table_name=table, operation="indexes")
+        return await driver.select(
+            self.get_query("indexes_by_table"), schema_name=schema_name, table_name=table, schema_type=IndexMetadata
+        )
+
+    async def get_foreign_keys(
+        self, driver: "CockroachAsyncpgDriver", table: "str | None" = None, schema: "str | None" = None
+    ) -> "list[ForeignKeyMetadata]":
+        schema_name = self.resolve_schema(schema)
+        if table is None:
+            self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="foreign_keys")
+            return await driver.select(
+                self.get_query("foreign_keys_by_schema"), schema_name=schema_name, schema_type=ForeignKeyMetadata
+            )
+
+        self._log_table_describe(driver, schema_name=schema_name, table_name=table, operation="foreign_keys")
+        return await driver.select(
+            self.get_query("foreign_keys_by_table"),
+            table_name=table,
+            schema_name=schema_name,
+            schema_type=ForeignKeyMetadata,
+        )

sqlspec/adapters/cockroach_asyncpg/driver.py

@@ -0,0 +1,144 @@
+"""CockroachDB AsyncPG driver implementation."""
+
+import asyncio
+import contextlib
+from typing import TYPE_CHECKING, Any, cast
+
+import asyncpg
+
+from sqlspec.adapters.asyncpg.core import create_mapped_exception, driver_profile
+from sqlspec.adapters.asyncpg.driver import AsyncpgDriver
+from sqlspec.adapters.cockroach_asyncpg._typing import CockroachAsyncpgSessionContext
+from sqlspec.adapters.cockroach_asyncpg.core import (
+    CockroachAsyncpgRetryConfig,
+    calculate_backoff_seconds,
+    is_retryable_error,
+)
+from sqlspec.adapters.cockroach_asyncpg.data_dictionary import CockroachAsyncpgDataDictionary
+from sqlspec.core import SQL, register_driver_profile
+from sqlspec.exceptions import SerializationConflictError, TransactionRetryError
+from sqlspec.utils.logging import get_logger
+from sqlspec.utils.type_guards import has_sqlstate
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from sqlspec.adapters.cockroach_asyncpg._typing import CockroachAsyncpgConnection
+    from sqlspec.core import StatementConfig
+    from sqlspec.driver import ExecutionResult
+
+__all__ = ("CockroachAsyncpgDriver", "CockroachAsyncpgExceptionHandler", "CockroachAsyncpgSessionContext")
+
+logger = get_logger("sqlspec.adapters.cockroach_asyncpg")
+
+
+class CockroachAsyncpgExceptionHandler:
+    """Async context manager for CockroachDB AsyncPG exceptions."""
+
+    __slots__ = ("pending_exception",)
+
+    def __init__(self) -> None:
+        self.pending_exception: Exception | None = None
+
+    async def __aenter__(self) -> "CockroachAsyncpgExceptionHandler":
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> bool:
+        if exc_val is None:
+            return False
+        if isinstance(exc_val, asyncpg.PostgresError) or has_sqlstate(exc_val):
+            if has_sqlstate(exc_val) and str(exc_val.sqlstate) == "40001":
+                self.pending_exception = SerializationConflictError(str(exc_val))
+                return True
+            self.pending_exception = create_mapped_exception(exc_val)
+            return True
+        return False
+
+
+class CockroachAsyncpgDriver(AsyncpgDriver):
+    """CockroachDB AsyncPG driver with retry support."""
+
+    __slots__ = ("_enable_retry", "_follower_staleness", "_retry_config")
+    dialect = "postgres"
+
+    def __init__(
+        self,
+        connection: "CockroachAsyncpgConnection",
+        statement_config: "StatementConfig | None" = None,
+        driver_features: "dict[str, Any] | None" = None,
+    ) -> None:
+        super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features)
+        self._retry_config = CockroachAsyncpgRetryConfig.from_features(self.driver_features)
+        self._enable_retry = bool(self.driver_features.get("enable_auto_retry", True))
+        self._follower_staleness = cast("str | None", self.driver_features.get("default_staleness"))
+        # Data dictionary is lazily initialized in property; use parent slot
+        self._data_dictionary = None
+
+    async def _execute_with_retry(self, operation: "Callable[[], Any]") -> "ExecutionResult":
+        if not self._enable_retry:
+            return cast("ExecutionResult", await operation())
+
+        last_error: Exception | None = None
+
+        async def attempt_operation() -> "tuple[ExecutionResult | None, Exception | None]":
+            try:
+                return await operation(), None
+            except Exception as exc:
+                return None, exc
+
+        for attempt in range(self._retry_config.max_retries + 1):
+            result, exc = await attempt_operation()
+            if exc is None:
+                return cast("ExecutionResult", result)
+            last_error = exc
+            if not is_retryable_error(exc) or attempt >= self._retry_config.max_retries:
+                raise exc
+            with contextlib.suppress(Exception):
+                await self.connection.execute("ROLLBACK")
+            delay = calculate_backoff_seconds(attempt, self._retry_config)
+            if self._retry_config.enable_logging:
+                logger.debug("CockroachDB retry %s/%s after %.3fs", attempt + 1, self._retry_config.max_retries, delay)
+            await asyncio.sleep(delay)
+
+        msg = "CockroachDB transaction retry limit exceeded"
+        raise TransactionRetryError(msg) from last_error
+
+    async def _apply_follower_reads(self, cursor: "CockroachAsyncpgConnection") -> None:
+        if not self.driver_features.get("enable_follower_reads", False):
+            return
+        if not self._follower_staleness:
+            return
+        await cursor.execute(f"SET TRANSACTION AS OF SYSTEM TIME {self._follower_staleness}")
+
+    async def dispatch_execute(self, cursor: "CockroachAsyncpgConnection", statement: SQL) -> "ExecutionResult":
+        async def operation() -> "ExecutionResult":
+            if statement.returns_rows():
+                await self._apply_follower_reads(cursor)
+            return await super(CockroachAsyncpgDriver, self).dispatch_execute(cursor, statement)
+
+        return await self._execute_with_retry(operation)
+
+    async def dispatch_execute_many(self, cursor: "CockroachAsyncpgConnection", statement: SQL) -> "ExecutionResult":
+        async def operation() -> "ExecutionResult":
+            return await super(CockroachAsyncpgDriver, self).dispatch_execute_many(cursor, statement)
+
+        return await self._execute_with_retry(operation)
+
+    async def dispatch_execute_script(self, cursor: "CockroachAsyncpgConnection", statement: SQL) -> "ExecutionResult":
+        async def operation() -> "ExecutionResult":
+            return await super(CockroachAsyncpgDriver, self).dispatch_execute_script(cursor, statement)
+
+        return await self._execute_with_retry(operation)
+
+    def handle_database_exceptions(self) -> "CockroachAsyncpgExceptionHandler":  # type: ignore[override]
+        return CockroachAsyncpgExceptionHandler()
+
+    @property
+    def data_dictionary(self) -> "CockroachAsyncpgDataDictionary":  # type: ignore[override]
+        if self._data_dictionary is None:
+            # Intentionally assign CockroachDB-specific data dictionary to parent slot
+            object.__setattr__(self, "_data_dictionary", CockroachAsyncpgDataDictionary())
+        return cast("CockroachAsyncpgDataDictionary", self._data_dictionary)
+
+
+register_driver_profile("cockroach_asyncpg", driver_profile)

sqlspec/adapters/cockroach_asyncpg/events/__init__.py

@@ -0,0 +1,3 @@
+from sqlspec.adapters.cockroach_asyncpg.events.store import CockroachAsyncpgEventQueueStore
+
+__all__ = ("CockroachAsyncpgEventQueueStore",)

sqlspec/adapters/cockroach_asyncpg/events/store.py

@@ -0,0 +1,20 @@
+"""CockroachDB event queue store for asyncpg driver."""
+
+from sqlspec.adapters.cockroach_asyncpg.config import CockroachAsyncpgConfig
+from sqlspec.extensions.events import BaseEventQueueStore
+
+__all__ = ("CockroachAsyncpgEventQueueStore",)
+
+
+class CockroachAsyncpgEventQueueStore(BaseEventQueueStore[CockroachAsyncpgConfig]):
+    """Queue DDL for CockroachDB asyncpg configs.
+
+    CockroachDB uses JSONB for efficient JSON storage and TIMESTAMPTZ for
+    timezone-aware timestamps.
+    """
+
+    __slots__ = ()
+
+    def _column_types(self) -> "tuple[str, str, str]":
+        """Return CockroachDB-optimized column types for the event queue."""
+        return "JSONB", "JSONB", "TIMESTAMPTZ"

sqlspec/adapters/cockroach_asyncpg/litestar/__init__.py

@@ -0,0 +1,3 @@
+from sqlspec.adapters.cockroach_asyncpg.litestar.store import CockroachAsyncpgStore
+
+__all__ = ("CockroachAsyncpgStore",)

sqlspec/adapters/cockroach_asyncpg/litestar/store.py

@@ -0,0 +1,142 @@
+"""CockroachDB session store for Litestar integration using asyncpg."""
+
+from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING
+
+from sqlspec.extensions.litestar.store import BaseSQLSpecStore
+
+if TYPE_CHECKING:
+    from sqlspec.adapters.cockroach_asyncpg.config import CockroachAsyncpgConfig
+
+
+__all__ = ("CockroachAsyncpgStore",)
+
+
+class CockroachAsyncpgStore(BaseSQLSpecStore["CockroachAsyncpgConfig"]):
+    """CockroachDB session store using asyncpg driver."""
+
+    __slots__ = ()
+
+    def __init__(self, config: "CockroachAsyncpgConfig") -> None:
+        super().__init__(config)
+
+    def _get_create_table_sql(self) -> str:
+        """Get CockroachDB CREATE TABLE SQL with optimized schema."""
+        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
+        );
+
+        CREATE INDEX IF NOT EXISTS idx_{self._table_name}_expires_at
+        ON {self._table_name}(expires_at) WHERE expires_at IS NOT NULL;
+        """
+
+    def _get_drop_table_sql(self) -> "list[str]":
+        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:
+        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":
+        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:
+        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:
+        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:
+        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:
+        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:
+            row = await conn.fetchrow(sql, key)
+            return row is not None
+
+    async def expires_in(self, key: str) -> "int | None":
+        sql = f"""
+        SELECT expires_at FROM {self._table_name}
+        WHERE session_id = $1
+        """
+
+        async with self._config.provide_connection() as conn:
+            row = await conn.fetchrow(sql, key)
+
+            if row is None or row["expires_at"] is None:
+                return None
+
+            expires_at = row["expires_at"]
+            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:
+        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 result else 0
+            if count > 0:
+                self._log_delete_expired(count)
+            return count

sqlspec/adapters/cockroach_psycopg/__init__.py

@@ -0,0 +1,38 @@
+from sqlspec.adapters.cockroach_psycopg._typing import (
+    CockroachAsyncConnection,
+    CockroachPsycopgAsyncSessionContext,
+    CockroachPsycopgSyncSessionContext,
+    CockroachSyncConnection,
+)
+from sqlspec.adapters.cockroach_psycopg.config import (
+    CockroachPsycopgAsyncConfig,
+    CockroachPsycopgConnectionConfig,
+    CockroachPsycopgDriverFeatures,
+    CockroachPsycopgPoolConfig,
+    CockroachPsycopgSyncConfig,
+)
+from sqlspec.adapters.cockroach_psycopg.core import build_statement_config, driver_profile
+from sqlspec.adapters.cockroach_psycopg.driver import (
+    CockroachPsycopgAsyncDriver,
+    CockroachPsycopgAsyncExceptionHandler,
+    CockroachPsycopgSyncDriver,
+    CockroachPsycopgSyncExceptionHandler,
+)
+
+__all__ = (
+    "CockroachAsyncConnection",
+    "CockroachPsycopgAsyncConfig",
+    "CockroachPsycopgAsyncDriver",
+    "CockroachPsycopgAsyncExceptionHandler",
+    "CockroachPsycopgAsyncSessionContext",
+    "CockroachPsycopgConnectionConfig",
+    "CockroachPsycopgDriverFeatures",
+    "CockroachPsycopgPoolConfig",
+    "CockroachPsycopgSyncConfig",
+    "CockroachPsycopgSyncDriver",
+    "CockroachPsycopgSyncExceptionHandler",
+    "CockroachPsycopgSyncSessionContext",
+    "CockroachSyncConnection",
+    "build_statement_config",
+    "driver_profile",
+)

sqlspec/adapters/cockroach_psycopg/_typing.py

@@ -0,0 +1,129 @@
+"""CockroachDB psycopg 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
+
+from psycopg import crdb as psycopg_crdb
+from psycopg.rows import DictRow as PsycopgDictRow
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from typing import TypeAlias
+
+    from psycopg.crdb import AsyncCrdbConnection, CrdbConnection
+
+    from sqlspec.adapters.cockroach_psycopg.driver import CockroachPsycopgAsyncDriver, CockroachPsycopgSyncDriver
+    from sqlspec.core import StatementConfig
+
+    # Parametrize with DictRow so type system knows rows are dict-like
+    CockroachSyncConnection: TypeAlias = CrdbConnection[PsycopgDictRow]
+    CockroachAsyncConnection: TypeAlias = AsyncCrdbConnection[PsycopgDictRow]
+else:
+    CockroachSyncConnection = psycopg_crdb.CrdbConnection
+    CockroachAsyncConnection = psycopg_crdb.AsyncCrdbConnection
+
+
+class CockroachPsycopgSyncSessionContext:
+    """Sync context manager for CockroachDB psycopg 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[[CockroachPsycopgSyncDriver], CockroachPsycopgSyncDriver]",
+    ) -> 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: CockroachPsycopgSyncDriver | None = None
+
+    def __enter__(self) -> "CockroachPsycopgSyncDriver":
+        from sqlspec.adapters.cockroach_psycopg.driver import CockroachPsycopgSyncDriver
+
+        self._connection = self._acquire_connection()
+        self._driver = CockroachPsycopgSyncDriver(
+            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
+
+
+class CockroachPsycopgAsyncSessionContext:
+    """Async context manager for CockroachDB psycopg 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[[CockroachPsycopgAsyncDriver], CockroachPsycopgAsyncDriver]",
+    ) -> 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: CockroachPsycopgAsyncDriver | None = None
+
+    async def __aenter__(self) -> "CockroachPsycopgAsyncDriver":
+        from sqlspec.adapters.cockroach_psycopg.driver import CockroachPsycopgAsyncDriver
+
+        self._connection = await self._acquire_connection()
+        self._driver = CockroachPsycopgAsyncDriver(
+            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__ = (
+    "CockroachAsyncConnection",
+    "CockroachPsycopgAsyncSessionContext",
+    "CockroachPsycopgSyncSessionContext",
+    "CockroachSyncConnection",
+    "PsycopgDictRow",
+)

sqlspec/adapters/cockroach_psycopg/adk/__init__.py

@@ -0,0 +1,13 @@
+from sqlspec.adapters.cockroach_psycopg.adk.store import (
+    CockroachPsycopgAsyncADKMemoryStore,
+    CockroachPsycopgAsyncADKStore,
+    CockroachPsycopgSyncADKMemoryStore,
+    CockroachPsycopgSyncADKStore,
+)
+
+__all__ = (
+    "CockroachPsycopgAsyncADKMemoryStore",
+    "CockroachPsycopgAsyncADKStore",
+    "CockroachPsycopgSyncADKMemoryStore",
+    "CockroachPsycopgSyncADKStore",
+)