sqlspec 0.9.1__py3-none-any.whl → 0.10.1__py3-none-any.whl

sqlspec/adapters/psqlpy/driver.py

@@ -5,22 +5,25 @@ import logging
 import re
 from typing import TYPE_CHECKING, Any, Optional, Union, cast, overload
 
+from psqlpy import Connection, QueryResult
 from psqlpy.exceptions import RustPSQLDriverPyBaseError
 
-from sqlspec.base import AsyncDriverAdapterProtocol, T
+from sqlspec.base import AsyncDriverAdapterProtocol
 from sqlspec.exceptions import SQLParsingError
+from sqlspec.mixins import SQLTranslatorMixin
 from sqlspec.statement import PARAM_REGEX, SQLStatement
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
-    from psqlpy import Connection, QueryResult
+    from psqlpy import QueryResult
 
-    from sqlspec.typing import ModelDTOT, StatementParameterType
+    from sqlspec.typing import ModelDTOT, StatementParameterType, T
 
-__all__ = ("PsqlpyDriver",)
+__all__ = ("PsqlpyConnection", "PsqlpyDriver")
 
 
+PsqlpyConnection = Connection
 # Regex to find '?' placeholders, skipping those inside quotes or SQL comments
 QMARK_REGEX = re.compile(
     r"""(?P<dquote>"[^"]*") | # Double-quoted strings
@@ -33,13 +36,16 @@ QMARK_REGEX = re.compile(
 logger = logging.getLogger("sqlspec")
 
 
-class PsqlpyDriver(AsyncDriverAdapterProtocol["Connection"]):
+class PsqlpyDriver(
+    SQLTranslatorMixin["PsqlpyConnection"],
+    AsyncDriverAdapterProtocol["PsqlpyConnection"],
+):
     """Psqlpy Postgres Driver Adapter."""
 
-    connection: "Connection"
+    connection: "PsqlpyConnection"
     dialect: str = "postgres"
 
-    def __init__(self, connection: "Connection") -> None:
+    def __init__(self, connection: "PsqlpyConnection") -> None:
         self.connection = connection
 
     def _process_sql_params(
@@ -175,7 +181,7 @@ class PsqlpyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[PsqlpyConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "Sequence[dict[str, Any]]": ...
@@ -186,7 +192,7 @@ class PsqlpyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[PsqlpyConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
     ) -> "Sequence[ModelDTOT]": ...
@@ -196,7 +202,7 @@ class PsqlpyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["PsqlpyConnection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Sequence[Union[ModelDTOT, dict[str, Any]]]":
@@ -217,7 +223,7 @@ class PsqlpyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[PsqlpyConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "dict[str, Any]": ...
@@ -228,7 +234,7 @@ class PsqlpyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[PsqlpyConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
     ) -> "ModelDTOT": ...
@@ -238,7 +244,7 @@ class PsqlpyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["PsqlpyConnection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Union[ModelDTOT, dict[str, Any]]":
@@ -260,7 +266,7 @@ class PsqlpyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[PsqlpyConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "Optional[dict[str, Any]]": ...
@@ -271,7 +277,7 @@ class PsqlpyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[PsqlpyConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
     ) -> "Optional[ModelDTOT]": ...
@@ -281,7 +287,7 @@ class PsqlpyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["PsqlpyConnection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[ModelDTOT, dict[str, Any]]]":
@@ -307,7 +313,7 @@ class PsqlpyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[PsqlpyConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "Any": ...
@@ -318,7 +324,7 @@ class PsqlpyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[PsqlpyConnection]" = None,
         schema_type: "type[T]",
         **kwargs: Any,
     ) -> "T": ...
@@ -328,7 +334,7 @@ class PsqlpyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[PsqlpyConnection]" = None,
         schema_type: "Optional[type[T]]" = None,
         **kwargs: Any,
     ) -> "Union[T, Any]":
@@ -349,7 +355,7 @@ class PsqlpyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[PsqlpyConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "Optional[Any]": ...
@@ -360,7 +366,7 @@ class PsqlpyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[PsqlpyConnection]" = None,
         schema_type: "type[T]",
         **kwargs: Any,
     ) -> "Optional[T]": ...
@@ -370,7 +376,7 @@ class PsqlpyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[PsqlpyConnection]" = None,
         schema_type: "Optional[type[T]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[T, Any]]":
@@ -394,7 +400,7 @@ class PsqlpyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["PsqlpyConnection"] = None,
         **kwargs: Any,
     ) -> int:
         connection = self._connection(connection)
@@ -413,7 +419,7 @@ class PsqlpyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[PsqlpyConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "dict[str, Any]": ...
@@ -424,7 +430,7 @@ class PsqlpyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[PsqlpyConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
     ) -> "ModelDTOT": ...
@@ -434,7 +440,7 @@ class PsqlpyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["PsqlpyConnection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":
@@ -459,7 +465,7 @@ class PsqlpyDriver(AsyncDriverAdapterProtocol["Connection"]):
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["PsqlpyConnection"] = None,
         **kwargs: Any,
     ) -> str:
         connection = self._connection(connection)
@@ -469,7 +475,7 @@ class PsqlpyDriver(AsyncDriverAdapterProtocol["Connection"]):
         await connection.execute(sql, parameters=parameters)
         return sql
 
-    def _connection(self, connection: Optional["Connection"] = None) -> "Connection":
+    def _connection(self, connection: Optional["PsqlpyConnection"] = None) -> "PsqlpyConnection":
         """Get the connection to use.
 
         Args:

sqlspec/adapters/psycopg/__init__.py

@@ -4,13 +4,20 @@ from sqlspec.adapters.psycopg.config import (
     PsycopgSyncConfig,
     PsycopgSyncPoolConfig,
 )
-from sqlspec.adapters.psycopg.driver import PsycopgAsyncDriver, PsycopgSyncDriver
+from sqlspec.adapters.psycopg.driver import (
+    PsycopgAsyncConnection,
+    PsycopgAsyncDriver,
+    PsycopgSyncConnection,
+    PsycopgSyncDriver,
+)
 
 __all__ = (
     "PsycopgAsyncConfig",
+    "PsycopgAsyncConnection",
     "PsycopgAsyncDriver",
     "PsycopgAsyncPoolConfig",
     "PsycopgSyncConfig",
+    "PsycopgSyncConnection",
     "PsycopgSyncDriver",
     "PsycopgSyncPoolConfig",
 )

sqlspec/adapters/psycopg/config/__init__.py

@@ -1,9 +1,19 @@
 from sqlspec.adapters.psycopg.config._async import PsycopgAsyncConfig, PsycopgAsyncPoolConfig
 from sqlspec.adapters.psycopg.config._sync import PsycopgSyncConfig, PsycopgSyncPoolConfig
+from sqlspec.adapters.psycopg.driver import (
+    PsycopgAsyncConnection,
+    PsycopgAsyncDriver,
+    PsycopgSyncConnection,
+    PsycopgSyncDriver,
+)
 
 __all__ = (
     "PsycopgAsyncConfig",
+    "PsycopgAsyncConnection",
+    "PsycopgAsyncDriver",
     "PsycopgAsyncPoolConfig",
     "PsycopgSyncConfig",
+    "PsycopgSyncConnection",
+    "PsycopgSyncDriver",
     "PsycopgSyncPoolConfig",
 )

sqlspec/adapters/psycopg/config/_async.py

@@ -2,11 +2,10 @@ from contextlib import asynccontextmanager
 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.driver import PsycopgAsyncDriver
+from sqlspec.adapters.psycopg.driver import PsycopgAsyncConnection, PsycopgAsyncDriver
 from sqlspec.base import AsyncDatabaseConfig
 from sqlspec.exceptions import ImproperConfigurationError
 from sqlspec.typing import dataclass_to_dict
@@ -22,12 +21,12 @@ __all__ = (
 
 
 @dataclass
-class PsycopgAsyncPoolConfig(PsycopgGenericPoolConfig[AsyncConnection, AsyncConnectionPool]):
+class PsycopgAsyncPoolConfig(PsycopgGenericPoolConfig[PsycopgAsyncConnection, AsyncConnectionPool]):
     """Async Psycopg Pool Config"""
 
 
 @dataclass
-class PsycopgAsyncConfig(AsyncDatabaseConfig[AsyncConnection, AsyncConnectionPool, PsycopgAsyncDriver]):
+class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnectionPool, PsycopgAsyncDriver]):
     """Async Psycopg database Configuration.
 
     This class provides the base configuration for Psycopg database connections, extending
@@ -41,7 +40,7 @@ class PsycopgAsyncConfig(AsyncDatabaseConfig[AsyncConnection, AsyncConnectionPoo
     """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]
+    connection_type: "type[PsycopgAsyncConnection]" = field(init=False, default_factory=lambda: PsycopgAsyncConnection)  # 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"""
@@ -93,7 +92,7 @@ class PsycopgAsyncConfig(AsyncDatabaseConfig[AsyncConnection, AsyncConnectionPoo
         msg = "'pool_config' methods can not be used when a 'pool_instance' is provided."
         raise ImproperConfigurationError(msg)
 
-    async def create_connection(self) -> "AsyncConnection":
+    async def create_connection(self) -> "PsycopgAsyncConnection":
         """Create and return a new psycopg async connection from the pool.
 
         Returns:
@@ -143,7 +142,7 @@ class PsycopgAsyncConfig(AsyncDatabaseConfig[AsyncConnection, AsyncConnectionPoo
         return self.create_pool()
 
     @asynccontextmanager
-    async def provide_connection(self, *args: "Any", **kwargs: "Any") -> "AsyncGenerator[AsyncConnection, None]":
+    async def provide_connection(self, *args: "Any", **kwargs: "Any") -> "AsyncGenerator[PsycopgAsyncConnection, None]":
         """Create and provide a database connection.
 
         Yields:

sqlspec/adapters/psycopg/config/_sync.py

@@ -2,11 +2,10 @@ from contextlib import contextmanager
 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.driver import PsycopgSyncDriver
+from sqlspec.adapters.psycopg.driver import PsycopgSyncConnection, PsycopgSyncDriver
 from sqlspec.base import SyncDatabaseConfig
 from sqlspec.exceptions import ImproperConfigurationError
 from sqlspec.typing import dataclass_to_dict
@@ -22,12 +21,12 @@ __all__ = (
 
 
 @dataclass
-class PsycopgSyncPoolConfig(PsycopgGenericPoolConfig[Connection, ConnectionPool]):
+class PsycopgSyncPoolConfig(PsycopgGenericPoolConfig[PsycopgSyncConnection, ConnectionPool]):
     """Sync Psycopg Pool Config"""
 
 
 @dataclass
-class PsycopgSyncConfig(SyncDatabaseConfig[Connection, ConnectionPool, PsycopgSyncDriver]):
+class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, 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))
@@ -40,7 +39,7 @@ class PsycopgSyncConfig(SyncDatabaseConfig[Connection, ConnectionPool, PsycopgSy
     """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]
+    connection_type: "type[PsycopgSyncConnection]" = field(init=False, default_factory=lambda: PsycopgSyncConnection)  # 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"""
@@ -92,7 +91,7 @@ class PsycopgSyncConfig(SyncDatabaseConfig[Connection, ConnectionPool, PsycopgSy
         msg = "'pool_config' methods can not be used when a 'pool_instance' is provided."
         raise ImproperConfigurationError(msg)
 
-    def create_connection(self) -> "Connection":
+    def create_connection(self) -> "PsycopgSyncConnection":
         """Create and return a new psycopg connection from the pool.
 
         Returns:
@@ -142,11 +141,11 @@ class PsycopgSyncConfig(SyncDatabaseConfig[Connection, ConnectionPool, PsycopgSy
         return self.create_pool()
 
     @contextmanager
-    def provide_connection(self, *args: "Any", **kwargs: "Any") -> "Generator[Connection, None, None]":
+    def provide_connection(self, *args: "Any", **kwargs: "Any") -> "Generator[PsycopgSyncConnection, None, None]":
         """Create and provide a database connection.
 
         Yields:
-            Connection: A database connection from the pool.
+            PsycopgSyncConnection: A database connection from the pool.
         """
         pool = self.provide_pool(*args, **kwargs)
         with pool, pool.connection() as connection: