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

sqlspec/adapters/oracledb/driver.py

@@ -1,38 +1,39 @@
 from contextlib import asynccontextmanager, contextmanager
 from typing import TYPE_CHECKING, Any, Optional, Union, cast, overload
 
-from sqlspec.base import (
-    AsyncArrowBulkOperationsMixin,
-    AsyncDriverAdapterProtocol,
-    SyncArrowBulkOperationsMixin,
-    SyncDriverAdapterProtocol,
-    T,
-)
-from sqlspec.typing import ArrowTable, StatementParameterType
+from oracledb import AsyncConnection, AsyncCursor, Connection, Cursor
+
+from sqlspec.base import AsyncDriverAdapterProtocol, SyncDriverAdapterProtocol
+from sqlspec.mixins import AsyncArrowBulkOperationsMixin, SQLTranslatorMixin, SyncArrowBulkOperationsMixin
+from sqlspec.typing import ArrowTable, StatementParameterType, T
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Generator, Sequence
 
-    from oracledb import AsyncConnection, AsyncCursor, Connection, Cursor
-
-    # Conditionally import ArrowTable for type checking
     from sqlspec.typing import ModelDTOT
 
-__all__ = ("OracleAsyncDriver", "OracleSyncDriver")
+__all__ = ("OracleAsyncConnection", "OracleAsyncDriver", "OracleSyncConnection", "OracleSyncDriver")
 
+OracleSyncConnection = Connection
+OracleAsyncConnection = AsyncConnection
 
-class OracleSyncDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAdapterProtocol["Connection"]):
+
+class OracleSyncDriver(
+    SyncArrowBulkOperationsMixin["OracleSyncConnection"],
+    SQLTranslatorMixin["OracleSyncConnection"],
+    SyncDriverAdapterProtocol["OracleSyncConnection"],
+):
     """Oracle Sync Driver Adapter."""
 
-    connection: "Connection"
+    connection: "OracleSyncConnection"
     dialect: str = "oracle"
 
-    def __init__(self, connection: "Connection") -> None:
+    def __init__(self, connection: "OracleSyncConnection") -> None:
         self.connection = connection
 
     @staticmethod
     @contextmanager
-    def _with_cursor(connection: "Connection") -> "Generator[Cursor, None, None]":
+    def _with_cursor(connection: "OracleSyncConnection") -> "Generator[Cursor, None, None]":
         cursor = connection.cursor()
         try:
             yield cursor
@@ -47,7 +48,7 @@ class OracleSyncDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAda
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[OracleSyncConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "Sequence[dict[str, Any]]": ...
@@ -58,7 +59,7 @@ class OracleSyncDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAda
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[OracleSyncConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
     ) -> "Sequence[ModelDTOT]": ...
@@ -68,7 +69,7 @@ class OracleSyncDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAda
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[OracleSyncConnection]" = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Sequence[Union[ModelDTOT, dict[str, Any]]]":
@@ -106,7 +107,7 @@ class OracleSyncDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAda
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[OracleSyncConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "dict[str, Any]": ...
@@ -117,7 +118,7 @@ class OracleSyncDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAda
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[OracleSyncConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
     ) -> "ModelDTOT": ...
@@ -127,7 +128,7 @@ class OracleSyncDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAda
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[OracleSyncConnection]" = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Union[ModelDTOT, dict[str, Any]]":
@@ -166,7 +167,7 @@ class OracleSyncDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAda
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[OracleSyncConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "Optional[dict[str, Any]]": ...
@@ -177,7 +178,7 @@ class OracleSyncDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAda
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[OracleSyncConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
     ) -> "Optional[ModelDTOT]": ...
@@ -187,7 +188,7 @@ class OracleSyncDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAda
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[OracleSyncConnection]" = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[ModelDTOT, dict[str, Any]]]":
@@ -221,7 +222,7 @@ class OracleSyncDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAda
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[OracleSyncConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "Any": ...
@@ -232,7 +233,7 @@ class OracleSyncDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAda
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[OracleSyncConnection]" = None,
         schema_type: "type[T]",
         **kwargs: Any,
     ) -> "T": ...
@@ -242,7 +243,7 @@ class OracleSyncDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAda
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[OracleSyncConnection]" = None,
         schema_type: "Optional[type[T]]" = None,
         **kwargs: Any,
     ) -> "Union[T, Any]":
@@ -270,7 +271,7 @@ class OracleSyncDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAda
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[OracleSyncConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "Optional[Any]": ...
@@ -281,7 +282,7 @@ class OracleSyncDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAda
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[OracleSyncConnection]" = None,
         schema_type: "type[T]",
         **kwargs: Any,
     ) -> "Optional[T]": ...
@@ -291,7 +292,7 @@ class OracleSyncDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAda
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[OracleSyncConnection]" = None,
         schema_type: "Optional[type[T]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[T, Any]]":
@@ -320,7 +321,7 @@ class OracleSyncDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAda
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[OracleSyncConnection]" = None,
         **kwargs: Any,
     ) -> int:
         """Insert, update, or delete data from the database.
@@ -342,7 +343,7 @@ class OracleSyncDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAda
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[OracleSyncConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "dict[str, Any]": ...
@@ -353,7 +354,7 @@ class OracleSyncDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAda
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[OracleSyncConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
     ) -> "ModelDTOT": ...
@@ -363,7 +364,7 @@ class OracleSyncDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAda
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[OracleSyncConnection]" = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":
@@ -396,7 +397,7 @@ class OracleSyncDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAda
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[OracleSyncConnection]" = None,
         **kwargs: Any,
     ) -> str:
         """Execute a script.
@@ -417,7 +418,7 @@ class OracleSyncDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAda
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[Connection]" = None,
+        connection: "Optional[OracleSyncConnection]" = None,
         **kwargs: Any,
     ) -> "ArrowTable":  # pyright: ignore[reportUnknownVariableType]
         """Execute a SQL query and return results as an Apache Arrow Table.
@@ -433,19 +434,21 @@ class OracleSyncDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAda
 
 
 class OracleAsyncDriver(
-    AsyncArrowBulkOperationsMixin["AsyncConnection"], AsyncDriverAdapterProtocol["AsyncConnection"]
+    AsyncArrowBulkOperationsMixin["OracleAsyncConnection"],
+    SQLTranslatorMixin["OracleAsyncConnection"],
+    AsyncDriverAdapterProtocol["OracleAsyncConnection"],
 ):
     """Oracle Async Driver Adapter."""
 
-    connection: "AsyncConnection"
+    connection: "OracleAsyncConnection"
     dialect: str = "oracle"
 
-    def __init__(self, connection: "AsyncConnection") -> None:
+    def __init__(self, connection: "OracleAsyncConnection") -> None:
         self.connection = connection
 
     @staticmethod
     @asynccontextmanager
-    async def _with_cursor(connection: "AsyncConnection") -> "AsyncGenerator[AsyncCursor, None]":
+    async def _with_cursor(connection: "OracleAsyncConnection") -> "AsyncGenerator[AsyncCursor, None]":
         cursor = connection.cursor()
         try:
             yield cursor
@@ -460,7 +463,7 @@ class OracleAsyncDriver(
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[AsyncConnection]" = None,
+        connection: "Optional[OracleAsyncConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "Sequence[dict[str, Any]]": ...
@@ -471,7 +474,7 @@ class OracleAsyncDriver(
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[AsyncConnection]" = None,
+        connection: "Optional[OracleAsyncConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
     ) -> "Sequence[ModelDTOT]": ...
@@ -481,7 +484,7 @@ class OracleAsyncDriver(
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[AsyncConnection]" = None,
+        connection: "Optional[OracleAsyncConnection]" = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Sequence[Union[ModelDTOT, dict[str, Any]]]":
@@ -513,7 +516,7 @@ class OracleAsyncDriver(
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[AsyncConnection]" = None,
+        connection: "Optional[OracleAsyncConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "dict[str, Any]": ...
@@ -524,7 +527,7 @@ class OracleAsyncDriver(
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[AsyncConnection]" = None,
+        connection: "Optional[OracleAsyncConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
     ) -> "ModelDTOT": ...
@@ -534,7 +537,7 @@ class OracleAsyncDriver(
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[AsyncConnection]" = None,
+        connection: "Optional[OracleAsyncConnection]" = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Union[ModelDTOT, dict[str, Any]]":
@@ -565,7 +568,7 @@ class OracleAsyncDriver(
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[AsyncConnection]" = None,
+        connection: "Optional[OracleAsyncConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "Optional[dict[str, Any]]": ...
@@ -576,7 +579,7 @@ class OracleAsyncDriver(
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[AsyncConnection]" = None,
+        connection: "Optional[OracleAsyncConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
     ) -> "Optional[ModelDTOT]": ...
@@ -586,7 +589,7 @@ class OracleAsyncDriver(
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[AsyncConnection]" = None,
+        connection: "Optional[OracleAsyncConnection]" = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[ModelDTOT, dict[str, Any]]]":
@@ -620,7 +623,7 @@ class OracleAsyncDriver(
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[AsyncConnection]" = None,
+        connection: "Optional[OracleAsyncConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "Any": ...
@@ -631,7 +634,7 @@ class OracleAsyncDriver(
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[AsyncConnection]" = None,
+        connection: "Optional[OracleAsyncConnection]" = None,
         schema_type: "type[T]",
         **kwargs: Any,
     ) -> "T": ...
@@ -641,7 +644,7 @@ class OracleAsyncDriver(
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[AsyncConnection]" = None,
+        connection: "Optional[OracleAsyncConnection]" = None,
         schema_type: "Optional[type[T]]" = None,
         **kwargs: Any,
     ) -> "Union[T, Any]":
@@ -669,7 +672,7 @@ class OracleAsyncDriver(
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[AsyncConnection]" = None,
+        connection: "Optional[OracleAsyncConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "Optional[Any]": ...
@@ -680,7 +683,7 @@ class OracleAsyncDriver(
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[AsyncConnection]" = None,
+        connection: "Optional[OracleAsyncConnection]" = None,
         schema_type: "type[T]",
         **kwargs: Any,
     ) -> "Optional[T]": ...
@@ -690,7 +693,7 @@ class OracleAsyncDriver(
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[AsyncConnection]" = None,
+        connection: "Optional[OracleAsyncConnection]" = None,
         schema_type: "Optional[type[T]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[T, Any]]":
@@ -719,7 +722,7 @@ class OracleAsyncDriver(
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[AsyncConnection]" = None,
+        connection: "Optional[OracleAsyncConnection]" = None,
         **kwargs: Any,
     ) -> int:
         """Insert, update, or delete data from the database.
@@ -741,7 +744,7 @@ class OracleAsyncDriver(
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[AsyncConnection]" = None,
+        connection: "Optional[OracleAsyncConnection]" = None,
         schema_type: None = None,
         **kwargs: Any,
     ) -> "dict[str, Any]": ...
@@ -752,7 +755,7 @@ class OracleAsyncDriver(
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[AsyncConnection]" = None,
+        connection: "Optional[OracleAsyncConnection]" = None,
         schema_type: "type[ModelDTOT]",
         **kwargs: Any,
     ) -> "ModelDTOT": ...
@@ -762,7 +765,7 @@ class OracleAsyncDriver(
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[AsyncConnection]" = None,
+        connection: "Optional[OracleAsyncConnection]" = None,
         schema_type: "Optional[type[ModelDTOT]]" = None,
         **kwargs: Any,
     ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":
@@ -795,7 +798,7 @@ class OracleAsyncDriver(
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[AsyncConnection]" = None,
+        connection: "Optional[OracleAsyncConnection]" = None,
         **kwargs: Any,
     ) -> str:
         """Execute a script.
@@ -816,7 +819,7 @@ class OracleAsyncDriver(
         parameters: "Optional[StatementParameterType]" = None,
         /,
         *,
-        connection: "Optional[AsyncConnection]" = None,
+        connection: "Optional[OracleAsyncConnection]" = None,
         **kwargs: Any,
     ) -> "ArrowTable":  # pyright: ignore[reportUnknownVariableType]
         """Execute a SQL query asynchronously and return results as an Apache Arrow Table.

sqlspec/adapters/psqlpy/__init__.py

@@ -0,0 +1,9 @@
+from sqlspec.adapters.psqlpy.config import PsqlpyConfig, PsqlpyPoolConfig
+from sqlspec.adapters.psqlpy.driver import PsqlpyConnection, PsqlpyDriver
+
+__all__ = (
+    "PsqlpyConfig",
+    "PsqlpyConnection",
+    "PsqlpyDriver",
+    "PsqlpyPoolConfig",
+)

sqlspec/adapters/psqlpy/config.py

@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Any, Optional, Union
 
 from psqlpy import Connection, ConnectionPool
 
-from sqlspec.adapters.psqlpy.driver import PsqlpyDriver
+from sqlspec.adapters.psqlpy.driver import PsqlpyConnection, PsqlpyDriver
 from sqlspec.base import AsyncDatabaseConfig, GenericPoolConfig
 from sqlspec.exceptions import ImproperConfigurationError
 from sqlspec.typing import Empty, EmptyType, dataclass_to_dict
@@ -94,7 +94,7 @@ class PsqlpyPoolConfig(GenericPoolConfig):
 
 
 @dataclass
-class PsqlpyConfig(AsyncDatabaseConfig[Connection, ConnectionPool, PsqlpyDriver]):
+class PsqlpyConfig(AsyncDatabaseConfig[PsqlpyConnection, ConnectionPool, PsqlpyDriver]):
     """Configuration for psqlpy database connections, managing a connection pool.
 
     This configuration class wraps `PsqlpyPoolConfig` and manages the lifecycle
@@ -105,7 +105,7 @@ class PsqlpyConfig(AsyncDatabaseConfig[Connection, ConnectionPool, PsqlpyDriver]
     """Psqlpy Pool configuration"""
     driver_type: type[PsqlpyDriver] = field(default=PsqlpyDriver, init=False, hash=False)
     """Type of the driver object"""
-    connection_type: type[Connection] = field(default=Connection, init=False, hash=False)
+    connection_type: type[PsqlpyConnection] = field(default=PsqlpyConnection, init=False, hash=False)
     """Type of the connection object"""
     pool_instance: Optional[ConnectionPool] = field(default=None, hash=False)
     """The connection pool instance. If set, this will be used instead of creating a new pool."""
@@ -204,7 +204,7 @@ class PsqlpyConfig(AsyncDatabaseConfig[Connection, ConnectionPool, PsqlpyDriver]
 
         return _create()
 
-    def create_connection(self) -> "Awaitable[Connection]":
+    def create_connection(self) -> "Awaitable[PsqlpyConnection]":
         """Create and return a new, standalone psqlpy connection using the configured parameters.
 
         Returns:
@@ -222,7 +222,7 @@ class PsqlpyConfig(AsyncDatabaseConfig[Connection, ConnectionPool, PsqlpyDriver]
         return _create()
 
     @asynccontextmanager
-    async def provide_connection(self, *args: "Any", **kwargs: "Any") -> "AsyncGenerator[Connection, None]":
+    async def provide_connection(self, *args: "Any", **kwargs: "Any") -> "AsyncGenerator[PsqlpyConnection, None]":
         """Acquire a connection from the pool.
 
         Yields: