sqlspec 0.7.1__py3-none-any.whl → 0.9.0__py3-none-any.whl

sqlspec/adapters/psycopg/config/_async.py

@@ -1,33 +1,33 @@
 from contextlib import asynccontextmanager
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Optional
 
 from psycopg import AsyncConnection
 from psycopg_pool import AsyncConnectionPool
 
-from sqlspec.adapters.psycopg.config._common import PsycoPgGenericPoolConfig
+from sqlspec.adapters.psycopg.config._common import PsycopgGenericPoolConfig
+from sqlspec.adapters.psycopg.driver import PsycopgAsyncDriver
 from sqlspec.base import AsyncDatabaseConfig
 from sqlspec.exceptions import ImproperConfigurationError
 from sqlspec.typing import dataclass_to_dict
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Awaitable
-    from typing import Any
 
 
 __all__ = (
-    "PsycoPgAsyncDatabaseConfig",
-    "PsycoPgAsyncPoolConfig",
+    "PsycopgAsyncConfig",
+    "PsycopgAsyncPoolConfig",
 )
 
 
 @dataclass
-class PsycoPgAsyncPoolConfig(PsycoPgGenericPoolConfig[AsyncConnection, AsyncConnectionPool]):
+class PsycopgAsyncPoolConfig(PsycopgGenericPoolConfig[AsyncConnection, AsyncConnectionPool]):
     """Async Psycopg Pool Config"""
 
 
 @dataclass
-class PsycoPgAsyncDatabaseConfig(AsyncDatabaseConfig[AsyncConnection, AsyncConnectionPool]):
+class PsycopgAsyncConfig(AsyncDatabaseConfig[AsyncConnection, AsyncConnectionPool, PsycopgAsyncDriver]):
     """Async Psycopg database Configuration.
 
     This class provides the base configuration for Psycopg database connections, extending
@@ -37,21 +37,88 @@ class PsycoPgAsyncDatabaseConfig(AsyncDatabaseConfig[AsyncConnection, AsyncConne
     with both synchronous and asynchronous connections.([2](https://www.psycopg.org/psycopg3/docs/api/connections.html))
     """
 
-    pool_config: "Optional[PsycoPgAsyncPoolConfig]" = None
+    pool_config: "Optional[PsycopgAsyncPoolConfig]" = None
     """Psycopg Pool configuration"""
     pool_instance: "Optional[AsyncConnectionPool]" = None
     """Optional pool to use"""
+    connection_type: "type[AsyncConnection]" = field(init=False, default_factory=lambda: AsyncConnection)  # type: ignore[assignment]
+    """Type of the connection object"""
+    driver_type: "type[PsycopgAsyncDriver]" = field(init=False, default_factory=lambda: PsycopgAsyncDriver)  # type: ignore[type-abstract,unused-ignore]
+    """Type of the driver object"""
+
+    @property
+    def connection_config_dict(self) -> "dict[str, Any]":
+        """Return the connection configuration as a dict.
+
+        Returns:
+            A string keyed dict of config kwargs for the psycopg.connect function.
+
+        Raises:
+            ImproperConfigurationError: If the connection configuration is not provided.
+        """
+        if self.pool_config:
+            # Filter out pool-specific parameters
+            pool_only_params = {
+                "min_size",
+                "max_size",
+                "name",
+                "timeout",
+                "reconnect_timeout",
+                "max_idle",
+                "max_lifetime",
+            }
+            return dataclass_to_dict(
+                self.pool_config,
+                exclude_empty=True,
+                convert_nested=False,
+                exclude=pool_only_params.union({"pool_instance", "connection_type", "driver_type", "open"}),
+            )
+        msg = "You must provide a 'pool_config' for this adapter."
+        raise ImproperConfigurationError(msg)
 
     @property
     def pool_config_dict(self) -> "dict[str, Any]":
-        """Return the pool configuration as a dict."""
+        """Return the pool configuration as a dict.
+
+        Raises:
+            ImproperConfigurationError: If pool_config is not set but pool_instance is provided.
+        """
         if self.pool_config:
-            return dataclass_to_dict(self.pool_config, exclude_empty=True, convert_nested=False)
+            return dataclass_to_dict(
+                self.pool_config,
+                exclude_empty=True,
+                convert_nested=False,
+                exclude={"pool_instance", "connection_type", "driver_type"},
+            )
         msg = "'pool_config' methods can not be used when a 'pool_instance' is provided."
         raise ImproperConfigurationError(msg)
 
+    async def create_connection(self) -> "AsyncConnection":
+        """Create and return a new psycopg async connection from the pool.
+
+        Returns:
+            An AsyncConnection instance.
+
+        Raises:
+            ImproperConfigurationError: If the connection could not be created.
+        """
+        try:
+            pool = await self.provide_pool()
+            return await pool.getconn()
+        except Exception as e:
+            msg = f"Could not configure the Psycopg connection. Error: {e!s}"
+            raise ImproperConfigurationError(msg) from e
+
     async def create_pool(self) -> "AsyncConnectionPool":
-        """Create and return a connection pool."""
+        """Create and return a connection pool.
+
+        Returns:
+            AsyncConnectionPool: The configured connection pool.
+
+        Raises:
+            ImproperConfigurationError: If neither pool_config nor pool_instance are provided
+                or if pool creation fails.
+        """
         if self.pool_instance is not None:
             return self.pool_instance
 
@@ -60,19 +127,44 @@ class PsycoPgAsyncDatabaseConfig(AsyncDatabaseConfig[AsyncConnection, AsyncConne
             raise ImproperConfigurationError(msg)
 
         pool_config = self.pool_config_dict
-        self.pool_instance = AsyncConnectionPool(**pool_config)
+        self.pool_instance = AsyncConnectionPool(open=False, **pool_config)
         if self.pool_instance is None:  # pyright: ignore[reportUnnecessaryComparison]
             msg = "Could not configure the 'pool_instance'. Please check your configuration."  # type: ignore[unreachable]
             raise ImproperConfigurationError(msg)
+        await self.pool_instance.open()
         return self.pool_instance
 
     def provide_pool(self, *args: "Any", **kwargs: "Any") -> "Awaitable[AsyncConnectionPool]":
-        """Create and return a connection pool."""
+        """Create and return a connection pool.
+
+        Returns:
+            Awaitable[AsyncConnectionPool]: The configured connection pool.
+        """
         return self.create_pool()
 
     @asynccontextmanager
     async def provide_connection(self, *args: "Any", **kwargs: "Any") -> "AsyncGenerator[AsyncConnection, None]":
-        """Create and provide a database connection."""
+        """Create and provide a database connection.
+
+        Yields:
+            AsyncConnection: A database connection from the pool.
+        """
         pool = await self.provide_pool(*args, **kwargs)
-        async with pool.connection() as connection:
+        async with pool, pool.connection() as connection:
             yield connection
+
+    @asynccontextmanager
+    async def provide_session(self, *args: "Any", **kwargs: "Any") -> "AsyncGenerator[PsycopgAsyncDriver, None]":
+        """Create and provide a database session.
+
+        Yields:
+            PsycopgAsyncDriver: A driver instance with an active connection.
+        """
+        async with self.provide_connection(*args, **kwargs) as connection:
+            yield self.driver_type(connection)
+
+    async def close_pool(self) -> None:
+        """Close the connection pool."""
+        if self.pool_instance is not None:
+            await self.pool_instance.close()
+            self.pool_instance = None

sqlspec/adapters/psycopg/config/_common.py

@@ -14,7 +14,7 @@ if TYPE_CHECKING:
     from sqlspec.typing import EmptyType
 
 
-__all__ = ("PsycoPgGenericPoolConfig",)
+__all__ = ("PsycopgGenericPoolConfig",)
 
 
 ConnectionT = TypeVar("ConnectionT", bound="Union[Connection, AsyncConnection]")
@@ -22,7 +22,7 @@ PoolT = TypeVar("PoolT", bound="Union[ConnectionPool, AsyncConnectionPool]")
 
 
 @dataclass
-class PsycoPgGenericPoolConfig(Generic[ConnectionT, PoolT], GenericPoolConfig):
+class PsycopgGenericPoolConfig(GenericPoolConfig, Generic[ConnectionT, PoolT]):
     """Configuration for Psycopg connection pools.
 
     This class provides configuration options for both synchronous and asynchronous Psycopg

sqlspec/adapters/psycopg/config/_sync.py

@@ -1,33 +1,33 @@
 from contextlib import contextmanager
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Optional
 
 from psycopg import Connection
 from psycopg_pool import ConnectionPool
 
-from sqlspec.adapters.psycopg.config._common import PsycoPgGenericPoolConfig
+from sqlspec.adapters.psycopg.config._common import PsycopgGenericPoolConfig
+from sqlspec.adapters.psycopg.driver import PsycopgSyncDriver
 from sqlspec.base import SyncDatabaseConfig
 from sqlspec.exceptions import ImproperConfigurationError
 from sqlspec.typing import dataclass_to_dict
 
 if TYPE_CHECKING:
     from collections.abc import Generator
-    from typing import Any
 
 
 __all__ = (
-    "PsycoPgSyncDatabaseConfig",
-    "PsycoPgSyncPoolConfig",
+    "PsycopgSyncConfig",
+    "PsycopgSyncPoolConfig",
 )
 
 
 @dataclass
-class PsycoPgSyncPoolConfig(PsycoPgGenericPoolConfig[Connection, ConnectionPool]):
+class PsycopgSyncPoolConfig(PsycopgGenericPoolConfig[Connection, ConnectionPool]):
     """Sync Psycopg Pool Config"""
 
 
 @dataclass
-class PsycoPgSyncDatabaseConfig(SyncDatabaseConfig[Connection, ConnectionPool]):
+class PsycopgSyncConfig(SyncDatabaseConfig[Connection, ConnectionPool, PsycopgSyncDriver]):
     """Sync Psycopg database Configuration.
     This class provides the base configuration for Psycopg database connections, extending
     the generic database configuration with Psycopg-specific settings.([1](https://www.psycopg.org/psycopg3/docs/api/connections.html))
@@ -36,21 +36,88 @@ class PsycoPgSyncDatabaseConfig(SyncDatabaseConfig[Connection, ConnectionPool]):
     with both synchronous and asynchronous connections.([2](https://www.psycopg.org/psycopg3/docs/api/connections.html))
     """
 
-    pool_config: "Optional[PsycoPgSyncPoolConfig]" = None
+    pool_config: "Optional[PsycopgSyncPoolConfig]" = None
     """Psycopg Pool configuration"""
     pool_instance: "Optional[ConnectionPool]" = None
     """Optional pool to use"""
+    connection_type: "type[Connection]" = field(init=False, default_factory=lambda: Connection)  # type: ignore[assignment]
+    """Type of the connection object"""
+    driver_type: "type[PsycopgSyncDriver]" = field(init=False, default_factory=lambda: PsycopgSyncDriver)  # type: ignore[type-abstract,unused-ignore]
+    """Type of the driver object"""
+
+    @property
+    def connection_config_dict(self) -> "dict[str, Any]":
+        """Return the connection configuration as a dict.
+
+        Returns:
+            A string keyed dict of config kwargs for the psycopg.connect function.
+
+        Raises:
+            ImproperConfigurationError: If the connection configuration is not provided.
+        """
+        if self.pool_config:
+            # Filter out pool-specific parameters
+            pool_only_params = {
+                "min_size",
+                "max_size",
+                "name",
+                "timeout",
+                "reconnect_timeout",
+                "max_idle",
+                "max_lifetime",
+            }
+            return dataclass_to_dict(
+                self.pool_config,
+                exclude_empty=True,
+                convert_nested=False,
+                exclude=pool_only_params.union({"pool_instance", "connection_type", "driver_type", "open"}),
+            )
+        msg = "You must provide a 'pool_config' for this adapter."
+        raise ImproperConfigurationError(msg)
 
     @property
     def pool_config_dict(self) -> "dict[str, Any]":
-        """Return the pool configuration as a dict."""
+        """Return the pool configuration as a dict.
+
+        Raises:
+            ImproperConfigurationError: If pool_config is not provided and instead pool_instance is used.
+        """
         if self.pool_config:
-            return dataclass_to_dict(self.pool_config, exclude_empty=True, convert_nested=False)
+            return dataclass_to_dict(
+                self.pool_config,
+                exclude_empty=True,
+                convert_nested=False,
+                exclude={"pool_instance", "connection_type", "driver_type", "open"},
+            )
         msg = "'pool_config' methods can not be used when a 'pool_instance' is provided."
         raise ImproperConfigurationError(msg)
 
+    def create_connection(self) -> "Connection":
+        """Create and return a new psycopg connection from the pool.
+
+        Returns:
+            A Connection instance.
+
+        Raises:
+            ImproperConfigurationError: If the connection could not be created.
+        """
+        try:
+            pool = self.provide_pool()
+            return pool.getconn()
+        except Exception as e:
+            msg = f"Could not configure the Psycopg connection. Error: {e!s}"
+            raise ImproperConfigurationError(msg) from e
+
     def create_pool(self) -> "ConnectionPool":
-        """Create and return a connection pool."""
+        """Create and return a connection pool.
+
+        Returns:
+            ConnectionPool: The configured connection pool instance.
+
+        Raises:
+            ImproperConfigurationError: If neither pool_config nor pool_instance is provided,
+                or if the pool could not be configured.
+        """
         if self.pool_instance is not None:
             return self.pool_instance
 
@@ -59,19 +126,44 @@ class PsycoPgSyncDatabaseConfig(SyncDatabaseConfig[Connection, ConnectionPool]):
             raise ImproperConfigurationError(msg)
 
         pool_config = self.pool_config_dict
-        self.pool_instance = ConnectionPool(**pool_config)
+        self.pool_instance = ConnectionPool(open=False, **pool_config)
         if self.pool_instance is None:  # pyright: ignore[reportUnnecessaryComparison]
             msg = "Could not configure the 'pool_instance'. Please check your configuration."  # type: ignore[unreachable]
             raise ImproperConfigurationError(msg)
+        self.pool_instance.open()
         return self.pool_instance
 
     def provide_pool(self, *args: "Any", **kwargs: "Any") -> "ConnectionPool":
-        """Create and return a connection pool."""
+        """Create and return a connection pool.
+
+        Returns:
+            ConnectionPool: The configured connection pool instance.
+        """
         return self.create_pool()
 
     @contextmanager
     def provide_connection(self, *args: "Any", **kwargs: "Any") -> "Generator[Connection, None, None]":
-        """Create and provide a database connection."""
+        """Create and provide a database connection.
+
+        Yields:
+            Connection: A database connection from the pool.
+        """
         pool = self.provide_pool(*args, **kwargs)
-        with pool.connection() as connection:
+        with pool, pool.connection() as connection:
             yield connection
+
+    @contextmanager
+    def provide_session(self, *args: "Any", **kwargs: "Any") -> "Generator[PsycopgSyncDriver, None, None]":
+        """Create and provide a database session.
+
+        Yields:
+            PsycopgSyncDriver: A driver instance with an active connection.
+        """
+        with self.provide_connection(*args, **kwargs) as connection:
+            yield self.driver_type(connection)
+
+    def close_pool(self) -> None:
+        """Close the connection pool."""
+        if self.pool_instance is not None:
+            self.pool_instance.close()
+            self.pool_instance = None