Source code for sqlspec.adapters.psycopg.config

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

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

from mypy_extensions import mypyc_attr
from psycopg import Connection as PsycopgConnection
from psycopg_pool import AsyncConnectionPool, ConnectionPool
from typing_extensions import NotRequired, Self

from sqlspec.adapters.psycopg._typing import (
    PsycopgAsyncConnection,
    PsycopgAsyncCursor,
    PsycopgAsyncSessionContext,
    PsycopgSyncConnection,
    PsycopgSyncCursor,
    PsycopgSyncSessionContext,
)
from sqlspec.adapters.psycopg.core import (
    apply_driver_features,
    build_postgres_extension_probe_names,
    default_statement_config,
    resolve_postgres_extension_state,
    resolve_runtime_statement_config,
)
from sqlspec.adapters.psycopg.driver import (
    PsycopgAsyncDriver,
    PsycopgAsyncExceptionHandler,
    PsycopgSyncDriver,
    PsycopgSyncExceptionHandler,
)
from sqlspec.adapters.psycopg.type_converter import register_pgvector_async, register_pgvector_sync
from sqlspec.config import AsyncDatabaseConfig, ExtensionConfigs, SyncDatabaseConfig
from sqlspec.driver._async import AsyncPoolConnectionContext, AsyncPoolSessionFactory
from sqlspec.driver._sync import SyncPoolConnectionContext, SyncPoolSessionFactory
from sqlspec.exceptions import ImproperConfigurationError, MissingDependencyError
from sqlspec.extensions.events import EventRuntimeHints
from sqlspec.typing import ALLOYDB_CONNECTOR_INSTALLED
from sqlspec.utils.config_tools import normalize_connection_config

if TYPE_CHECKING:
    from collections.abc import Awaitable, Callable
    from types import TracebackType

    from psycopg import AsyncConnection, AsyncCursor, Connection, Cursor
    from psycopg.abc import AdaptContext
    from psycopg.rows import AsyncRowFactory, RowFactory
    from psycopg_pool.abc import AsyncConnectFailedCB, AsyncConnectionCB, ConnectFailedCB, ConnectionCB

    from sqlspec.core import StatementConfig

__all__ = (
    "PsycopgAsyncConfig",
    "PsycopgAsyncCursor",
    "PsycopgConnectionParams",
    "PsycopgDriverFeatures",
    "PsycopgPoolParams",
    "PsycopgSyncConfig",
    "PsycopgSyncCursor",
)


PsycopgSSLMode = Literal["disable", "allow", "prefer", "require", "verify-ca", "verify-full"]
_ALLOYDB_DIRECT_CONNECTION_KEYS = ("conninfo", "host", "hostaddr", "port", "user", "password", "dbname", "db")
_ALLOYDB_CONNECTOR_PACKAGE = "google-cloud-alloydb-connector"



[docs]
class PsycopgConnectionParams(TypedDict):
    """Psycopg connection parameters."""

    conninfo: NotRequired[str]
    host: NotRequired[str]
    port: NotRequired[int]
    user: NotRequired[str]
    password: NotRequired[str]
    dbname: NotRequired[str]
    connect_timeout: NotRequired[int]
    options: NotRequired[str]
    application_name: NotRequired[str]
    sslmode: NotRequired[PsycopgSSLMode]
    sslcert: NotRequired[str]
    sslkey: NotRequired[str]
    sslrootcert: NotRequired[str]
    autocommit: NotRequired[bool]
    prepare_threshold: NotRequired[int | None]
    context: NotRequired["AdaptContext | None"]
    row_factory: NotRequired["RowFactory[Any] | AsyncRowFactory[Any] | None"]
    cursor_factory: NotRequired["type[Cursor[Any] | AsyncCursor[Any]] | None"]
    extra: NotRequired["dict[str, Any]"]

[docs]

class PsycopgPoolParams(PsycopgConnectionParams): """Psycopg pool parameters.""" min_size: NotRequired[int] max_size: NotRequired[int] connection_class: NotRequired["type[Connection[Any] | AsyncConnection[Any]]"] name: NotRequired[str] timeout: NotRequired[float] max_waiting: NotRequired[int] max_lifetime: NotRequired[float] max_idle: NotRequired[float] reconnect_timeout: NotRequired[float] num_workers: NotRequired[int] open: NotRequired[bool | None] configure: NotRequired["ConnectionCB[PsycopgSyncConnection] | AsyncConnectionCB[PsycopgAsyncConnection] | None"] check: NotRequired["ConnectionCB[PsycopgSyncConnection] | AsyncConnectionCB[PsycopgAsyncConnection] | None"] reset: NotRequired["ConnectionCB[PsycopgSyncConnection] | AsyncConnectionCB[PsycopgAsyncConnection] | None"] close_returns: NotRequired[bool] reconnect_failed: NotRequired["ConnectFailedCB | AsyncConnectFailedCB | None"] kwargs: NotRequired["dict[str, Any]"]

class PsycopgDriverFeatures(TypedDict): """Psycopg driver feature flags. enable_pgvector: Enable automatic 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. Set to False to disable pgvector support even when package is available. 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. json_serializer: Custom JSON serializer for StatementConfig parameter handling. json_deserializer: Custom JSON deserializer reference stored alongside the serializer for parity with asyncpg. on_connection_create: Callback executed when a connection is created/acquired from the pool. Receives the raw psycopg connection for low-level driver configuration. Runs after internal setup (pgvector registration). For sync config: Callable[[PsycopgSyncConnection], None] For async config: Callable[[PsycopgAsyncConnection], Awaitable[None]] 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". enable_alloydb: Enable sync-only Google AlloyDB connector integration. Requires google-cloud-alloydb-connector. Defaults to False. alloydb_instance_uri: AlloyDB instance URI in projects/... resource path form. Required when enable_alloydb is True. enable_alloydb_iam_auth: Enable AlloyDB IAM database authentication for sync connector connections. Defaults to False. alloydb_ip_type: AlloyDB connector IP type. Defaults to PRIVATE. """ enable_pgvector: NotRequired[bool] enable_paradedb: NotRequired[bool] json_serializer: NotRequired["Callable[[Any], str]"] json_deserializer: NotRequired["Callable[[str], Any]"] on_connection_create: NotRequired["Callable[..., Any]"] enable_events: NotRequired[bool] events_backend: NotRequired[Literal["notify", "notify_queue", "poll_queue"]] enable_alloydb: NotRequired[bool] alloydb_instance_uri: NotRequired[str] enable_alloydb_iam_auth: NotRequired[bool] alloydb_ip_type: NotRequired[str] def _make_alloydb_connection_class( *, connector: Any, instance_uri: str, enable_iam_auth: bool, ip_type: str, user: str | None, password: str | None, database: str | None, ) -> "type[PsycopgConnection[Any]]": """Build a psycopg connection class backed by the AlloyDB connector.""" class _AlloyDBPsycopgConnection(PsycopgConnection[Any]): @classmethod def connect( cls, conninfo: str = "", *, autocommit: bool = False, prepare_threshold: int | None = 5, context: "AdaptContext | None" = None, row_factory: "RowFactory[Any] | None" = None, cursor_factory: "type[Cursor[Any]] | None" = None, **kwargs: str | int | None, ) -> Self: _ = (cls, conninfo, autocommit, prepare_threshold, context, row_factory, cursor_factory) connector_kwargs = dict(kwargs) connector_kwargs["enable_iam_auth"] = enable_iam_auth connector_kwargs["ip_type"] = ip_type if user is not None: connector_kwargs["user"] = user if password is not None: connector_kwargs["password"] = password if database is not None: connector_kwargs["db"] = database return cast("Self", connector.connect(instance_uri, "psycopg", **connector_kwargs)) return _AlloyDBPsycopgConnection class PsycopgSyncConnectionContext(SyncPoolConnectionContext): """Context manager for Psycopg connections.""" __slots__ = () def __init__(self, config: "PsycopgSyncConfig") -> None: super().__init__(config) def __enter__(self) -> "PsycopgSyncConnection": if self._config.connection_instance: self._ctx = self._config.connection_instance.connection() return cast("PsycopgSyncConnection", self._ctx.__enter__()) # Fallback for no pool self._ctx = self._config.create_connection() return cast("PsycopgSyncConnection", self._ctx) def __exit__( self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None" ) -> bool | None: if self._config.connection_instance and self._ctx: return cast("bool | None", self._ctx.__exit__(exc_type, exc_val, exc_tb)) if self._ctx: self._ctx.close() return None class _PsycopgSyncSessionConnectionHandler(SyncPoolSessionFactory): __slots__ = ("_conn",) def __init__(self, config: "PsycopgSyncConfig") -> None: super().__init__(config) self._conn: PsycopgSyncConnection | None = None def acquire_connection(self) -> "PsycopgSyncConnection": if self._config.connection_instance: self._ctx = self._config.connection_instance.connection() return cast("PsycopgSyncConnection", self._ctx.__enter__()) self._conn = self._config.create_connection() return cast("PsycopgSyncConnection", self._conn) def release_connection(self, _conn: "PsycopgSyncConnection", **kwargs: Any) -> None: if self._ctx is not None: self._ctx.__exit__(None, None, None) self._ctx = None return if self._conn is not None: self._conn.close() self._conn = None

[docs]

@mypyc_attr(native_class=False) class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, ConnectionPool, PsycopgSyncDriver]): """Configuration for Psycopg synchronous database connections with direct field-based configuration.""" driver_type: "ClassVar[type[PsycopgSyncDriver]]" = PsycopgSyncDriver connection_type: "ClassVar[type[PsycopgSyncConnection]]" = PsycopgSyncConnection 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[PsycopgSyncConnectionContext]]" = PsycopgSyncConnectionContext _session_factory_class: "ClassVar[type[_PsycopgSyncSessionConnectionHandler]]" = ( _PsycopgSyncSessionConnectionHandler ) _session_context_class: "ClassVar[type[PsycopgSyncSessionContext]]" = PsycopgSyncSessionContext _default_statement_config = default_statement_config

[docs]

def __init__( self, *, connection_config: "PsycopgPoolParams | dict[str, Any] | None" = None, connection_instance: "ConnectionPool | None" = None, migration_config: "dict[str, Any] | None" = None, statement_config: "StatementConfig | None" = None, driver_features: "PsycopgDriverFeatures | dict[str, Any] | None" = None, bind_key: "str | None" = None, extension_config: "ExtensionConfigs | None" = None, **kwargs: Any, ) -> None: """Initialize Psycopg synchronous 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: Default SQL statement configuration driver_features: Optional driver feature configuration bind_key: Optional unique identifier for this configuration extension_config: Extension-specific configuration **kwargs: Additional keyword arguments """ 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) # Extract user connection hook before storing driver_features features_dict = dict(driver_features) if driver_features else {} self._user_connection_hook: Callable[[PsycopgSyncConnection], None] | None = features_dict.pop( "on_connection_create", None ) features_dict.setdefault("enable_alloydb", False) features_dict.setdefault("enable_alloydb_iam_auth", False) features_dict.setdefault("alloydb_ip_type", "PRIVATE") self._pgvector_available: bool | None = None self._paradedb_available: bool | None = None self._alloydb_connector: Any | None = 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, **kwargs, ) self._validate_alloydb_config()

def _validate_alloydb_config(self) -> None: """Validate sync-only AlloyDB connector configuration.""" if not self.driver_features.get("enable_alloydb", False): return if not ALLOYDB_CONNECTOR_INSTALLED: raise MissingDependencyError(package=_ALLOYDB_CONNECTOR_PACKAGE, install_package=_ALLOYDB_CONNECTOR_PACKAGE) instance_uri = self.driver_features.get("alloydb_instance_uri") if not instance_uri: msg = ( "alloydb_instance_uri required when enable_alloydb is True. " "Format: 'projects/PROJECT/locations/REGION/clusters/CLUSTER/instances/INSTANCE'" ) raise ImproperConfigurationError(msg) if not str(instance_uri).startswith("projects/"): msg = ( f"Invalid AlloyDB instance URI format: {instance_uri}. " "Expected format: 'projects/PROJECT/locations/REGION/clusters/CLUSTER/instances/INSTANCE'" ) raise ImproperConfigurationError(msg) def _setup_alloydb_connector( self, config: "dict[str, Any]", pool_parameters: "dict[str, Any] | None" = None ) -> None: """Setup AlloyDB connector and configure psycopg-pool connection_class.""" from google.cloud.alloydb.connector import Connector # type: ignore[import-untyped,unused-ignore] self._alloydb_connector = Connector() user = config.get("user") password = config.get("password") database = config.get("dbname") or config.get("db") for key in _ALLOYDB_DIRECT_CONNECTION_KEYS: config.pop(key, None) if pool_parameters is not None: pool_parameters["connection_class"] = _make_alloydb_connection_class( connector=self._alloydb_connector, instance_uri=str(self.driver_features["alloydb_instance_uri"]), enable_iam_auth=bool(self.driver_features.get("enable_alloydb_iam_auth", False)), ip_type=str(self.driver_features.get("alloydb_ip_type", "PRIVATE")), user=cast("str | None", user), password=cast("str | None", password), database=cast("str | None", database), ) def _create_pool(self) -> "ConnectionPool": """Create the actual connection pool.""" all_config = dict(self.connection_config) pool_parameters = { "connection_class": all_config.pop("connection_class", None), "min_size": all_config.pop("min_size", 4), "max_size": all_config.pop("max_size", None), "name": all_config.pop("name", None), "timeout": all_config.pop("timeout", 30.0), "max_waiting": all_config.pop("max_waiting", 0), "max_lifetime": all_config.pop("max_lifetime", 3600.0), "max_idle": all_config.pop("max_idle", 600.0), "reconnect_timeout": all_config.pop("reconnect_timeout", 300.0), "reconnect_failed": all_config.pop("reconnect_failed", None), "num_workers": all_config.pop("num_workers", 3), "open": all_config.pop("open", True), "check": all_config.pop("check", None), "reset": all_config.pop("reset", None), "close_returns": all_config.pop("close_returns", False), } pool_parameters["configure"] = all_config.pop("configure", self._configure_connection) pool_parameters = {k: v for k, v in pool_parameters.items() if v is not None} conninfo = all_config.pop("conninfo", None) kwargs = all_config.pop("kwargs", {}) all_config.update(kwargs) if self.driver_features.get("enable_alloydb", False): if conninfo is not None: all_config["conninfo"] = conninfo self._setup_alloydb_connector(all_config, pool_parameters) conninfo = None if conninfo: pool = ConnectionPool(conninfo, kwargs=all_config, **pool_parameters) else: pool = ConnectionPool("", kwargs=all_config, **pool_parameters) return pool def _configure_connection(self, conn: "PsycopgSyncConnection") -> None: autocommit_setting = self.connection_config.get("autocommit") if autocommit_setting is not None: conn.autocommit = autocommit_setting # Detect extensions on first connection, update dialect if self._pgvector_available is None: detected_extensions: set[str] = set() extensions = build_postgres_extension_probe_names(self.driver_features) if extensions: try: cursor = conn.execute( "SELECT extname FROM pg_extension WHERE extname = ANY(%s::text[])", (extensions,) ) results = cursor.fetchall() detected_extensions = {r[0] for r in results} # type: ignore[index] except Exception: detected_extensions = set() self.statement_config, self._pgvector_available, self._paradedb_available = ( resolve_postgres_extension_state(self.statement_config, self.driver_features, detected_extensions) ) if self._pgvector_available: register_pgvector_sync(conn) # Ensure connection is not left in INTRANS state from extension detection or registration if not conn.autocommit: conn.rollback() # Call user-provided callback after internal setup if self._user_connection_hook is not None: self._user_connection_hook(conn) def _close_pool(self) -> None: """Close the actual connection pool and cleanup connectors.""" try: if self.connection_instance: self.connection_instance.close() finally: self.connection_instance = None if self._alloydb_connector is not None: self._alloydb_connector.close() self._alloydb_connector = None

[docs]

def create_connection(self) -> "PsycopgSyncConnection": """Create a single connection (not from pool). Returns: A psycopg Connection instance. """ if self.connection_instance is None: self.connection_instance = self.create_pool() return cast("PsycopgSyncConnection", self.connection_instance.getconn()) # pyright: ignore