Source code for sqlspec.adapters.psqlpy.config

"""Psqlpy database configuration."""

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

from mypy_extensions import mypyc_attr
from typing_extensions import NotRequired

from sqlspec.adapters.psqlpy._typing import PsqlpyConnection, PsqlpyCursor, PsqlpySessionContext
from sqlspec.adapters.psqlpy.core import (
    apply_driver_features,
    build_connection_config,
    build_postgres_extension_probe_names,
    default_statement_config,
    resolve_postgres_extension_state,
    resolve_runtime_statement_config,
)
from sqlspec.adapters.psqlpy.driver import PsqlpyDriver, PsqlpyExceptionHandler
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

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

    from psqlpy import ConnectionPool

    from sqlspec.core import StatementConfig

__all__ = ("PsqlpyConfig", "PsqlpyConnectionParams", "PsqlpyCursor", "PsqlpyDriverFeatures", "PsqlpyPoolParams")




class PsqlpyConnectionParams(TypedDict):
    """Psqlpy connection parameters."""

    dsn: NotRequired[str]
    username: NotRequired[str]
    password: NotRequired[str]
    db_name: NotRequired[str]
    host: NotRequired[str]
    port: NotRequired[int]
    connect_timeout_sec: NotRequired[int]
    connect_timeout_nanosec: NotRequired[int]
    tcp_user_timeout_sec: NotRequired[int]
    tcp_user_timeout_nanosec: NotRequired[int]
    keepalives: NotRequired[bool]
    keepalives_idle_sec: NotRequired[int]
    keepalives_idle_nanosec: NotRequired[int]
    keepalives_interval_sec: NotRequired[int]
    keepalives_interval_nanosec: NotRequired[int]
    keepalives_retries: NotRequired[int]
    ssl_mode: NotRequired[str]
    ca_file: NotRequired[str]
    target_session_attrs: NotRequired[str]
    options: NotRequired[str]
    application_name: NotRequired[str]
    client_encoding: NotRequired[str]
    gssencmode: NotRequired[str]
    sslnegotiation: NotRequired[str]
    sslcompression: NotRequired[str]
    sslcert: NotRequired[str]
    sslkey: NotRequired[str]
    sslpassword: NotRequired[str]
    sslrootcert: NotRequired[str]
    sslcrl: NotRequired[str]
    require_auth: NotRequired[str]
    channel_binding: NotRequired[str]
    krbsrvname: NotRequired[str]
    gsslib: NotRequired[str]
    gssdelegation: NotRequired[str]
    service: NotRequired[str]
    load_balance_hosts: NotRequired[str]

class PsqlpyPoolParams(PsqlpyConnectionParams): """Psqlpy pool parameters.""" hosts: NotRequired[list[str]] ports: NotRequired[list[int]] conn_recycling_method: NotRequired[str] max_db_pool_size: NotRequired[int] configure: NotRequired["Callable[..., Any]"] extra: NotRequired["dict[str, Any]"] class PsqlpyDriverFeatures(TypedDict): """Psqlpy driver feature flags. enable_cast_detection: Enable cast-aware parameter processing. Defaults to True. When enabled, SQL casts in prepared statements guide psqlpy parameter coercion for JSON, UUID, and timestamp-like values. enable_pgvector: Enable pgvector extension detection for vector similarity search. Defaults to True when the pgvector Python package is installed. Detects the PostgreSQL vector extension and promotes the dialect to "pgvector". It does not register psqlpy type handlers; use psqlpy.extra_types.PgVector or explicit SQL casts for vector values. 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 applied to the statement configuration. json_deserializer: Custom JSON deserializer retained alongside the serializer for parity with asyncpg. on_connection_create: Async callback executed when a connection is acquired from pool. Receives the raw psqlpy connection for low-level driver configuration. Called exactly once per physical connection using WeakSet tracking. 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_cast_detection: NotRequired[bool] 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[[PsqlpyConnection], Awaitable[None]]]" enable_events: NotRequired[bool] events_backend: NotRequired[Literal["notify", "notify_queue", "poll_queue"]] class _PsqlpySessionFactory(AsyncPoolSessionFactory): __slots__ = ("_ctx",) def __init__(self, config: "PsqlpyConfig") -> None: super().__init__(config) self._ctx: Any | None = None async def acquire_connection(self) -> "PsqlpyConnection": 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 connection = cast("PsqlpyConnection", await ctx.__aenter__()) await self._config._ensure_connection(connection) # pyright: ignore[reportPrivateUsage] return connection async def release_connection(self, _conn: "PsqlpyConnection", **kwargs: Any) -> None: if self._ctx is not None: await self._ctx.__aexit__(None, None, None) self._ctx = None class PsqlpyConnectionContext(AsyncPoolConnectionContext): """Async context manager for Psqlpy connections.""" __slots__ = ("_ctx",) def __init__(self, config: "PsqlpyConfig") -> None: super().__init__(config) self._ctx: Any = None async def __aenter__(self) -> PsqlpyConnection: pool = self._config.connection_instance if pool is None: pool = await self._config.create_pool() self._config.connection_instance = pool self._ctx = pool.acquire() connection = await self._ctx.__aenter__() await self._config._ensure_connection(connection) # pyright: ignore[reportPrivateUsage] return connection # type: ignore[no-any-return] async def __aexit__( self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None" ) -> bool | None: if self._ctx: return await self._ctx.__aexit__(exc_type, exc_val, exc_tb) # type: ignore[no-any-return] return None

@mypyc_attr(native_class=False) class PsqlpyConfig(AsyncDatabaseConfig[PsqlpyConnection, "ConnectionPool", PsqlpyDriver]): """Configuration for Psqlpy asynchronous database connections.""" driver_type: ClassVar[type[PsqlpyDriver]] = PsqlpyDriver connection_type: "ClassVar[type[PsqlpyConnection]]" = PsqlpyConnection 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[PsqlpyConnectionContext]]" = PsqlpyConnectionContext _session_factory_class: "ClassVar[type[_PsqlpySessionFactory]]" = _PsqlpySessionFactory _session_context_class: "ClassVar[type[PsqlpySessionContext]]" = PsqlpySessionContext _default_statement_config = default_statement_config

def __init__( self, *, connection_config: "PsqlpyPoolParams | dict[str, Any] | None" = None, connection_instance: "ConnectionPool | None" = None, migration_config: "dict[str, Any] | None" = None, statement_config: "StatementConfig | None" = None, driver_features: "PsqlpyDriverFeatures | dict[str, Any] | None" = None, bind_key: str | None = None, extension_config: "ExtensionConfigs | None" = None, **kwargs: Any, ) -> None: """Initialize Psqlpy configuration. Extracts the 'on_connection_create' hook from driver_features before storing them. Initializes a set to track initialized connection IDs because psqlpy connections do not support weak references. Args: connection_config: Connection and pool configuration parameters. connection_instance: Existing connection pool instance to use. migration_config: Migration configuration. statement_config: SQL statement configuration. driver_features: Driver feature configuration (TypedDict or dict). 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) features_dict = dict(driver_features) if driver_features else {} self._user_connection_hook: Callable[[PsqlpyConnection], Awaitable[None]] | None = features_dict.pop( "on_connection_create", None ) self._initialized_connection_ids: set[int] = set() self._pgvector_available: bool | None = None self._paradedb_available: bool | 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, )