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

sqlspec/adapters/psqlpy/driver.py

@@ -3,22 +3,27 @@
 
 import logging
 import re
-from typing import TYPE_CHECKING, Any, Optional, Union, cast
+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 psqlpy import Connection, QueryResult
+    from collections.abc import Sequence
 
-    from sqlspec.typing import ModelDTOT, StatementParameterType
+    from psqlpy import QueryResult
 
-__all__ = ("PsqlpyDriver",)
+    from sqlspec.typing import ModelDTOT, StatementParameterType, T
 
+__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
@@ -31,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(
@@ -52,6 +60,17 @@ class PsqlpyDriver(AsyncDriverAdapterProtocol["Connection"]):
         psqlpy uses $1, $2 style parameters natively.
         This method converts '?' (tuple/list) and ':name' (dict) styles to $n.
         It relies on SQLStatement for initial parameter validation and merging.
+
+        Args:
+            sql: The SQL to process.
+            parameters: The parameters to process.
+            kwargs: Additional keyword arguments.
+
+        Raises:
+            SQLParsingError: If the SQL is invalid.
+
+        Returns:
+            A tuple of the processed SQL and parameters.
         """
         stmt = SQLStatement(sql=sql, parameters=parameters, dialect=self.dialect, kwargs=kwargs or None)
         sql, parameters = stmt.process()
@@ -154,16 +173,39 @@ class PsqlpyDriver(AsyncDriverAdapterProtocol["Connection"]):
 
         return sql, ()
 
+    # --- Public API Methods --- #
+    @overload
+    async def select(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsqlpyConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Sequence[dict[str, Any]]": ...
+    @overload
+    async def select(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsqlpyConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "Sequence[ModelDTOT]": ...
     async def select(
         self,
         sql: str,
         parameters: Optional["StatementParameterType"] = None,
         /,
         *,
-        connection: Optional["Connection"] = None,
+        connection: Optional["PsqlpyConnection"] = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
-    ) -> "list[Union[ModelDTOT, dict[str, Any]]]":
+    ) -> "Sequence[Union[ModelDTOT, dict[str, Any]]]":
         connection = self._connection(connection)
         sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
         parameters = parameters or []  # psqlpy expects a list/tuple
@@ -171,16 +213,38 @@ class PsqlpyDriver(AsyncDriverAdapterProtocol["Connection"]):
         results: QueryResult = await connection.fetch(sql, parameters=parameters)
 
         if schema_type is None:
-            return cast("list[dict[str, Any]]", results.result())  # type: ignore[return-value]
+            return cast("list[dict[str, Any]]", results.result())
         return results.as_class(as_class=schema_type)
 
+    @overload
+    async def select_one(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsqlpyConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "dict[str, Any]": ...
+    @overload
+    async def select_one(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsqlpyConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "ModelDTOT": ...
     async def select_one(
         self,
         sql: str,
         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]]":
@@ -195,13 +259,35 @@ class PsqlpyDriver(AsyncDriverAdapterProtocol["Connection"]):
             return cast("dict[str, Any]", result[0])  # type: ignore[index]
         return result.as_class(as_class=schema_type)[0]
 
+    @overload
+    async def select_one_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsqlpyConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Optional[dict[str, Any]]": ...
+    @overload
+    async def select_one_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsqlpyConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "Optional[ModelDTOT]": ...
     async def select_one_or_none(
         self,
         sql: str,
         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]]]":
@@ -220,13 +306,35 @@ class PsqlpyDriver(AsyncDriverAdapterProtocol["Connection"]):
             return None
         return cast("ModelDTOT", result[0])  # type: ignore[index]
 
+    @overload
+    async def select_value(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsqlpyConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Any": ...
+    @overload
     async def select_value(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[PsqlpyConnection]" = None,
+        schema_type: "type[T]",
+        **kwargs: Any,
+    ) -> "T": ...
+    async def select_value(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsqlpyConnection]" = None,
         schema_type: "Optional[type[T]]" = None,
         **kwargs: Any,
     ) -> "Union[T, Any]":
@@ -240,13 +348,35 @@ class PsqlpyDriver(AsyncDriverAdapterProtocol["Connection"]):
             return value
         return schema_type(value)  # type: ignore[call-arg]
 
+    @overload
     async def select_value_or_none(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[PsqlpyConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "Optional[Any]": ...
+    @overload
+    async def select_value_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsqlpyConnection]" = None,
+        schema_type: "type[T]",
+        **kwargs: Any,
+    ) -> "Optional[T]": ...
+    async def select_value_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsqlpyConnection]" = None,
         schema_type: "Optional[type[T]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[T, Any]]":
@@ -270,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)
@@ -282,13 +412,35 @@ class PsqlpyDriver(AsyncDriverAdapterProtocol["Connection"]):
         # if no error was raised
         return 1
 
+    @overload
+    async def insert_update_delete_returning(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsqlpyConnection]" = None,
+        schema_type: None = None,
+        **kwargs: Any,
+    ) -> "dict[str, Any]": ...
+    @overload
+    async def insert_update_delete_returning(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[PsqlpyConnection]" = None,
+        schema_type: "type[ModelDTOT]",
+        **kwargs: Any,
+    ) -> "ModelDTOT": ...
     async def insert_update_delete_returning(
         self,
         sql: str,
         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]]":
@@ -313,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)
@@ -323,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: