Source code for sqlspec.adapters.asyncpg.config

"""AsyncPG database configuration with direct field-based configuration."""

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

from asyncpg import Connection, Record
from asyncpg import create_pool as asyncpg_create_pool
from asyncpg.connection import ConnectionMeta
from asyncpg.pool import Pool, PoolConnectionProxy, PoolConnectionProxyMeta
from mypy_extensions import mypyc_attr
from typing_extensions import NotRequired

from sqlspec.adapters.asyncpg._typing import (
    AsyncpgConnection,
    AsyncpgCursor,
    AsyncpgPool,
    AsyncpgPreparedStatement,
    AsyncpgSessionContext,
)
from sqlspec.adapters.asyncpg.core import (
    apply_driver_features,
    build_connection_config,
    build_postgres_extension_probe_names,
    default_statement_config,
    register_json_codecs,
    register_pgvector_support,
    resolve_postgres_extension_state,
    resolve_runtime_statement_config,
)
from sqlspec.adapters.asyncpg.driver import AsyncpgDriver, AsyncpgExceptionHandler
from sqlspec.config import AsyncDatabaseConfig, ExtensionConfigs
from sqlspec.driver._async import AsyncPoolConnectionContext, AsyncPoolSessionFactory
from sqlspec.exceptions import ImproperConfigurationError, MissingDependencyError
from sqlspec.extensions.events import EventRuntimeHints
from sqlspec.typing import ALLOYDB_CONNECTOR_INSTALLED, CLOUD_SQL_CONNECTOR_INSTALLED, PGVECTOR_INSTALLED
from sqlspec.utils.config_tools import normalize_connection_config
from sqlspec.utils.logging import get_logger
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__ = (
    "PGVECTOR_INSTALLED",
    "AsyncpgConfig",
    "AsyncpgConnectionConfig",
    "AsyncpgDriverFeatures",
    "AsyncpgGSSLib",
    "AsyncpgPoolConfig",
    "AsyncpgTargetSessionAttrs",
    "register_json_codecs",
    "register_pgvector_support",
)


logger = get_logger(__name__)


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



[docs]
class AsyncpgConnectionConfig(TypedDict):
    """TypedDict for AsyncPG connection parameters."""

    dsn: NotRequired[str]
    host: NotRequired[str]
    port: NotRequired[int]
    user: NotRequired[str]
    password: NotRequired[str]
    database: NotRequired[str]
    ssl: NotRequired[Any]
    passfile: NotRequired[str]
    service: NotRequired[str]
    servicefile: 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[AsyncpgTargetSessionAttrs]
    krbsrvname: NotRequired[str]
    gsslib: NotRequired[AsyncpgGSSLib]

[docs]

class AsyncpgPoolConfig(AsyncpgConnectionConfig): """TypedDict for AsyncPG pool parameters, inheriting connection parameters.""" min_size: NotRequired[int] max_size: NotRequired[int] max_queries: NotRequired[int] max_inactive_connection_lifetime: NotRequired[float] connect: NotRequired["Callable[..., Awaitable[AsyncpgConnection]]"] setup: NotRequired["Callable[[AsyncpgConnection], Awaitable[None]]"] init: NotRequired["Callable[[AsyncpgConnection], Awaitable[None]]"] reset: NotRequired["Callable[[AsyncpgConnection], Awaitable[None]]"] loop: NotRequired["AbstractEventLoop"] connection_class: NotRequired[type["AsyncpgConnection"]] record_class: NotRequired[type[Record]] extra: NotRequired["dict[str, Any]"]

class AsyncpgDriverFeatures(TypedDict): """AsyncPG driver feature flags. json_serializer: Custom JSON serializer function for PostgreSQL JSON/JSONB types. Defaults to sqlspec.utils.serializers.to_json. Use for performance optimization or custom encoding behavior. Applied when enable_json_codecs is True. json_deserializer: Custom JSON deserializer function for PostgreSQL JSON/JSONB types. Defaults to sqlspec.utils.serializers.from_json. Use for performance optimization or custom decoding behavior. Applied when enable_json_codecs is True. enable_json_codecs: Enable automatic JSON/JSONB codec registration on connections. Defaults to True for seamless Python dict/list to PostgreSQL JSON/JSONB conversion. Set to False to disable automatic codec registration (manual handling required). enable_pgvector: Enable pgvector extension support for vector similarity search. Requires pgvector-python package (pip install pgvector) and PostgreSQL with pgvector extension. Defaults to True when pgvector-python is installed. Provides automatic conversion between Python objects and PostgreSQL vector types. Enables vector similarity operations and index support. enable_paradedb: Enable ParadeDB (pg_search) extension detection. When enabled and the pg_search extension is detected, the SQL dialect switches to "paradedb" which supports search operators (@@@, &&&, etc.) and inherits all pgvector distance operators. Defaults to True. Independent of enable_pgvector. enable_cloud_sql: Enable Google Cloud SQL connector integration. Requires cloud-sql-python-connector package. Defaults to False (explicit opt-in required). Auto-configures IAM authentication, SSL, and IP routing. Mutually exclusive with enable_alloydb. cloud_sql_instance: Cloud SQL instance connection name. Format: "project:region:instance" Required when enable_cloud_sql is True. cloud_sql_enable_iam_auth: Enable IAM database authentication. Defaults to False for passwordless authentication. When False, requires user/password in connection_config. cloud_sql_ip_type: IP address type for connection. Options: "PUBLIC", "PRIVATE", "PSC" Defaults to "PRIVATE". enable_alloydb: Enable Google AlloyDB connector integration. Requires cloud-alloydb-python-connector package. Defaults to False (explicit opt-in required). Auto-configures IAM authentication and private networking. Mutually exclusive with enable_cloud_sql. alloydb_instance_uri: AlloyDB instance URI. Format: "projects/PROJECT/locations/REGION/clusters/CLUSTER/instances/INSTANCE" Required when enable_alloydb is True. enable_alloydb_iam_auth: Enable IAM database authentication. Defaults to False for passwordless authentication. alloydb_ip_type: IP address type for connection. Options: "PUBLIC", "PRIVATE", "PSC" Defaults to "PRIVATE". enable_events: Enable database event channel support. Defaults to True when extension_config["events"] is configured. Provides pub/sub capabilities via LISTEN/NOTIFY or table-backed fallback. Requires extension_config["events"] for migration setup when using poll_queue or notify_queue. events_backend: Event channel backend selection. Options: "notify", "notify_queue", "poll_queue" - "notify": Transient PostgreSQL LISTEN/NOTIFY with no replay or retry - "notify_queue": Durable queue plus a PostgreSQL notification wakeup hint - "poll_queue": Durable queue discovered by polling Defaults to "notify". """ json_serializer: NotRequired["Callable[[Any], str]"] json_deserializer: NotRequired["Callable[[str], Any]"] enable_json_codecs: NotRequired[bool] enable_pgvector: NotRequired[bool] enable_paradedb: NotRequired[bool] enable_cloud_sql: NotRequired[bool] cloud_sql_instance: NotRequired[str] cloud_sql_enable_iam_auth: NotRequired[bool] cloud_sql_ip_type: NotRequired[str] enable_alloydb: NotRequired[bool] alloydb_instance_uri: NotRequired[str] enable_alloydb_iam_auth: NotRequired[bool] alloydb_ip_type: NotRequired[str] enable_events: NotRequired[bool] events_backend: NotRequired[Literal["notify", "notify_queue", "poll_queue"]] connection_instance: NotRequired["AsyncpgPool"] on_connection_create: NotRequired["Callable[[AsyncpgConnection], Awaitable[None]]"] class _AsyncpgCloudSqlConnector: __slots__ = ("_config", "_database", "_password", "_user") def __init__(self, config: "AsyncpgConfig", user: str | None, password: str | None, database: str | None) -> None: self._config = config self._user = user self._password = password self._database = database async def __call__(self) -> "AsyncpgConnection": connector = self._config.get_cloud_sql_connector() if connector is None: msg = "Cloud SQL connector is not initialized" raise ImproperConfigurationError(msg) conn_kwargs: dict[str, Any] = { "instance_connection_string": self._config.driver_features["cloud_sql_instance"], "driver": "asyncpg", "enable_iam_auth": self._config.driver_features.get("cloud_sql_enable_iam_auth", False), "ip_type": self._config.driver_features.get("cloud_sql_ip_type", "PRIVATE"), } if self._user: conn_kwargs["user"] = self._user if self._password: conn_kwargs["password"] = self._password if self._database: conn_kwargs["db"] = self._database return cast("AsyncpgConnection", await connector.connect_async(**conn_kwargs)) class _AsyncpgAlloydbConnector: __slots__ = ("_config", "_database", "_password", "_user") def __init__(self, config: "AsyncpgConfig", user: str | None, password: str | None, database: str | None) -> None: self._config = config self._user = user self._password = password self._database = database async def __call__(self) -> "AsyncpgConnection": connector = self._config.get_alloydb_connector() if connector is None: msg = "AlloyDB connector is not initialized" raise ImproperConfigurationError(msg) conn_kwargs: dict[str, Any] = { "instance_uri": self._config.driver_features["alloydb_instance_uri"], "driver": "asyncpg", "enable_iam_auth": self._config.driver_features.get("enable_alloydb_iam_auth", False), "ip_type": self._config.driver_features.get("alloydb_ip_type", "PRIVATE"), } if self._user: conn_kwargs["user"] = self._user if self._password: conn_kwargs["password"] = self._password if self._database: conn_kwargs["db"] = self._database return cast("AsyncpgConnection", await connector.connect(**conn_kwargs)) class _AsyncpgSessionFactory(AsyncPoolSessionFactory): __slots__ = () class AsyncpgConnectionContext(AsyncPoolConnectionContext): """Async context manager for AsyncPG connections.""" __slots__ = ()

[docs]

@mypyc_attr(native_class=False) class AsyncpgConfig(AsyncDatabaseConfig[AsyncpgConnection, "Pool[Record]", AsyncpgDriver]): """Configuration for AsyncPG database connections using TypedDict.""" driver_type: "ClassVar[type[AsyncpgDriver]]" = AsyncpgDriver connection_type: "ClassVar[type[AsyncpgConnection]]" = AsyncpgConnection # 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[AsyncpgConnectionContext]]" = AsyncpgConnectionContext _session_factory_class: "ClassVar[type[_AsyncpgSessionFactory]]" = _AsyncpgSessionFactory _session_context_class: "ClassVar[type[AsyncpgSessionContext]]" = AsyncpgSessionContext _default_statement_config = default_statement_config

[docs]

def __init__( self, *, connection_config: "AsyncpgPoolConfig | dict[str, Any] | None" = None, connection_instance: "Pool[Record] | None" = None, migration_config: "dict[str, Any] | None" = None, statement_config: "StatementConfig | None" = None, driver_features: "AsyncpgDriverFeatures | dict[str, Any] | None" = None, bind_key: "str | None" = None, extension_config: "ExtensionConfigs | None" = None, observability_config: "ObservabilityConfig | None" = None, **kwargs: Any, ) -> None: """Initialize AsyncPG configuration. Args: connection_config: Connection and pool configuration parameters (TypedDict or dict) connection_instance: Existing pool instance to use migration_config: Migration configuration statement_config: Statement configuration override driver_features: Driver features configuration (TypedDict or dict) bind_key: Optional unique identifier for this configuration extension_config: Extension-specific configuration observability_config: Adapter-level observability overrides for lifecycle hooks and observers **kwargs: Additional keyword arguments """ statement_config = statement_config or default_statement_config statement_config, driver_features = apply_driver_features(statement_config, driver_features) # Extract user connection hook before storing driver_features features_dict = dict(driver_features) self._user_connection_hook: Callable[[AsyncpgConnection], Awaitable[None]] | None = features_dict.pop( "on_connection_create", None ) super().__init__( connection_config=normalize_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, ) self._cloud_sql_connector: Any | None = None self._alloydb_connector: Any | None = None self._pgvector_available: bool | None = None self._paradedb_available: bool | None = None self._validate_connector_config()

[docs]

def get_cloud_sql_connector(self) -> Any | None: """Return the configured Cloud SQL connector instance.""" return self._cloud_sql_connector