sqlspec 0.32.0__py3-none-any.whl

sqlspec/adapters/psqlpy/config.py

@@ -0,0 +1,271 @@
+"""Psqlpy database configuration."""
+
+import logging
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, Any, ClassVar, TypedDict, cast
+
+from psqlpy import ConnectionPool
+from typing_extensions import NotRequired
+
+from sqlspec.adapters.psqlpy._types import PsqlpyConnection
+from sqlspec.adapters.psqlpy.driver import (
+    PsqlpyCursor,
+    PsqlpyDriver,
+    PsqlpyExceptionHandler,
+    build_psqlpy_statement_config,
+)
+from sqlspec.config import AsyncDatabaseConfig, ExtensionConfigs
+from sqlspec.core import StatementConfig
+from sqlspec.typing import PGVECTOR_INSTALLED
+from sqlspec.utils.serializers import to_json
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+logger = logging.getLogger("sqlspec.adapters.psqlpy")
+
+
+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.
+    json_serializer: Custom JSON serializer applied to the statement configuration.
+    json_deserializer: Custom JSON deserializer retained alongside the serializer for parity with asyncpg.
+    """
+
+    enable_pgvector: NotRequired[bool]
+    json_serializer: NotRequired["Callable[[Any], str]"]
+    json_deserializer: NotRequired["Callable[[str], Any]"]
+
+
+__all__ = ("PsqlpyConfig", "PsqlpyConnectionParams", "PsqlpyCursor", "PsqlpyDriverFeatures", "PsqlpyPoolParams")
+
+
+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_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
+
+    def __init__(
+        self,
+        *,
+        pool_config: PsqlpyPoolParams | dict[str, Any] | None = None,
+        pool_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,
+    ) -> None:
+        """Initialize Psqlpy configuration.
+
+        Args:
+            pool_config: Pool configuration parameters.
+            pool_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).
+        """
+        processed_pool_config: dict[str, Any] = dict(pool_config) if pool_config else {}
+        if "extra" in processed_pool_config:
+            extras = processed_pool_config.pop("extra")
+            processed_pool_config.update(extras)
+
+        processed_driver_features: dict[str, Any] = dict(driver_features) if driver_features else {}
+        serializer = processed_driver_features.get("json_serializer")
+        serializer_callable = to_json if serializer is None else cast("Callable[[Any], str]", serializer)
+        processed_driver_features.setdefault("json_serializer", serializer_callable)
+        processed_driver_features.setdefault("enable_pgvector", PGVECTOR_INSTALLED)
+
+        super().__init__(
+            pool_config=processed_pool_config,
+            pool_instance=pool_instance,
+            migration_config=migration_config,
+            statement_config=statement_config or build_psqlpy_statement_config(json_serializer=serializer_callable),
+            driver_features=processed_driver_features,
+            bind_key=bind_key,
+            extension_config=extension_config,
+        )
+
+    def _get_pool_config_dict(self) -> dict[str, Any]:
+        """Get pool configuration as plain dict for external library.
+
+        Returns:
+            Dictionary with pool parameters, filtering out None values.
+        """
+        return {k: v for k, v in self.pool_config.items() if v is not None}
+
+    async def _create_pool(self) -> "ConnectionPool":
+        """Create the actual async connection pool."""
+        logger.info("Creating psqlpy connection pool", extra={"adapter": "psqlpy"})
+
+        try:
+            config = self._get_pool_config_dict()
+
+            pool = ConnectionPool(**config)
+            logger.info("Psqlpy connection pool created successfully", extra={"adapter": "psqlpy"})
+        except Exception as e:
+            logger.exception("Failed to create psqlpy connection pool", extra={"adapter": "psqlpy", "error": str(e)})
+            raise
+        return pool
+
+    async def _close_pool(self) -> None:
+        """Close the actual async connection pool."""
+        if not self.pool_instance:
+            return
+
+        logger.info("Closing psqlpy connection pool", extra={"adapter": "psqlpy"})
+
+        try:
+            self.pool_instance.close()
+            logger.info("Psqlpy connection pool closed successfully", extra={"adapter": "psqlpy"})
+        except Exception as e:
+            logger.exception("Failed to close psqlpy connection pool", extra={"adapter": "psqlpy", "error": str(e)})
+            raise
+
+    async def close_pool(self) -> None:
+        """Close the connection pool."""
+        await self._close_pool()
+
+    async def create_connection(self) -> "PsqlpyConnection":
+        """Create a single async connection (not from pool).
+
+        Returns:
+            A psqlpy Connection instance.
+        """
+        if not self.pool_instance:
+            self.pool_instance = await self._create_pool()
+
+        return await self.pool_instance.connection()
+
+    @asynccontextmanager
+    async def provide_connection(self, *args: Any, **kwargs: Any) -> AsyncGenerator[PsqlpyConnection, None]:
+        """Provide an async connection context manager.
+
+        Args:
+            *args: Additional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Yields:
+            A psqlpy Connection instance.
+        """
+        if not self.pool_instance:
+            self.pool_instance = await self._create_pool()
+
+        async with self.pool_instance.acquire() as conn:
+            yield conn
+
+    @asynccontextmanager
+    async def provide_session(
+        self, *args: Any, statement_config: "StatementConfig | None" = None, **kwargs: Any
+    ) -> AsyncGenerator[PsqlpyDriver, None]:
+        """Provide an async driver session context manager.
+
+        Args:
+            *args: Additional arguments.
+            statement_config: Optional statement configuration override.
+            **kwargs: Additional keyword arguments.
+
+        Yields:
+            A PsqlpyDriver instance.
+        """
+        async with self.provide_connection(*args, **kwargs) as conn:
+            driver = self.driver_type(
+                connection=conn,
+                statement_config=statement_config or self.statement_config,
+                driver_features=self.driver_features,
+            )
+            yield self._prepare_driver(driver)
+
+    async def provide_pool(self, *args: Any, **kwargs: Any) -> ConnectionPool:
+        """Provide async pool instance.
+
+        Returns:
+            The async connection pool.
+        """
+        if not self.pool_instance:
+            self.pool_instance = await self.create_pool()
+        return self.pool_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({
+            "PsqlpyConnection": PsqlpyConnection,
+            "PsqlpyConnectionParams": PsqlpyConnectionParams,
+            "PsqlpyCursor": PsqlpyCursor,
+            "PsqlpyDriver": PsqlpyDriver,
+            "PsqlpyExceptionHandler": PsqlpyExceptionHandler,
+            "PsqlpyPoolParams": PsqlpyPoolParams,
+        })
+        return namespace

sqlspec/adapters/psqlpy/data_dictionary.py

@@ -0,0 +1,179 @@
+"""PostgreSQL-specific data dictionary for metadata queries via psqlpy."""
+
+import re
+from typing import TYPE_CHECKING, Any, cast
+
+from sqlspec.driver import AsyncDataDictionaryBase, AsyncDriverAdapterBase, VersionInfo
+from sqlspec.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from sqlspec.adapters.psqlpy.driver import PsqlpyDriver
+
+logger = get_logger("adapters.psqlpy.data_dictionary")
+
+# Compiled regex patterns
+POSTGRES_VERSION_PATTERN = re.compile(r"PostgreSQL (\d+)\.(\d+)(?:\.(\d+))?")
+
+__all__ = ("PsqlpyAsyncDataDictionary",)
+
+
+class PsqlpyAsyncDataDictionary(AsyncDataDictionaryBase):
+    """PostgreSQL-specific async data dictionary via psqlpy."""
+
+    async def get_version(self, driver: AsyncDriverAdapterBase) -> "VersionInfo | None":
+        """Get PostgreSQL database version information.
+
+        Args:
+            driver: Async database driver instance
+
+        Returns:
+            PostgreSQL version information or None if detection fails
+        """
+        version_str = await cast("PsqlpyDriver", driver).select_value("SELECT version()")
+        if not version_str:
+            logger.warning("No PostgreSQL version information found")
+            return None
+
+        # Parse version like "PostgreSQL 15.3 on x86_64-pc-linux-gnu..."
+        version_match = POSTGRES_VERSION_PATTERN.search(str(version_str))
+        if not version_match:
+            logger.warning("Could not parse PostgreSQL version: %s", version_str)
+            return None
+
+        major = int(version_match.group(1))
+        minor = int(version_match.group(2))
+        patch = int(version_match.group(3)) if version_match.group(3) else 0
+
+        version_info = VersionInfo(major, minor, patch)
+        logger.debug("Detected PostgreSQL version: %s", version_info)
+        return version_info
+
+    async def get_feature_flag(self, driver: AsyncDriverAdapterBase, feature: str) -> bool:
+        """Check if PostgreSQL database supports a specific feature.
+
+        Args:
+            driver: Async database driver instance
+            feature: Feature name to check
+
+        Returns:
+            True if feature is supported, False otherwise
+        """
+        version_info = await self.get_version(driver)
+        if not version_info:
+            return False
+
+        feature_checks: dict[str, Callable[[VersionInfo], bool]] = {
+            "supports_json": lambda v: v >= VersionInfo(9, 2, 0),
+            "supports_jsonb": lambda v: v >= VersionInfo(9, 4, 0),
+            "supports_uuid": lambda _: True,  # UUID extension widely available
+            "supports_arrays": lambda _: True,  # PostgreSQL has excellent array support
+            "supports_returning": lambda v: v >= VersionInfo(8, 2, 0),
+            "supports_upsert": lambda v: v >= VersionInfo(9, 5, 0),  # ON CONFLICT
+            "supports_window_functions": lambda v: v >= VersionInfo(8, 4, 0),
+            "supports_cte": lambda v: v >= VersionInfo(8, 4, 0),
+            "supports_transactions": lambda _: True,
+            "supports_prepared_statements": lambda _: True,
+            "supports_schemas": lambda _: True,
+            "supports_partitioning": lambda v: v >= VersionInfo(10, 0, 0),
+        }
+
+        if feature in feature_checks:
+            return bool(feature_checks[feature](version_info))
+
+        return False
+
+    async def get_optimal_type(self, driver: AsyncDriverAdapterBase, type_category: str) -> str:
+        """Get optimal PostgreSQL type for a category.
+
+        Args:
+            driver: Async database driver instance
+            type_category: Type category
+
+        Returns:
+            PostgreSQL-specific type name
+        """
+        version_info = await self.get_version(driver)
+
+        if type_category == "json":
+            if version_info and version_info >= VersionInfo(9, 4, 0):
+                return "JSONB"  # Prefer JSONB over JSON
+            if version_info and version_info >= VersionInfo(9, 2, 0):
+                return "JSON"
+            return "TEXT"
+
+        type_map = {
+            "uuid": "UUID",
+            "boolean": "BOOLEAN",
+            "timestamp": "TIMESTAMP WITH TIME ZONE",
+            "text": "TEXT",
+            "blob": "BYTEA",
+            "array": "ARRAY",
+        }
+        return type_map.get(type_category, "TEXT")
+
+    async def get_columns(
+        self, driver: AsyncDriverAdapterBase, table: str, schema: "str | None" = None
+    ) -> "list[dict[str, Any]]":
+        """Get column information for a table using pg_catalog.
+
+        Args:
+            driver: Psqlpy async driver instance
+            table: Table name to query columns for
+            schema: Schema name (None for default 'public')
+
+        Returns:
+            List of column metadata dictionaries with keys:
+                - column_name: Name of the column
+                - data_type: PostgreSQL data type
+                - is_nullable: Whether column allows NULL (YES/NO)
+                - column_default: Default value if any
+
+        Notes:
+            Uses pg_catalog instead of information_schema to avoid psqlpy's
+            inability to handle the PostgreSQL 'name' type returned by information_schema.
+        """
+        psqlpy_driver = cast("PsqlpyDriver", driver)
+
+        schema_name = schema or "public"
+        sql = """
+            SELECT
+                a.attname::text AS column_name,
+                pg_catalog.format_type(a.atttypid, a.atttypmod) AS data_type,
+                CASE WHEN a.attnotnull THEN 'NO' ELSE 'YES' END AS is_nullable,
+                pg_catalog.pg_get_expr(d.adbin, d.adrelid)::text AS column_default
+            FROM pg_catalog.pg_attribute a
+            JOIN pg_catalog.pg_class c ON a.attrelid = c.oid
+            JOIN pg_catalog.pg_namespace n ON c.relnamespace = n.oid
+            LEFT JOIN pg_catalog.pg_attrdef d ON a.attrelid = d.adrelid AND a.attnum = d.adnum
+            WHERE c.relname = $1
+                AND n.nspname = $2
+                AND a.attnum > 0
+                AND NOT a.attisdropped
+            ORDER BY a.attnum
+        """
+
+        result = await psqlpy_driver.execute(sql, (table, schema_name))
+        return result.data or []
+
+    def list_available_features(self) -> "list[str]":
+        """List available PostgreSQL feature flags.
+
+        Returns:
+            List of supported feature names
+        """
+        return [
+            "supports_json",
+            "supports_jsonb",
+            "supports_uuid",
+            "supports_arrays",
+            "supports_returning",
+            "supports_upsert",
+            "supports_window_functions",
+            "supports_cte",
+            "supports_transactions",
+            "supports_prepared_statements",
+            "supports_schemas",
+            "supports_partitioning",
+        ]