sqlspec 0.47.0__cp314-cp314-win_amd64.whl

sqlspec/adapters/psqlpy/config.py

@@ -0,0 +1,376 @@
+"""Psqlpy database configuration."""
+
+from typing import TYPE_CHECKING, Any, ClassVar, TypedDict, cast
+
+from mypy_extensions import mypyc_attr
+from psqlpy import ConnectionPool
+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.adapters.psqlpy.type_converter import register_pgvector
+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 sqlspec.core import StatementConfig
+
+
+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_pgvector: Enable automatic pgvector extension support for vector similarity search.
+        Requires pgvector-python package installed.
+        Defaults to True when pgvector is installed.
+        Provides automatic conversion between NumPy arrays and PostgreSQL vector types.
+    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 table_queue backend.
+    events_backend: Event channel backend selection.
+        Options: "listen_notify", "table_queue", "listen_notify_durable"
+        - "listen_notify": Zero-copy PostgreSQL LISTEN/NOTIFY (ephemeral, real-time) - coming soon
+        - "table_queue": Durable table-backed queue with retries and exactly-once delivery (current default)
+        - "listen_notify_durable": Hybrid - real-time + durable (available when native support lands)
+        Defaults to "table_queue" until native LISTEN/NOTIFY support is implemented.
+    """
+
+    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[str]
+
+
+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_initialized(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
+
+
+__all__ = ("PsqlpyConfig", "PsqlpyConnectionParams", "PsqlpyCursor", "PsqlpyDriverFeatures", "PsqlpyPoolParams")
+
+
+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_initialized(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.
+
+    Example::
+
+        config = PsqlpyConfig(
+            connection_config=PsqlpyPoolParams(
+                dsn="postgresql://user:pass@localhost/db"
+            )
+        )
+    """
+
+    driver_type: ClassVar[type[PsqlpyDriver]] = PsqlpyDriver
+    connection_type: "ClassVar[type[PsqlpyConnection]]" = PsqlpyConnection
+    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[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.
+
+        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 (e.g., Litestar plugin settings).
+            **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[[PsqlpyConnection], Awaitable[None]] | None = features_dict.pop(
+            "on_connection_create", None
+        )
+        # Track initialized connections by ID (psqlpy connections don't support weak refs)
+        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,
+        )
+
+    async def _ensure_connection_initialized(self, connection: "PsqlpyConnection") -> None:
+        """Ensure connection callback has been called exactly once for this connection."""
+        # 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:
+                    result = await connection.fetch(
+                        "SELECT extname FROM pg_extension WHERE extname = ANY($1::text[])", [extensions]
+                    )
+                    rows = result.result() if result else []
+                    detected_extensions = {r["extname"] for r in rows}
+                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)
+            )
+
+        conn_id = id(connection)
+        if conn_id in self._initialized_connection_ids:
+            return
+        if self._pgvector_available:
+            register_pgvector(connection)
+        if self._user_connection_hook is not None:
+            await self._user_connection_hook(connection)
+        self._initialized_connection_ids.add(conn_id)
+
+    async def _create_pool(self) -> "ConnectionPool":
+        """Create the actual async connection pool."""
+        return ConnectionPool(**build_connection_config(self.connection_config))
+
+    def _update_dialect_for_extensions(self) -> None:
+        """Update statement_config dialect based on detected extensions.
+
+        Priority: paradedb > pgvector > postgres (default).
+        """
+        current_dialect = getattr(self.statement_config, "dialect", "postgres")
+        if current_dialect != "postgres":
+            return
+
+        if self._paradedb_available:
+            self.statement_config = self.statement_config.replace(dialect="paradedb")
+        elif self._pgvector_available:
+            self.statement_config = self.statement_config.replace(dialect="pgvector")
+
+    async def _close_pool(self) -> None:
+        """Close the actual async connection pool."""
+        if not self.connection_instance:
+            return
+
+        self.connection_instance.close()
+        self.connection_instance = None
+
+    async def create_connection(self) -> "PsqlpyConnection":
+        """Create a single async connection (not from pool).
+
+        Returns:
+            A psqlpy Connection instance.
+        """
+        pool = self.connection_instance
+        if pool is None:
+            pool = await self.create_pool()
+            self.connection_instance = pool
+
+        return await pool.connection()
+
+    def provide_session(
+        self, *_args: Any, statement_config: "StatementConfig | None" = None, **_kwargs: Any
+    ) -> "PsqlpySessionContext":
+        """Provide an async driver session context manager.
+
+        Args:
+            *_args: Additional arguments.
+            statement_config: Optional statement configuration override.
+            **_kwargs: Additional keyword arguments.
+
+        Returns:
+            A PsqlpyDriver session context manager.
+        """
+        factory = _PsqlpySessionFactory(self)
+        return PsqlpySessionContext(
+            acquire_connection=factory.acquire_connection,
+            release_connection=factory.release_connection,
+            statement_config=statement_config
+            or (lambda: resolve_runtime_statement_config(None, self.statement_config, default_statement_config)),
+            driver_features=self.driver_features,
+            prepare_driver=self._prepare_driver,
+        )
+
+    async def provide_pool(self, *args: Any, **kwargs: Any) -> ConnectionPool:
+        """Provide async pool instance.
+
+        Returns:
+            The async connection pool.
+        """
+        if not self.connection_instance:
+            self.connection_instance = await self.create_pool()
+        return self.connection_instance
+
+    def get_signature_namespace(self) -> "dict[str, Any]":
+        """Get the signature namespace for Psqlpy types.
+
+        Returns:
+            Dictionary mapping type names to types.
+        """
+        namespace = super().get_signature_namespace()
+        namespace.update({
+            "PsqlpyConnectionContext": PsqlpyConnectionContext,
+            "PsqlpyConnection": PsqlpyConnection,
+            "PsqlpyConnectionParams": PsqlpyConnectionParams,
+            "PsqlpyCursor": PsqlpyCursor,
+            "PsqlpyDriver": PsqlpyDriver,
+            "PsqlpyDriverFeatures": PsqlpyDriverFeatures,
+            "PsqlpyExceptionHandler": PsqlpyExceptionHandler,
+            "PsqlpyPoolParams": PsqlpyPoolParams,
+            "PsqlpySessionContext": PsqlpySessionContext,
+        })
+        return namespace
+
+    def get_event_runtime_hints(self) -> "EventRuntimeHints":
+        """Return LISTEN/NOTIFY defaults for Psqlpy adapters."""
+
+        return EventRuntimeHints(poll_interval=0.5, select_for_update=True, skip_locked=True, json_passthrough=True)