sqlspec 0.26.0__py3-none-any.whl → 0.28.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,16 +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,
-        bind_key: "Optional[str]" = 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.
 
@@ -99,19 +117,26 @@ class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, ConnectionPool
             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":
@@ -140,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)
 
@@ -219,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.
 
@@ -233,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.
@@ -264,16 +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,
-        bind_key: "Optional[str]" = 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.
 
@@ -284,19 +308,26 @@ class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnec
             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":
@@ -323,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)
 
@@ -394,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.
 
@@ -408,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

@@ -1,7 +1,7 @@
 """PostgreSQL-specific data dictionary for metadata queries via psycopg."""
 
 import re
-from typing import TYPE_CHECKING, Optional, cast
+from typing import TYPE_CHECKING, Any, cast
 
 from sqlspec.driver import (
     AsyncDataDictionaryBase,
@@ -28,7 +28,7 @@ __all__ = ("PostgresAsyncDataDictionary", "PostgresSyncDataDictionary")
 class PostgresSyncDataDictionary(SyncDataDictionaryBase):
     """PostgreSQL-specific sync data dictionary."""
 
-    def get_version(self, driver: SyncDriverAdapterBase) -> "Optional[VersionInfo]":
+    def get_version(self, driver: SyncDriverAdapterBase) -> "VersionInfo | None":
         """Get PostgreSQL database version information.
 
         Args:
@@ -119,6 +119,50 @@ class PostgresSyncDataDictionary(SyncDataDictionaryBase):
         }
         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 pg_catalog.
+
+        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
+
+        Notes:
+            Uses pg_catalog instead of information_schema to avoid potential
+            issues with PostgreSQL 'name' type in some drivers.
+        """
+        psycopg_driver = cast("PsycopgSyncDriver", 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 = %s
+                AND n.nspname = %s
+                AND a.attnum > 0
+                AND NOT a.attisdropped
+            ORDER BY a.attnum
+        """
+
+        result = psycopg_driver.execute(sql, (table, schema_name))
+        return result.data or []
+
     def list_available_features(self) -> "list[str]":
         """List available PostgreSQL feature flags.
 
@@ -144,7 +188,7 @@ class PostgresSyncDataDictionary(SyncDataDictionaryBase):
 class PostgresAsyncDataDictionary(AsyncDataDictionaryBase):
     """PostgreSQL-specific async data dictionary."""
 
-    async def get_version(self, driver: AsyncDriverAdapterBase) -> "Optional[VersionInfo]":
+    async def get_version(self, driver: AsyncDriverAdapterBase) -> "VersionInfo | None":
         """Get PostgreSQL database version information.
 
         Args:
@@ -235,6 +279,50 @@ class PostgresAsyncDataDictionary(AsyncDataDictionaryBase):
         }
         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: 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
+
+        Notes:
+            Uses pg_catalog instead of information_schema to avoid potential
+            issues with PostgreSQL 'name' type in some drivers.
+        """
+        psycopg_driver = cast("PsycopgAsyncDriver", 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 = %s
+                AND n.nspname = %s
+                AND a.attnum > 0
+                AND NOT a.attisdropped
+            ORDER BY a.attnum
+        """
+
+        result = await psycopg_driver.execute(sql, (table, schema_name))
+        return result.data or []
+
     def list_available_features(self) -> "list[str]":
         """List available PostgreSQL feature flags.