sqlspec 0.47.0__cp314-cp314-win_amd64.whl

sqlspec/adapters/cockroach_asyncpg/config.py

@@ -0,0 +1,248 @@
+"""CockroachDB AsyncPG configuration."""
+
+from typing import TYPE_CHECKING, Any, ClassVar, TypedDict, cast
+
+from asyncpg import create_pool as asyncpg_create_pool
+from typing_extensions import NotRequired
+
+from sqlspec.adapters.asyncpg.core import (
+    apply_driver_features,
+    build_connection_config,
+    default_statement_config,
+    register_json_codecs,
+)
+from sqlspec.adapters.cockroach_asyncpg._typing import (
+    CockroachAsyncpgConnection,
+    CockroachAsyncpgPool,
+    CockroachAsyncpgSessionContext,
+)
+from sqlspec.adapters.cockroach_asyncpg.driver import CockroachAsyncpgDriver, CockroachAsyncpgExceptionHandler
+from sqlspec.config import AsyncDatabaseConfig, ExtensionConfigs
+from sqlspec.driver._async import AsyncPoolConnectionContext, AsyncPoolSessionFactory
+from sqlspec.extensions.events import EventRuntimeHints
+from sqlspec.utils.config_tools import normalize_connection_config
+from sqlspec.utils.serializers import from_json, to_json
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
+    from sqlspec.core import StatementConfig
+    from sqlspec.observability import ObservabilityConfig
+
+__all__ = (
+    "CockroachAsyncpgConfig",
+    "CockroachAsyncpgConnectionConfig",
+    "CockroachAsyncpgDriverFeatures",
+    "CockroachAsyncpgPoolConfig",
+)
+
+
+class CockroachAsyncpgConnectionConfig(TypedDict):
+    """AsyncPG connection parameters for CockroachDB."""
+
+    dsn: NotRequired[str]
+    host: NotRequired[str]
+    port: NotRequired[int]
+    user: NotRequired[str]
+    password: NotRequired[str]
+    database: NotRequired[str]
+    ssl: NotRequired[Any]
+    passfile: NotRequired[str]
+    direct_tls: NotRequired[bool]
+    connect_timeout: NotRequired[float]
+    command_timeout: NotRequired[float]
+    statement_cache_size: NotRequired[int]
+    max_cached_statement_lifetime: NotRequired[int]
+    max_cacheable_statement_size: NotRequired[int]
+    server_settings: NotRequired["dict[str, str]"]
+
+
+class CockroachAsyncpgPoolConfig(CockroachAsyncpgConnectionConfig):
+    """AsyncPG pool parameters for CockroachDB."""
+
+    min_size: NotRequired[int]
+    max_size: NotRequired[int]
+    max_queries: NotRequired[int]
+    max_inactive_connection_lifetime: NotRequired[float]
+    setup: NotRequired["Callable[[CockroachAsyncpgConnection], Awaitable[None]]"]
+    init: NotRequired["Callable[[CockroachAsyncpgConnection], Awaitable[None]]"]
+    extra: NotRequired["dict[str, Any]"]
+
+
+class CockroachAsyncpgDriverFeatures(TypedDict):
+    """Driver feature flags for CockroachDB AsyncPG adapter.
+
+    on_connection_create: Async callback executed when a connection is acquired from pool.
+        Receives the raw asyncpg connection for low-level driver configuration.
+        Called after internal setup (JSON codecs, pgvector registration).
+    """
+
+    enable_auto_retry: NotRequired[bool]
+    max_retries: NotRequired[int]
+    retry_delay_base_ms: NotRequired[float]
+    retry_delay_max_ms: NotRequired[float]
+    enable_retry_logging: NotRequired[bool]
+    enable_follower_reads: NotRequired[bool]
+    default_staleness: NotRequired[str]
+    json_serializer: NotRequired["Callable[[Any], str]"]
+    json_deserializer: NotRequired["Callable[[str], Any]"]
+    enable_json_codecs: NotRequired[bool]
+    enable_pgvector: NotRequired[bool]
+    on_connection_create: "NotRequired[Callable[[CockroachAsyncpgConnection], Awaitable[None]]]"
+    enable_events: NotRequired[bool]
+    events_backend: NotRequired[str]
+
+
+class _CockroachAsyncpgSessionFactory(AsyncPoolSessionFactory):
+    """Uses pool.acquire() context manager pattern instead of direct acquire/release."""
+
+    __slots__ = ("_ctx",)
+
+    def __init__(self, config: "CockroachAsyncpgConfig") -> None:
+        super().__init__(config)
+        self._ctx: Any | None = None
+
+    async def acquire_connection(self) -> "CockroachAsyncpgConnection":
+        pool = self._config.connection_instance
+        if pool is None:
+            pool = await self._config.create_pool()
+            self._config.connection_instance = pool
+        ctx = pool.acquire()
+        self._ctx = ctx
+        return cast("CockroachAsyncpgConnection", await ctx.__aenter__())
+
+    async def release_connection(self, _conn: "CockroachAsyncpgConnection", **kwargs: Any) -> None:
+        if self._ctx is not None:
+            await self._ctx.__aexit__(None, None, None)
+            self._ctx = None
+
+
+class CockroachAsyncpgConnectionContext(AsyncPoolConnectionContext):
+    """Async context manager for CockroachDB AsyncPG connections."""
+
+    __slots__ = ()
+
+
+class CockroachAsyncpgConfig(
+    AsyncDatabaseConfig[CockroachAsyncpgConnection, CockroachAsyncpgPool, CockroachAsyncpgDriver]
+):
+    """Configuration for CockroachDB using AsyncPG."""
+
+    driver_type: "ClassVar[type[CockroachAsyncpgDriver]]" = CockroachAsyncpgDriver
+    connection_type: "ClassVar[type[CockroachAsyncpgConnection]]" = CockroachAsyncpgConnection  # type: ignore[assignment]
+    supports_transactional_ddl: "ClassVar[bool]" = True
+    supports_native_arrow_export: "ClassVar[bool]" = True
+    supports_native_arrow_import: "ClassVar[bool]" = True
+    supports_native_parquet_export: "ClassVar[bool]" = True
+    supports_native_parquet_import: "ClassVar[bool]" = True
+    _connection_context_class: "ClassVar[type[CockroachAsyncpgConnectionContext]]" = CockroachAsyncpgConnectionContext
+    _session_factory_class: "ClassVar[type[_CockroachAsyncpgSessionFactory]]" = _CockroachAsyncpgSessionFactory
+    _session_context_class: "ClassVar[type[CockroachAsyncpgSessionContext]]" = CockroachAsyncpgSessionContext
+    _default_statement_config = default_statement_config
+
+    def __init__(
+        self,
+        *,
+        connection_config: "CockroachAsyncpgPoolConfig | dict[str, Any] | None" = None,
+        connection_instance: "CockroachAsyncpgPool | None" = None,
+        migration_config: "dict[str, Any] | None" = None,
+        statement_config: "StatementConfig | None" = None,
+        driver_features: "CockroachAsyncpgDriverFeatures | dict[str, Any] | None" = None,
+        bind_key: "str | None" = None,
+        extension_config: "ExtensionConfigs | None" = None,
+        observability_config: "ObservabilityConfig | None" = None,
+        **kwargs: Any,
+    ) -> None:
+        connection_config = normalize_connection_config(connection_config)
+        statement_config = statement_config or default_statement_config
+        statement_config, driver_features = apply_driver_features(statement_config, driver_features)
+
+        driver_features.setdefault("enable_auto_retry", True)
+
+        # Extract user connection hook before storing driver_features
+        features_dict = dict(driver_features) if driver_features else {}
+        self._user_connection_hook: Callable[[CockroachAsyncpgConnection], Awaitable[None]] | None = features_dict.pop(
+            "on_connection_create", None
+        )
+
+        super().__init__(
+            connection_config=connection_config,
+            connection_instance=connection_instance,
+            migration_config=migration_config,
+            statement_config=statement_config,
+            driver_features=features_dict,
+            bind_key=bind_key,
+            extension_config=extension_config,
+            observability_config=observability_config,
+            **kwargs,
+        )
+
+    async def _create_pool(self) -> "CockroachAsyncpgPool":
+        config = build_connection_config(self.connection_config)
+        config.setdefault("init", self._init_connection)
+        return await asyncpg_create_pool(**config)
+
+    async def _init_connection(self, connection: "CockroachAsyncpgConnection") -> None:
+        """Initialize connection with JSON codecs and user callback."""
+        if self.driver_features.get("enable_json_codecs", True):
+            await register_json_codecs(
+                connection,
+                encoder=self.driver_features.get("json_serializer", to_json),
+                decoder=self.driver_features.get("json_deserializer", from_json),
+            )
+
+        if self._user_connection_hook is not None:
+            await self._user_connection_hook(connection)
+
+    async def _close_pool(self) -> None:
+        if not self.connection_instance:
+            return
+        await self.connection_instance.close()
+        self.connection_instance = None
+
+    async def create_connection(self) -> "CockroachAsyncpgConnection":
+        if self.connection_instance is None:
+            self.connection_instance = await self.create_pool()
+        return cast("CockroachAsyncpgConnection", await self.connection_instance.acquire())
+
+    def provide_session(
+        self,
+        *_args: Any,
+        statement_config: "StatementConfig | None" = None,
+        follower_reads: bool | None = None,
+        staleness: str | None = None,
+        **_kwargs: Any,
+    ) -> "CockroachAsyncpgSessionContext":
+        factory = _CockroachAsyncpgSessionFactory(self)
+        driver_features = dict(self.driver_features)
+        if follower_reads is not None:
+            driver_features["enable_follower_reads"] = follower_reads
+        if staleness is not None:
+            driver_features["default_staleness"] = staleness
+
+        return CockroachAsyncpgSessionContext(
+            acquire_connection=factory.acquire_connection,
+            release_connection=factory.release_connection,
+            statement_config=statement_config or self.statement_config or default_statement_config,
+            driver_features=driver_features,
+            prepare_driver=self._prepare_driver,
+        )
+
+    async def provide_pool(self, *args: Any, **kwargs: Any) -> "CockroachAsyncpgPool":
+        if not self.connection_instance:
+            self.connection_instance = await self.create_pool()
+        return self.connection_instance
+
+    def get_signature_namespace(self) -> "dict[str, Any]":
+        namespace = super().get_signature_namespace()
+        namespace.update({
+            "CockroachAsyncpgConnectionConfig": CockroachAsyncpgConnectionConfig,
+            "CockroachAsyncpgPoolConfig": CockroachAsyncpgPoolConfig,
+            "CockroachAsyncpgDriver": CockroachAsyncpgDriver,
+            "CockroachAsyncpgExceptionHandler": CockroachAsyncpgExceptionHandler,
+            "CockroachAsyncpgSessionContext": CockroachAsyncpgSessionContext,
+        })
+        return namespace
+
+    def get_event_runtime_hints(self) -> "EventRuntimeHints":
+        return EventRuntimeHints(poll_interval=0.5, select_for_update=True, skip_locked=True)

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,110 @@
+"""CockroachDB AsyncPG data dictionary."""
+
+from typing import TYPE_CHECKING, ClassVar
+
+from mypy_extensions import mypyc_attr
+
+from sqlspec.data_dictionary.dialects.cockroachdb import resolve_cockroachdb_json_type
+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()
+        if type_category == "json":
+            return resolve_cockroachdb_json_type(await self.get_version(driver))
+        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,142 @@
+"""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.driver import BaseAsyncExceptionHandler
+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(BaseAsyncExceptionHandler):
+    """Async context manager for CockroachDB AsyncPG exceptions."""
+
+    __slots__ = ()
+
+    def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "BaseException") -> bool:
+        _ = exc_type
+        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]", *args: Any) -> "ExecutionResult":
+        if not self._enable_retry:
+            return cast("ExecutionResult", await operation(*args))
+
+        last_error: Exception | None = None
+
+        for attempt in range(self._retry_config.max_retries + 1):
+            try:
+                return cast("ExecutionResult", await operation(*args))
+            except Exception as exc:
+                last_error = exc
+                if not is_retryable_error(exc) or attempt >= self._retry_config.max_retries:
+                    raise
+            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_impl(self, cursor: "CockroachAsyncpgConnection", statement: SQL) -> "ExecutionResult":
+        if statement.returns_rows():
+            await self._apply_follower_reads(cursor)
+        return await super().dispatch_execute(cursor, statement)
+
+    async def _dispatch_execute_many_impl(
+        self, cursor: "CockroachAsyncpgConnection", statement: SQL
+    ) -> "ExecutionResult":
+        return await super().dispatch_execute_many(cursor, statement)
+
+    async def _dispatch_execute_script_impl(
+        self, cursor: "CockroachAsyncpgConnection", statement: SQL
+    ) -> "ExecutionResult":
+        return await super().dispatch_execute_script(cursor, statement)
+
+    async def dispatch_execute(self, cursor: Any, statement: SQL) -> "ExecutionResult":
+        if not self._enable_retry:
+            return await self._dispatch_execute_impl(cursor, statement)
+        return await self._execute_with_retry(self._dispatch_execute_impl, cursor, statement)
+
+    async def dispatch_execute_many(self, cursor: Any, statement: SQL) -> "ExecutionResult":
+        if not self._enable_retry:
+            return await super().dispatch_execute_many(cursor, statement)
+        return await self._execute_with_retry(self._dispatch_execute_many_impl, cursor, statement)
+
+    async def dispatch_execute_script(self, cursor: Any, statement: SQL) -> "ExecutionResult":
+        if not self._enable_retry:
+            return await super().dispatch_execute_script(cursor, statement)
+        return await self._execute_with_retry(self._dispatch_execute_script_impl, cursor, statement)
+
+    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",)