sqlspec 0.25.0__py3-none-any.whl → 0.27.0__py3-none-any.whl

sqlspec/adapters/psycopg/config.py

@@ -3,7 +3,7 @@
 import contextlib
 import logging
 from contextlib import asynccontextmanager
-from typing import TYPE_CHECKING, Any, ClassVar, Optional, TypedDict, Union, cast
+from typing import TYPE_CHECKING, Any, ClassVar, TypedDict, cast
 
 from psycopg.rows import dict_row
 from psycopg_pool import AsyncConnectionPool, ConnectionPool
@@ -18,6 +18,7 @@ from sqlspec.adapters.psycopg.driver import (
     psycopg_statement_config,
 )
 from sqlspec.config import AsyncDatabaseConfig, SyncDatabaseConfig
+from sqlspec.typing import PGVECTOR_INSTALLED
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Callable, Generator
@@ -28,7 +29,7 @@ if TYPE_CHECKING:
 logger = logging.getLogger("sqlspec.adapters.psycopg")
 
 
-class PsycopgConnectionParams(TypedDict, total=False):
+class PsycopgConnectionParams(TypedDict):
     """Psycopg connection parameters."""
 
     conninfo: NotRequired[str]
@@ -48,7 +49,7 @@ class PsycopgConnectionParams(TypedDict, total=False):
     extra: NotRequired[dict[str, Any]]
 
 
-class PsycopgPoolParams(PsycopgConnectionParams, total=False):
+class PsycopgPoolParams(PsycopgConnectionParams):
     """Psycopg pool parameters."""
 
     min_size: NotRequired[int]
@@ -64,10 +65,25 @@ class PsycopgPoolParams(PsycopgConnectionParams, total=False):
     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_pgvector: NotRequired[bool]
+
+
 __all__ = (
     "PsycopgAsyncConfig",
     "PsycopgAsyncCursor",
     "PsycopgConnectionParams",
+    "PsycopgDriverFeatures",
     "PsycopgPoolParams",
     "PsycopgSyncConfig",
     "PsycopgSyncCursor",
@@ -79,15 +95,18 @@ class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, ConnectionPool
 
     driver_type: "ClassVar[type[PsycopgSyncDriver]]" = PsycopgSyncDriver
     connection_type: "ClassVar[type[PsycopgSyncConnection]]" = PsycopgSyncConnection
+    supports_transactional_ddl: "ClassVar[bool]" = True
 
     def __init__(
         self,
         *,
-        pool_config: "Optional[Union[PsycopgPoolParams, dict[str, Any]]]" = None,
-        pool_instance: Optional["ConnectionPool"] = None,
-        migration_config: Optional[dict[str, Any]] = None,
-        statement_config: "Optional[StatementConfig]" = None,
-        driver_features: "Optional[dict[str, Any]]" = None,
+        pool_config: "PsycopgPoolParams | dict[str, Any] | None" = None,
+        pool_instance: "ConnectionPool | None" = None,
+        migration_config: dict[str, Any] | None = None,
+        statement_config: "StatementConfig | None" = None,
+        driver_features: "dict[str, Any] | None" = None,
+        bind_key: "str | None" = None,
+        extension_config: "dict[str, dict[str, Any]] | None" = None,
     ) -> None:
         """Initialize Psycopg synchronous configuration.
 
@@ -97,18 +116,27 @@ class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, ConnectionPool
             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 (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)
 
+        if driver_features is None:
+            driver_features = {}
+        if "enable_pgvector" not in driver_features:
+            driver_features["enable_pgvector"] = PGVECTOR_INSTALLED
+
         super().__init__(
             pool_config=processed_pool_config,
             pool_instance=pool_instance,
             migration_config=migration_config,
             statement_config=statement_config or psycopg_statement_config,
-            driver_features=driver_features or {},
+            driver_features=driver_features,
+            bind_key=bind_key,
+            extension_config=extension_config,
         )
 
     def _create_pool(self) -> "ConnectionPool":
@@ -137,15 +165,10 @@ class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, ConnectionPool
                 if autocommit_setting is not None:
                     conn.autocommit = autocommit_setting
 
-                try:
-                    import pgvector.psycopg
+                if self.driver_features.get("enable_pgvector", False):
+                    from sqlspec.adapters.psycopg._type_handlers import register_pgvector_sync
 
-                    pgvector.psycopg.register_vector(conn)
-                    logger.debug("pgvector registered successfully for psycopg sync connection")
-                except ImportError:
-                    pass
-                except Exception as e:
-                    logger.debug("Failed to register pgvector for psycopg sync: %s", e)
+                    register_pgvector_sync(conn)
 
             pool_parameters["configure"] = all_config.pop("configure", configure_connection)
 
@@ -216,7 +239,7 @@ class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, ConnectionPool
 
     @contextlib.contextmanager
     def provide_session(
-        self, *args: Any, statement_config: "Optional[StatementConfig]" = None, **kwargs: Any
+        self, *args: Any, statement_config: "StatementConfig | None" = None, **kwargs: Any
     ) -> "Generator[PsycopgSyncDriver, None, None]":
         """Provide a driver session context manager.
 
@@ -230,7 +253,9 @@ class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, ConnectionPool
         """
         with self.provide_connection(*args, **kwargs) as conn:
             final_statement_config = statement_config or self.statement_config
-            yield self.driver_type(connection=conn, statement_config=final_statement_config)
+            yield self.driver_type(
+                connection=conn, statement_config=final_statement_config, driver_features=self.driver_features
+            )
 
     def provide_pool(self, *args: Any, **kwargs: Any) -> "ConnectionPool":
         """Provide pool instance.
@@ -261,15 +286,18 @@ class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnec
 
     driver_type: ClassVar[type[PsycopgAsyncDriver]] = PsycopgAsyncDriver
     connection_type: "ClassVar[type[PsycopgAsyncConnection]]" = PsycopgAsyncConnection
+    supports_transactional_ddl: "ClassVar[bool]" = True
 
     def __init__(
         self,
         *,
-        pool_config: "Optional[Union[PsycopgPoolParams, dict[str, Any]]]" = None,
-        pool_instance: "Optional[AsyncConnectionPool]" = None,
-        migration_config: "Optional[dict[str, Any]]" = None,
-        statement_config: "Optional[StatementConfig]" = None,
-        driver_features: "Optional[dict[str, Any]]" = None,
+        pool_config: "PsycopgPoolParams | dict[str, Any] | None" = None,
+        pool_instance: "AsyncConnectionPool | None" = None,
+        migration_config: "dict[str, Any] | None" = None,
+        statement_config: "StatementConfig | None" = None,
+        driver_features: "dict[str, Any] | None" = None,
+        bind_key: "str | None" = None,
+        extension_config: "dict[str, dict[str, Any]] | None" = None,
     ) -> None:
         """Initialize Psycopg asynchronous configuration.
 
@@ -279,18 +307,27 @@ class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnec
             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 (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)
 
+        if driver_features is None:
+            driver_features = {}
+        if "enable_pgvector" not in driver_features:
+            driver_features["enable_pgvector"] = PGVECTOR_INSTALLED
+
         super().__init__(
             pool_config=processed_pool_config,
             pool_instance=pool_instance,
             migration_config=migration_config,
             statement_config=statement_config or psycopg_statement_config,
-            driver_features=driver_features or {},
+            driver_features=driver_features,
+            bind_key=bind_key,
+            extension_config=extension_config,
         )
 
     async def _create_pool(self) -> "AsyncConnectionPool":
@@ -317,15 +354,10 @@ class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnec
             if autocommit_setting is not None:
                 await conn.set_autocommit(autocommit_setting)
 
-            try:
-                from pgvector.psycopg import register_vector_async
+            if self.driver_features.get("enable_pgvector", False):
+                from sqlspec.adapters.psycopg._type_handlers import register_pgvector_async
 
-                await register_vector_async(conn)
-                logger.debug("pgvector registered successfully for psycopg async connection")
-            except ImportError:
-                pass
-            except Exception as e:
-                logger.debug("Failed to register pgvector for psycopg async: %s", e)
+                await register_pgvector_async(conn)
 
         pool_parameters["configure"] = all_config.pop("configure", configure_connection)
 
@@ -388,7 +420,7 @@ class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnec
 
     @asynccontextmanager
     async def provide_session(
-        self, *args: Any, statement_config: "Optional[StatementConfig]" = None, **kwargs: Any
+        self, *args: Any, statement_config: "StatementConfig | None" = None, **kwargs: Any
     ) -> "AsyncGenerator[PsycopgAsyncDriver, None]":
         """Provide an async driver session context manager.
 
@@ -402,7 +434,9 @@ class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnec
         """
         async with self.provide_connection(*args, **kwargs) as conn:
             final_statement_config = statement_config or psycopg_statement_config
-            yield self.driver_type(connection=conn, statement_config=final_statement_config)
+            yield self.driver_type(
+                connection=conn, statement_config=final_statement_config, driver_features=self.driver_features
+            )
 
     async def provide_pool(self, *args: Any, **kwargs: Any) -> "AsyncConnectionPool":
         """Provide async pool instance.

sqlspec/adapters/psycopg/data_dictionary.py

@@ -0,0 +1,331 @@
+"""PostgreSQL-specific data dictionary for metadata queries via psycopg."""
+
+import re
+from typing import TYPE_CHECKING, Any, cast
+
+from sqlspec.driver import (
+    AsyncDataDictionaryBase,
+    AsyncDriverAdapterBase,
+    SyncDataDictionaryBase,
+    SyncDriverAdapterBase,
+    VersionInfo,
+)
+from sqlspec.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from sqlspec.adapters.psycopg.driver import PsycopgAsyncDriver, PsycopgSyncDriver
+
+logger = get_logger("adapters.psycopg.data_dictionary")
+
+# Compiled regex patterns
+POSTGRES_VERSION_PATTERN = re.compile(r"PostgreSQL (\d+)\.(\d+)(?:\.(\d+))?")
+
+__all__ = ("PostgresAsyncDataDictionary", "PostgresSyncDataDictionary")
+
+
+class PostgresSyncDataDictionary(SyncDataDictionaryBase):
+    """PostgreSQL-specific sync data dictionary."""
+
+    def get_version(self, driver: SyncDriverAdapterBase) -> "VersionInfo | None":
+        """Get PostgreSQL database version information.
+
+        Args:
+            driver: Sync database driver instance
+
+        Returns:
+            PostgreSQL version information or None if detection fails
+        """
+        version_str = cast("PsycopgSyncDriver", 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
+
+    def get_feature_flag(self, driver: SyncDriverAdapterBase, feature: str) -> bool:
+        """Check if PostgreSQL database supports a specific feature.
+
+        Args:
+            driver: Sync database driver instance
+            feature: Feature name to check
+
+        Returns:
+            True if feature is supported, False otherwise
+        """
+        version_info = 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
+
+    def get_optimal_type(self, driver: SyncDriverAdapterBase, type_category: str) -> str:
+        """Get optimal PostgreSQL type for a category.
+
+        Args:
+            driver: Sync database driver instance
+            type_category: Type category
+
+        Returns:
+            PostgreSQL-specific type name
+        """
+        version_info = 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")
+
+    def get_columns(
+        self, driver: SyncDriverAdapterBase, table: str, schema: "str | None" = None
+    ) -> "list[dict[str, Any]]":
+        """Get column information for a table using information_schema.
+
+        Args:
+            driver: Psycopg sync 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
+        """
+        psycopg_driver = cast("PsycopgSyncDriver", driver)
+
+        if schema:
+            sql = f"""
+                SELECT column_name, data_type, is_nullable, column_default
+                FROM information_schema.columns
+                WHERE table_name = '{table}' AND table_schema = '{schema}'
+                ORDER BY ordinal_position
+            """
+        else:
+            sql = f"""
+                SELECT column_name, data_type, is_nullable, column_default
+                FROM information_schema.columns
+                WHERE table_name = '{table}' AND table_schema = 'public'
+                ORDER BY ordinal_position
+            """
+
+        result = psycopg_driver.execute(sql)
+        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",
+        ]
+
+
+class PostgresAsyncDataDictionary(AsyncDataDictionaryBase):
+    """PostgreSQL-specific async data dictionary."""
+
+    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("PsycopgAsyncDriver", 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 information_schema.
+
+        Args:
+            driver: Psycopg 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
+        """
+        psycopg_driver = cast("PsycopgAsyncDriver", driver)
+
+        if schema:
+            sql = f"""
+                SELECT column_name, data_type, is_nullable, column_default
+                FROM information_schema.columns
+                WHERE table_name = '{table}' AND table_schema = '{schema}'
+                ORDER BY ordinal_position
+            """
+        else:
+            sql = f"""
+                SELECT column_name, data_type, is_nullable, column_default
+                FROM information_schema.columns
+                WHERE table_name = '{table}' AND table_schema = 'public'
+                ORDER BY ordinal_position
+            """
+
+        result = await psycopg_driver.execute(sql)
+        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",
+        ]