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

sqlspec/adapters/oracledb/driver.py

@@ -0,0 +1,571 @@
+from contextlib import asynccontextmanager, contextmanager
+from typing import TYPE_CHECKING, Any, Optional, Union, cast
+
+from sqlspec.base import (
+    AsyncArrowBulkOperationsMixin,
+    AsyncDriverAdapterProtocol,
+    SyncArrowBulkOperationsMixin,
+    SyncDriverAdapterProtocol,
+    T,
+)
+from sqlspec.typing import ArrowTable, StatementParameterType
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator, Generator
+
+    from oracledb import AsyncConnection, AsyncCursor, Connection, Cursor
+
+    # Conditionally import ArrowTable for type checking
+    from sqlspec.typing import ModelDTOT
+
+__all__ = ("OracleAsyncDriver", "OracleSyncDriver")
+
+
+class OracleSyncDriver(SyncArrowBulkOperationsMixin["Connection"], SyncDriverAdapterProtocol["Connection"]):
+    """Oracle Sync Driver Adapter."""
+
+    connection: "Connection"
+    dialect: str = "oracle"
+
+    def __init__(self, connection: "Connection") -> None:
+        self.connection = connection
+
+    @staticmethod
+    @contextmanager
+    def _with_cursor(connection: "Connection") -> "Generator[Cursor, None, None]":
+        cursor = connection.cursor()
+        try:
+            yield cursor
+        finally:
+            cursor.close()
+
+    def select(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[Connection]" = None,
+        schema_type: "Optional[type[ModelDTOT]]" = None,
+        **kwargs: Any,
+    ) -> "list[Union[ModelDTOT, dict[str, Any]]]":
+        """Fetch data from the database.
+
+        Args:
+            sql: The SQL query string.
+            parameters: The parameters for the query (dict, tuple, list, or None).
+            connection: Optional connection override.
+            schema_type: Optional schema class for the result.
+            **kwargs: Additional keyword arguments to merge with parameters if parameters is a dict.
+
+        Returns:
+            List of row data as either model instances or dictionaries.
+        """
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+        with self._with_cursor(connection) as cursor:
+            cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
+            results = cursor.fetchall()  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+            if not results:
+                return []
+            # Get column names
+            column_names = [col[0] for col in cursor.description or []]  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+
+            if schema_type:
+                return [cast("ModelDTOT", schema_type(**dict(zip(column_names, row)))) for row in results]  # pyright: ignore
+
+            return [dict(zip(column_names, row)) for row in results]  # pyright: ignore
+
+    def select_one(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[Connection]" = None,
+        schema_type: "Optional[type[ModelDTOT]]" = None,
+        **kwargs: Any,
+    ) -> "Union[ModelDTOT, dict[str, Any]]":
+        """Fetch one row from the database.
+
+        Args:
+            sql: The SQL query string.
+            parameters: The parameters for the query (dict, tuple, list, or None).
+            connection: Optional connection override.
+            schema_type: Optional schema class for the result.
+            **kwargs: Additional keyword arguments to merge with parameters if parameters is a dict.
+
+        Returns:
+            The first row of the query results.
+        """
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+
+        with self._with_cursor(connection) as cursor:
+            cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
+            result = cursor.fetchone()  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+            result = self.check_not_found(result)  # pyright: ignore[reportUnknownArgumentType]
+
+            # Get column names
+            column_names = [col[0] for col in cursor.description or []]  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+
+            if schema_type is not None:
+                return cast("ModelDTOT", schema_type(**dict(zip(column_names, result))))  # pyright: ignore[reportUnknownArgumentType]
+            # Always return dictionaries
+            return dict(zip(column_names, result))  # pyright: ignore[reportUnknownArgumentType,reportUnknownVariableType]
+
+    def select_one_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[Connection]" = None,
+        schema_type: "Optional[type[ModelDTOT]]" = None,
+        **kwargs: Any,
+    ) -> "Optional[Union[ModelDTOT, dict[str, Any]]]":
+        """Fetch one row from the database.
+
+        Returns:
+            The first row of the query results.
+        """
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+
+        with self._with_cursor(connection) as cursor:
+            cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
+            result = cursor.fetchone()  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+
+            if result is None:
+                return None
+
+            # Get column names
+            column_names = [col[0] for col in cursor.description or []]  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+
+            if schema_type is not None:
+                return cast("ModelDTOT", schema_type(**dict(zip(column_names, result))))  # pyright: ignore[reportUnknownArgumentType]
+            # Always return dictionaries
+            return dict(zip(column_names, result))  # pyright: ignore[reportUnknownArgumentType,reportUnknownVariableType]
+
+    def select_value(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[Connection]" = None,
+        schema_type: "Optional[type[T]]" = None,
+        **kwargs: Any,
+    ) -> "Union[T, Any]":
+        """Fetch a single value from the database.
+
+        Returns:
+            The first value from the first row of results, or None if no results.
+        """
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+
+        with self._with_cursor(connection) as cursor:
+            cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
+            result = cursor.fetchone()  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+            result = self.check_not_found(result)  # pyright: ignore[reportUnknownArgumentType]
+
+            if schema_type is None:
+                return result[0]  # pyright: ignore[reportUnknownArgumentType]
+            return schema_type(result[0])  # type: ignore[call-arg]
+
+    def select_value_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[Connection]" = None,
+        schema_type: "Optional[type[T]]" = None,
+        **kwargs: Any,
+    ) -> "Optional[Union[T, Any]]":
+        """Fetch a single value from the database.
+
+        Returns:
+            The first value from the first row of results, or None if no results.
+        """
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+
+        with self._with_cursor(connection) as cursor:
+            cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
+            result = cursor.fetchone()  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+
+            if result is None:
+                return None
+
+            if schema_type is None:
+                return result[0]  # pyright: ignore[reportUnknownArgumentType]
+            return schema_type(result[0])  # type: ignore[call-arg]
+
+    def insert_update_delete(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[Connection]" = None,
+        **kwargs: Any,
+    ) -> int:
+        """Insert, update, or delete data from the database.
+
+        Returns:
+            Row count affected by the operation.
+        """
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+
+        with self._with_cursor(connection) as cursor:
+            cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
+            return cursor.rowcount  # pyright: ignore[reportUnknownMemberType]
+
+    def insert_update_delete_returning(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[Connection]" = None,
+        schema_type: "Optional[type[ModelDTOT]]" = None,
+        **kwargs: Any,
+    ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":
+        """Insert, update, or delete data from the database and return result.
+
+        Returns:
+            The first row of results.
+        """
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+
+        with self._with_cursor(connection) as cursor:
+            cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
+            result = cursor.fetchone()  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+
+            if result is None:
+                return None
+
+            # Get column names
+            column_names = [col[0] for col in cursor.description or []]  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+
+            if schema_type is not None:
+                return cast("ModelDTOT", schema_type(**dict(zip(column_names, result))))  # pyright: ignore[reportUnknownArgumentType]
+            # Always return dictionaries
+            return dict(zip(column_names, result))  # pyright: ignore[reportUnknownArgumentType,reportUnknownVariableType]
+
+    def execute_script(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[Connection]" = None,
+        **kwargs: Any,
+    ) -> str:
+        """Execute a script.
+
+        Returns:
+            Status message for the operation.
+        """
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+
+        with self._with_cursor(connection) as cursor:
+            cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
+            return str(cursor.rowcount)  # pyright: ignore[reportUnknownMemberType]
+
+    def select_arrow(  # pyright: ignore[reportUnknownParameterType]
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[Connection]" = None,
+        **kwargs: Any,
+    ) -> "ArrowTable":  # pyright: ignore[reportUnknownVariableType]
+        """Execute a SQL query and return results as an Apache Arrow Table.
+
+        Returns:
+            An Apache Arrow Table containing the query results.
+        """
+
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+        results = connection.fetch_df_all(sql, parameters)
+        return cast("ArrowTable", ArrowTable.from_arrays(arrays=results.column_arrays(), names=results.column_names()))  # pyright: ignore
+
+
+class OracleAsyncDriver(
+    AsyncArrowBulkOperationsMixin["AsyncConnection"], AsyncDriverAdapterProtocol["AsyncConnection"]
+):
+    """Oracle Async Driver Adapter."""
+
+    connection: "AsyncConnection"
+    dialect: str = "oracle"
+
+    def __init__(self, connection: "AsyncConnection") -> None:
+        self.connection = connection
+
+    @staticmethod
+    @asynccontextmanager
+    async def _with_cursor(connection: "AsyncConnection") -> "AsyncGenerator[AsyncCursor, None]":
+        cursor = connection.cursor()
+        try:
+            yield cursor
+        finally:
+            cursor.close()
+
+    async def select(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncConnection]" = None,
+        schema_type: "Optional[type[ModelDTOT]]" = None,
+        **kwargs: Any,
+    ) -> "list[Union[ModelDTOT, dict[str, Any]]]":
+        """Fetch data from the database.
+
+        Returns:
+            List of row data as either model instances or dictionaries.
+        """
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+
+        async with self._with_cursor(connection) as cursor:
+            await cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
+            results = await cursor.fetchall()  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+            if not results:
+                return []
+            # Get column names
+            column_names = [col[0] for col in cursor.description or []]  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+
+            if schema_type:
+                return [cast("ModelDTOT", schema_type(**dict(zip(column_names, row)))) for row in results]  # pyright: ignore
+
+            return [dict(zip(column_names, row)) for row in results]  # pyright: ignore
+
+    async def select_one(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncConnection]" = None,
+        schema_type: "Optional[type[ModelDTOT]]" = None,
+        **kwargs: Any,
+    ) -> "Union[ModelDTOT, dict[str, Any]]":
+        """Fetch one row from the database.
+
+        Returns:
+            The first row of the query results.
+        """
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+
+        async with self._with_cursor(connection) as cursor:
+            await cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
+            result = await cursor.fetchone()  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+            result = self.check_not_found(result)  # pyright: ignore[reportUnknownArgumentType]
+            # Get column names
+            column_names = [col[0] for col in cursor.description or []]  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+
+            if schema_type is not None:
+                return cast("ModelDTOT", schema_type(**dict(zip(column_names, result))))  # pyright: ignore[reportUnknownArgumentType]
+            # Always return dictionaries
+            return dict(zip(column_names, result))  # pyright: ignore[reportUnknownArgumentType,reportUnknownVariableType]
+
+    async def select_one_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncConnection]" = None,
+        schema_type: "Optional[type[ModelDTOT]]" = None,
+        **kwargs: Any,
+    ) -> "Optional[Union[ModelDTOT, dict[str, Any]]]":
+        """Fetch one row from the database.
+
+        Returns:
+            The first row of the query results.
+        """
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+
+        async with self._with_cursor(connection) as cursor:
+            await cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
+            result = await cursor.fetchone()  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+
+            if result is None:
+                return None
+
+            # Get column names
+            column_names = [col[0] for col in cursor.description or []]  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+
+            if schema_type is not None:
+                return cast("ModelDTOT", schema_type(**dict(zip(column_names, result))))  # pyright: ignore[reportUnknownArgumentType]
+            # Always return dictionaries
+            return dict(zip(column_names, result))  # pyright: ignore[reportUnknownArgumentType,reportUnknownVariableType]
+
+    async def select_value(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncConnection]" = None,
+        schema_type: "Optional[type[T]]" = None,
+        **kwargs: Any,
+    ) -> "Union[T, Any]":
+        """Fetch a single value from the database.
+
+        Returns:
+            The first value from the first row of results, or None if no results.
+        """
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+
+        async with self._with_cursor(connection) as cursor:
+            await cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
+            result = await cursor.fetchone()  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+            result = self.check_not_found(result)  # pyright: ignore[reportUnknownArgumentType]
+
+            if schema_type is None:
+                return result[0]  # pyright: ignore[reportUnknownArgumentType]
+            return schema_type(result[0])  # type: ignore[call-arg]
+
+    async def select_value_or_none(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncConnection]" = None,
+        schema_type: "Optional[type[T]]" = None,
+        **kwargs: Any,
+    ) -> "Optional[Union[T, Any]]":
+        """Fetch a single value from the database.
+
+        Returns:
+            The first value from the first row of results, or None if no results.
+        """
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+
+        async with self._with_cursor(connection) as cursor:
+            await cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
+            result = await cursor.fetchone()  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+
+            if result is None:
+                return None
+
+            if schema_type is None:
+                return result[0]  # pyright: ignore[reportUnknownArgumentType]
+            return schema_type(result[0])  # type: ignore[call-arg]
+
+    async def insert_update_delete(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncConnection]" = None,
+        **kwargs: Any,
+    ) -> int:
+        """Insert, update, or delete data from the database.
+
+        Returns:
+            Row count affected by the operation.
+        """
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+
+        async with self._with_cursor(connection) as cursor:
+            await cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
+            return cursor.rowcount  # pyright: ignore[reportUnknownMemberType]
+
+    async def insert_update_delete_returning(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncConnection]" = None,
+        schema_type: "Optional[type[ModelDTOT]]" = None,
+        **kwargs: Any,
+    ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":
+        """Insert, update, or delete data from the database and return result.
+
+        Returns:
+            The first row of results.
+        """
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+
+        async with self._with_cursor(connection) as cursor:
+            await cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
+            result = await cursor.fetchone()  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+
+            if result is None:
+                return None
+
+            # Get column names
+            column_names = [col[0] for col in cursor.description or []]  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+
+            if schema_type is not None:
+                return cast("ModelDTOT", schema_type(**dict(zip(column_names, result))))  # pyright: ignore[reportUnknownArgumentType]
+            # Always return dictionaries
+            return dict(zip(column_names, result))  # pyright: ignore[reportUnknownArgumentType,reportUnknownVariableType]
+
+    async def execute_script(
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncConnection]" = None,
+        **kwargs: Any,
+    ) -> str:
+        """Execute a script.
+
+        Returns:
+            Status message for the operation.
+        """
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+
+        async with self._with_cursor(connection) as cursor:
+            await cursor.execute(sql, parameters)  # pyright: ignore[reportUnknownMemberType]
+            return str(cursor.rowcount)  # pyright: ignore[reportUnknownMemberType]
+
+    async def select_arrow(  # pyright: ignore[reportUnknownParameterType]
+        self,
+        sql: str,
+        parameters: "Optional[StatementParameterType]" = None,
+        /,
+        *,
+        connection: "Optional[AsyncConnection]" = None,
+        **kwargs: Any,
+    ) -> "ArrowTable":  # pyright: ignore[reportUnknownVariableType]
+        """Execute a SQL query asynchronously and return results as an Apache Arrow Table.
+
+        Args:
+            sql: The SQL query string.
+            parameters: Parameters for the query.
+            connection: Optional connection override.
+            **kwargs: Additional keyword arguments to merge with parameters if parameters is a dict.
+
+        Returns:
+            An Apache Arrow Table containing the query results.
+        """
+
+        connection = self._connection(connection)
+        sql, parameters = self._process_sql_params(sql, parameters, **kwargs)
+        results = await connection.fetch_df_all(sql, parameters)
+        return ArrowTable.from_arrays(arrays=results.column_arrays(), names=results.column_names())  # pyright: ignore