sqlspec 0.14.1__py3-none-any.whl → 0.15.0__py3-none-any.whl

sqlspec/adapters/asyncpg/config.py

@@ -1,32 +1,34 @@
 """AsyncPG database configuration with direct field-based configuration."""
 
 import logging
-from collections.abc import AsyncGenerator, Awaitable, Callable
+from collections.abc import Callable
 from contextlib import asynccontextmanager
-from typing import TYPE_CHECKING, Any, ClassVar, TypedDict
+from typing import TYPE_CHECKING, Any, ClassVar, Optional, TypedDict, Union
 
 from asyncpg import Connection, Record
 from asyncpg import create_pool as asyncpg_create_pool
 from asyncpg.connection import ConnectionMeta
 from asyncpg.pool import Pool, PoolConnectionProxy, PoolConnectionProxyMeta
-from typing_extensions import NotRequired, Unpack
+from typing_extensions import NotRequired
 
-from sqlspec.adapters.asyncpg.driver import AsyncpgConnection, AsyncpgDriver
+from sqlspec.adapters.asyncpg._types import AsyncpgConnection
+from sqlspec.adapters.asyncpg.driver import AsyncpgCursor, AsyncpgDriver, asyncpg_statement_config
 from sqlspec.config import AsyncDatabaseConfig
-from sqlspec.statement.sql import SQLConfig
-from sqlspec.typing import DictRow, Empty
 from sqlspec.utils.serializers import from_json, to_json
 
 if TYPE_CHECKING:
     from asyncio.events import AbstractEventLoop
+    from collections.abc import AsyncGenerator, Awaitable
 
+    from sqlspec.core.statement import StatementConfig
 
-__all__ = ("CONNECTION_FIELDS", "POOL_FIELDS", "AsyncpgConfig")
+
+__all__ = ("AsyncpgConfig", "AsyncpgConnectionConfig", "AsyncpgDriverFeatures", "AsyncpgPoolConfig")
 
 logger = logging.getLogger("sqlspec")
 
 
-class AsyncpgConnectionParams(TypedDict, total=False):
+class AsyncpgConnectionConfig(TypedDict, total=False):
     """TypedDict for AsyncPG connection parameters."""
 
     dsn: NotRequired[str]
@@ -35,7 +37,7 @@ class AsyncpgConnectionParams(TypedDict, total=False):
     user: NotRequired[str]
     password: NotRequired[str]
     database: NotRequired[str]
-    ssl: NotRequired[Any]  # Can be bool, SSLContext, or specific string
+    ssl: NotRequired[Any]
     passfile: NotRequired[str]
     direct_tls: NotRequired[bool]
     connect_timeout: NotRequired[float]
@@ -46,7 +48,7 @@ class AsyncpgConnectionParams(TypedDict, total=False):
     server_settings: NotRequired[dict[str, str]]
 
 
-class AsyncpgPoolParams(AsyncpgConnectionParams, total=False):
+class AsyncpgPoolConfig(AsyncpgConnectionConfig, total=False):
     """TypedDict for AsyncPG pool parameters, inheriting connection parameters."""
 
     min_size: NotRequired[int]
@@ -58,189 +60,92 @@ class AsyncpgPoolParams(AsyncpgConnectionParams, total=False):
     loop: NotRequired["AbstractEventLoop"]
     connection_class: NotRequired[type["AsyncpgConnection"]]
     record_class: NotRequired[type[Record]]
+    extra: NotRequired[dict[str, Any]]
 
 
-class DriverParameters(AsyncpgPoolParams, total=False):
-    """TypedDict for additional parameters that can be passed to AsyncPG."""
+class AsyncpgDriverFeatures(TypedDict, total=False):
+    """TypedDict for AsyncPG driver features configuration."""
 
-    statement_config: NotRequired[SQLConfig]
-    default_row_type: NotRequired[type[DictRow]]
     json_serializer: NotRequired[Callable[[Any], str]]
     json_deserializer: NotRequired[Callable[[str], Any]]
-    pool_instance: NotRequired["Pool[Record]"]
-    extras: NotRequired[dict[str, Any]]
-
-
-CONNECTION_FIELDS = {
-    "dsn",
-    "host",
-    "port",
-    "user",
-    "password",
-    "database",
-    "ssl",
-    "passfile",
-    "direct_tls",
-    "connect_timeout",
-    "command_timeout",
-    "statement_cache_size",
-    "max_cached_statement_lifetime",
-    "max_cacheable_statement_size",
-    "server_settings",
-}
-POOL_FIELDS = CONNECTION_FIELDS.union(
-    {
-        "min_size",
-        "max_size",
-        "max_queries",
-        "max_inactive_connection_lifetime",
-        "setup",
-        "init",
-        "loop",
-        "connection_class",
-        "record_class",
-    }
-)
 
 
 class AsyncpgConfig(AsyncDatabaseConfig[AsyncpgConnection, "Pool[Record]", AsyncpgDriver]):
     """Configuration for AsyncPG database connections using TypedDict."""
 
-    driver_type: type[AsyncpgDriver] = AsyncpgDriver
-    connection_type: type[AsyncpgConnection] = type(AsyncpgConnection)  # type: ignore[assignment]
-    supported_parameter_styles: ClassVar[tuple[str, ...]] = ("numeric",)
-    default_parameter_style: ClassVar[str] = "numeric"
-
-    def __init__(self, **kwargs: "Unpack[DriverParameters]") -> None:
-        """Initialize AsyncPG configuration."""
-        # Known fields that are part of the config
-        known_fields = {
-            "dsn",
-            "host",
-            "port",
-            "user",
-            "password",
-            "database",
-            "ssl",
-            "passfile",
-            "direct_tls",
-            "connect_timeout",
-            "command_timeout",
-            "statement_cache_size",
-            "max_cached_statement_lifetime",
-            "max_cacheable_statement_size",
-            "server_settings",
-            "min_size",
-            "max_size",
-            "max_queries",
-            "max_inactive_connection_lifetime",
-            "setup",
-            "init",
-            "loop",
-            "connection_class",
-            "record_class",
-            "extras",
-            "statement_config",
-            "default_row_type",
-            "json_serializer",
-            "json_deserializer",
-            "pool_instance",
-        }
-
-        self.dsn = kwargs.get("dsn")
-        self.host = kwargs.get("host")
-        self.port = kwargs.get("port")
-        self.user = kwargs.get("user")
-        self.password = kwargs.get("password")
-        self.database = kwargs.get("database")
-        self.ssl = kwargs.get("ssl")
-        self.passfile = kwargs.get("passfile")
-        self.direct_tls = kwargs.get("direct_tls")
-        self.connect_timeout = kwargs.get("connect_timeout")
-        self.command_timeout = kwargs.get("command_timeout")
-        self.statement_cache_size = kwargs.get("statement_cache_size")
-        self.max_cached_statement_lifetime = kwargs.get("max_cached_statement_lifetime")
-        self.max_cacheable_statement_size = kwargs.get("max_cacheable_statement_size")
-        self.server_settings = kwargs.get("server_settings")
-        self.min_size = kwargs.get("min_size")
-        self.max_size = kwargs.get("max_size")
-        self.max_queries = kwargs.get("max_queries")
-        self.max_inactive_connection_lifetime = kwargs.get("max_inactive_connection_lifetime")
-        self.setup = kwargs.get("setup")
-        self.init = kwargs.get("init")
-        self.loop = kwargs.get("loop")
-        self.connection_class = kwargs.get("connection_class")
-        self.record_class = kwargs.get("record_class")
-
-        # Collect unknown parameters into extras
-        provided_extras = kwargs.get("extras", {})
-        unknown_params = {k: v for k, v in kwargs.items() if k not in known_fields}
-        self.extras = {**provided_extras, **unknown_params}
-
-        self.statement_config = (
-            SQLConfig() if kwargs.get("statement_config") is None else kwargs.get("statement_config")
-        )
-        self.default_row_type = kwargs.get("default_row_type", dict[str, Any])
-        self.json_serializer = kwargs.get("json_serializer", to_json)
-        self.json_deserializer = kwargs.get("json_deserializer", from_json)
-        pool_instance_from_kwargs = kwargs.get("pool_instance")
-
-        super().__init__()
-
-        # Override prepared statements to True for PostgreSQL since it supports them well
-        self.enable_prepared_statements = kwargs.get("enable_prepared_statements", True)  # type: ignore[assignment]
-
-        if pool_instance_from_kwargs is not None:
-            self.pool_instance = pool_instance_from_kwargs
+    driver_type: "ClassVar[type[AsyncpgDriver]]" = AsyncpgDriver
+    connection_type: "ClassVar[type[AsyncpgConnection]]" = type(AsyncpgConnection)  # type: ignore[assignment]
 
-    @property
-    def connection_config_dict(self) -> dict[str, Any]:
-        """Return the connection configuration as a dict for asyncpg.connect().
+    def __init__(
+        self,
+        *,
+        pool_config: "Optional[Union[AsyncpgPoolConfig, dict[str, Any]]]" = None,
+        pool_instance: "Optional[Pool[Record]]" = None,
+        migration_config: "Optional[dict[str, Any]]" = None,
+        statement_config: "Optional[StatementConfig]" = None,
+        driver_features: "Optional[Union[AsyncpgDriverFeatures, dict[str, Any]]]" = None,
+    ) -> None:
+        """Initialize AsyncPG configuration.
 
-        This method filters out pool-specific parameters that are not valid for asyncpg.connect().
+        Args:
+            pool_config: Pool configuration parameters (TypedDict or dict)
+            pool_instance: Existing pool instance to use
+            migration_config: Migration configuration
+            statement_config: Statement configuration override
+            driver_features: Driver features configuration (TypedDict or dict)
         """
-        # Gather non-None connection parameters
-        config = {
-            field: getattr(self, field)
-            for field in CONNECTION_FIELDS
-            if getattr(self, field, None) is not None and getattr(self, field) is not Empty
-        }
-
-        config.update(self.extras)
-
-        return config
+        features_dict: dict[str, Any] = dict(driver_features) if driver_features else {}
+
+        if "json_serializer" not in features_dict:
+            features_dict["json_serializer"] = to_json
+        if "json_deserializer" not in features_dict:
+            features_dict["json_deserializer"] = from_json
+        super().__init__(
+            pool_config=dict(pool_config) if pool_config else {},
+            pool_instance=pool_instance,
+            migration_config=migration_config,
+            statement_config=statement_config or asyncpg_statement_config,
+            driver_features=features_dict,
+        )
 
-    @property
-    def pool_config_dict(self) -> dict[str, Any]:
-        """Return the full pool configuration as a dict for asyncpg.create_pool().
+    def _get_pool_config_dict(self) -> "dict[str, Any]":
+        """Get pool configuration as plain dict for external library.
 
         Returns:
-            A dictionary containing all pool configuration parameters.
+            Dictionary with pool parameters, filtering out None values.
         """
-        # All AsyncPG parameter names (connection + pool)
-        config = {
-            field: getattr(self, field)
-            for field in POOL_FIELDS
-            if getattr(self, field, None) is not None and getattr(self, field) is not Empty
-        }
-
-        # Merge extras parameters
-        config.update(self.extras)
-
-        return config
+        config: dict[str, Any] = dict(self.pool_config)
+        extras = config.pop("extra", {})
+        config.update(extras)
+        return {k: v for k, v in config.items() if v is not None}
 
     async def _create_pool(self) -> "Pool[Record]":
         """Create the actual async connection pool."""
-        pool_args = self.pool_config_dict
-        return await asyncpg_create_pool(**pool_args)
+        config = self._get_pool_config_dict()
+
+        if "init" not in config:
+            config["init"] = self._init_pgvector_connection
+
+        return await asyncpg_create_pool(**config)
+
+    async def _init_pgvector_connection(self, connection: "AsyncpgConnection") -> None:
+        """Initialize pgvector support for asyncpg connections."""
+        try:
+            import pgvector.asyncpg
+
+            await pgvector.asyncpg.register_vector(connection)
+        except ImportError:
+            pass
+        except Exception as e:
+            logger.debug("Failed to register pgvector for asyncpg: %s", e)
 
     async def _close_pool(self) -> None:
         """Close the actual async connection pool."""
         if self.pool_instance:
             await self.pool_instance.close()
 
-    async def create_connection(self) -> AsyncpgConnection:
-        """Create a single async connection (not from pool).
+    async def create_connection(self) -> "AsyncpgConnection":
+        """Create a single async connection from the pool.
 
         Returns:
             An AsyncPG connection instance.
@@ -250,7 +155,7 @@ class AsyncpgConfig(AsyncDatabaseConfig[AsyncpgConnection, "Pool[Record]", Async
         return await self.pool_instance.acquire()
 
     @asynccontextmanager
-    async def provide_connection(self, *args: Any, **kwargs: Any) -> AsyncGenerator[AsyncpgConnection, None]:
+    async def provide_connection(self, *args: Any, **kwargs: Any) -> "AsyncGenerator[AsyncpgConnection, None]":
         """Provide an async connection context manager.
 
         Args:
@@ -271,28 +176,22 @@ class AsyncpgConfig(AsyncDatabaseConfig[AsyncpgConnection, "Pool[Record]", Async
                 await self.pool_instance.release(connection)
 
     @asynccontextmanager
-    async def provide_session(self, *args: Any, **kwargs: Any) -> AsyncGenerator[AsyncpgDriver, None]:
+    async def provide_session(
+        self, *args: Any, statement_config: "Optional[StatementConfig]" = None, **kwargs: Any
+    ) -> "AsyncGenerator[AsyncpgDriver, None]":
         """Provide an async driver session context manager.
 
         Args:
             *args: Additional arguments.
+            statement_config: Optional statement configuration override.
             **kwargs: Additional keyword arguments.
 
         Yields:
             An AsyncpgDriver instance.
         """
         async with self.provide_connection(*args, **kwargs) as connection:
-            statement_config = self.statement_config
-            # Inject parameter style info if not already set
-            if statement_config is not None and statement_config.allowed_parameter_styles is None:
-                from dataclasses import replace
-
-                statement_config = replace(
-                    statement_config,
-                    allowed_parameter_styles=self.supported_parameter_styles,
-                    default_parameter_style=self.default_parameter_style,
-                )
-            yield self.driver_type(connection=connection, config=statement_config)
+            final_statement_config = statement_config or self.statement_config or asyncpg_statement_config
+            yield self.driver_type(connection=connection, statement_config=final_statement_config)
 
     async def provide_pool(self, *args: Any, **kwargs: Any) -> "Pool[Record]":
         """Provide async pool instance.
@@ -313,6 +212,7 @@ class AsyncpgConfig(AsyncDatabaseConfig[AsyncpgConnection, "Pool[Record]", Async
         Returns:
             Dictionary mapping type names to types.
         """
+
         namespace = super().get_signature_namespace()
         namespace.update(
             {
@@ -322,7 +222,8 @@ class AsyncpgConfig(AsyncDatabaseConfig[AsyncpgConnection, "Pool[Record]", Async
                 "PoolConnectionProxyMeta": PoolConnectionProxyMeta,
                 "ConnectionMeta": ConnectionMeta,
                 "Record": Record,
-                "AsyncpgConnection": type(AsyncpgConnection),
+                "AsyncpgConnection": AsyncpgConnection,  # type: ignore[dict-item]
+                "AsyncpgCursor": AsyncpgCursor,
             }
         )
         return namespace