Source code for sqlspec.adapters.cockroach_asyncpg.config

"""CockroachDB AsyncPG configuration."""

from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypedDict, cast

from asyncpg import Record
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,
    register_pgvector_support,
    resolve_runtime_statement_config,
)
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 asyncio.events import AbstractEventLoop
    from collections.abc import Awaitable, Callable

    from sqlspec.core import StatementConfig
    from sqlspec.observability import ObservabilityConfig

__all__ = (
    "CockroachAsyncpgConfig",
    "CockroachAsyncpgConnectionConfig",
    "CockroachAsyncpgDriverFeatures",
    "CockroachAsyncpgGSSLib",
    "CockroachAsyncpgPoolConfig",
    "CockroachAsyncpgTargetSessionAttrs",
)


CockroachAsyncpgTargetSessionAttrs = Literal["any", "primary", "standby", "read-write", "read-only", "prefer-standby"]
CockroachAsyncpgGSSLib = Literal["gssapi", "sspi"]


class CockroachAsyncpgConnectionConfig(TypedDict):
    """AsyncPG connection parameters for CockroachDB."""

    dsn: NotRequired[str]
    host: NotRequired[str]
    port: NotRequired[int]
    user: NotRequired[str]
    password: NotRequired[str]
    service: NotRequired[str]
    servicefile: NotRequired[str]
    database: NotRequired[str]
    ssl: NotRequired[Any]
    passfile: NotRequired[str]
    direct_tls: NotRequired[bool]
    timeout: NotRequired[float]
    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]"]
    target_session_attrs: NotRequired[CockroachAsyncpgTargetSessionAttrs]
    krbsrvname: NotRequired[str]
    gsslib: NotRequired[CockroachAsyncpgGSSLib]




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]
    connect: NotRequired["Callable[..., Awaitable[CockroachAsyncpgConnection]]"]
    setup: NotRequired["Callable[[CockroachAsyncpgConnection], Awaitable[None]]"]
    init: NotRequired["Callable[[CockroachAsyncpgConnection], Awaitable[None]]"]
    reset: NotRequired["Callable[[CockroachAsyncpgConnection], Awaitable[None]]"]
    loop: NotRequired["AbstractEventLoop"]
    connection_class: NotRequired[type["CockroachAsyncpgConnection"]]
    record_class: NotRequired[type[Record]]
    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[Literal["poll_queue"]] class _CockroachAsyncpgSessionFactory(AsyncPoolSessionFactory): """Uses pool.acquire() context manager pattern instead of direct acquire/release.""" # _connection inherited from AsyncPoolSessionFactory.__slots__ is never written; this class uses _ctx exclusively via the pool.acquire() context manager pattern. __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_migration_schemas: "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 supports_native_row_streaming: "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: raw_enable_pgvector = bool(driver_features and driver_features.get("enable_pgvector") is True) 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["enable_pgvector"] = raw_enable_pgvector driver_features.setdefault("enable_auto_retry", True) # Extract user connection hook before storing driver_features features_dict = dict(driver_features) 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, )