polars-runtime-compat 1.34.0b3__cp39-abi3-macosx_11_0_arm64.whl → 1.34.0b4__cp39-abi3-macosx_11_0_arm64.whl

polars/io/database/_executor.py

@@ -1,578 +0,0 @@
-from __future__ import annotations
-
-import re
-from collections.abc import Coroutine, Sequence
-from contextlib import suppress
-from inspect import Parameter, signature
-from typing import TYPE_CHECKING, Any
-
-from polars import functions as F
-from polars._utils.various import parse_version
-from polars.convert import from_arrow
-from polars.datatypes import N_INFER_DEFAULT
-from polars.exceptions import (
-    DuplicateError,
-    ModuleUpgradeRequiredError,
-    UnsuitableSQLError,
-)
-from polars.io.database._arrow_registry import ARROW_DRIVER_REGISTRY
-from polars.io.database._cursor_proxies import ODBCCursorProxy, SurrealDBCursorProxy
-from polars.io.database._inference import dtype_from_cursor_description
-from polars.io.database._utils import _run_async
-
-if TYPE_CHECKING:
-    import sys
-    from collections.abc import Iterable, Iterator
-    from types import TracebackType
-
-    import pyarrow as pa
-
-    from polars.io.database._arrow_registry import ArrowDriverProperties
-
-    if sys.version_info >= (3, 11):
-        from typing import Self
-    else:
-        from typing_extensions import Self
-
-    from sqlalchemy.sql.elements import TextClause
-    from sqlalchemy.sql.expression import Selectable
-
-    from polars import DataFrame
-    from polars._typing import ConnectionOrCursor, Cursor, SchemaDict
-
-_INVALID_QUERY_TYPES = {
-    "ALTER",
-    "ANALYZE",
-    "CREATE",
-    "DELETE",
-    "DROP",
-    "INSERT",
-    "REPLACE",
-    "UPDATE",
-    "UPSERT",
-    "USE",
-    "VACUUM",
-}
-
-
-class CloseAfterFrameIter:
-    """Allows cursor close to be deferred until the last batch is returned."""
-
-    def __init__(self, frames: Any, *, cursor: Cursor) -> None:
-        self._iter_frames = frames
-        self._cursor = cursor
-
-    def __iter__(self) -> Iterator[DataFrame]:
-        yield from self._iter_frames
-
-        if hasattr(self._cursor, "close"):
-            self._cursor.close()
-
-
-class ConnectionExecutor:
-    """Abstraction for querying databases with user-supplied connection objects."""
-
-    # indicate if we can/should close the cursor on scope exit. note that we
-    # should never close the underlying connection, or a user-supplied cursor.
-    can_close_cursor: bool = False
-
-    def __init__(self, connection: ConnectionOrCursor) -> None:
-        self.driver_name = (
-            "arrow_odbc_proxy"
-            if isinstance(connection, ODBCCursorProxy)
-            else type(connection).__module__.split(".", 1)[0].lower()
-        )
-        if self.driver_name == "surrealdb":
-            connection = SurrealDBCursorProxy(client=connection)
-
-        self.cursor = self._normalise_cursor(connection)
-        self.result: Any = None
-
-    def __enter__(self) -> Self:
-        return self
-
-    def __exit__(
-        self,
-        exc_type: type[BaseException] | None,
-        exc_val: BaseException | None,
-        exc_tb: TracebackType | None,
-    ) -> None:
-        # if we created it and are finished with it, we can
-        # close the cursor (but NOT the connection)
-        if self._is_alchemy_async(self.cursor):
-            from sqlalchemy.ext.asyncio import AsyncConnection
-
-            if isinstance(self.cursor, AsyncConnection):
-                _run_async(self._close_async_cursor())
-        elif self.can_close_cursor and hasattr(self.cursor, "close"):
-            self.cursor.close()
-
-    def __repr__(self) -> str:
-        return f"<{type(self).__name__} module={self.driver_name!r}>"
-
-    @staticmethod
-    def _apply_overrides(df: DataFrame, schema_overrides: SchemaDict) -> DataFrame:
-        """Apply schema overrides to a DataFrame."""
-        existing_schema = df.schema
-        if cast_cols := [
-            F.col(col).cast(dtype)
-            for col, dtype in schema_overrides.items()
-            if col in existing_schema and dtype != existing_schema[col]
-        ]:
-            df = df.with_columns(cast_cols)
-        return df
-
-    async def _close_async_cursor(self) -> None:
-        if self.can_close_cursor and hasattr(self.cursor, "close"):
-            from sqlalchemy.ext.asyncio.exc import AsyncContextNotStarted
-
-            with suppress(AsyncContextNotStarted):
-                await self.cursor.close()
-
-    @staticmethod
-    def _check_module_version(module_name: str, minimum_version: str) -> None:
-        """Check the module version against a minimum required version."""
-        mod = __import__(module_name)
-        with suppress(AttributeError):
-            module_version: tuple[int, ...] | None = None
-            for version_attr in ("__version__", "version"):
-                if isinstance(ver := getattr(mod, version_attr, None), str):
-                    module_version = parse_version(ver)
-                    break
-            if module_version and module_version < parse_version(minimum_version):
-                msg = f"`read_database` queries require at least {module_name} version {minimum_version}"
-                raise ModuleUpgradeRequiredError(msg)
-
-    def _fetch_arrow(
-        self,
-        driver_properties: ArrowDriverProperties,
-        *,
-        batch_size: int | None,
-        iter_batches: bool,
-    ) -> Iterable[pa.RecordBatch]:
-        """Yield Arrow data as a generator of one or more RecordBatches or Tables."""
-        fetch_batches = driver_properties["fetch_batches"]
-        if not iter_batches or fetch_batches is None:
-            fetch_method = driver_properties["fetch_all"]
-            yield getattr(self.result, fetch_method)()
-        else:
-            size = [batch_size] if driver_properties["exact_batch_size"] else []
-            repeat_batch_calls = driver_properties["repeat_batch_calls"]
-            fetchmany_arrow = getattr(self.result, fetch_batches)
-            if not repeat_batch_calls:
-                yield from fetchmany_arrow(*size)
-            else:
-                while True:
-                    arrow = fetchmany_arrow(*size)
-                    if not arrow:
-                        break
-                    yield arrow
-
-    @staticmethod
-    def _fetchall_rows(result: Cursor, *, is_alchemy: bool) -> Iterable[Sequence[Any]]:
-        """Fetch row data in a single call, returning the complete result set."""
-        rows = result.fetchall()
-        return (
-            rows
-            if rows and (is_alchemy or isinstance(rows[0], (list, tuple, dict)))
-            else [tuple(row) for row in rows]
-        )
-
-    def _fetchmany_rows(
-        self, result: Cursor, *, batch_size: int | None, is_alchemy: bool
-    ) -> Iterable[Sequence[Any]]:
-        """Fetch row data incrementally, yielding over the complete result set."""
-        while True:
-            rows = result.fetchmany(batch_size)
-            if not rows:
-                break
-            elif is_alchemy or isinstance(rows[0], (list, tuple, dict)):
-                yield rows
-            else:
-                yield [tuple(row) for row in rows]
-
-    def _from_arrow(
-        self,
-        *,
-        batch_size: int | None,
-        iter_batches: bool,
-        schema_overrides: SchemaDict | None,
-        infer_schema_length: int | None,
-    ) -> DataFrame | Iterator[DataFrame] | None:
-        """Return resultset data in Arrow format for frame init."""
-        from polars import DataFrame
-
-        try:
-            for driver, driver_properties in ARROW_DRIVER_REGISTRY.items():
-                if re.match(f"^{driver}$", self.driver_name):
-                    if ver := driver_properties["minimum_version"]:
-                        self._check_module_version(self.driver_name, ver)
-
-                    if iter_batches and (
-                        driver_properties["exact_batch_size"] and not batch_size
-                    ):
-                        msg = f"Cannot set `iter_batches` for {self.driver_name} without also setting a non-zero `batch_size`"
-                        raise ValueError(msg)  # noqa: TRY301
-
-                    frames = (
-                        self._apply_overrides(batch, (schema_overrides or {}))
-                        if isinstance(batch, DataFrame)
-                        else from_arrow(batch, schema_overrides=schema_overrides)
-                        for batch in self._fetch_arrow(
-                            driver_properties,
-                            iter_batches=iter_batches,
-                            batch_size=batch_size,
-                        )
-                    )
-                    return frames if iter_batches else next(frames)  # type: ignore[arg-type,return-value]
-        except Exception as err:
-            # eg: valid turbodbc/snowflake connection, but no arrow support
-            # compiled in to the underlying driver (or on this connection)
-            arrow_not_supported = (
-                "does not support Apache Arrow",
-                "Apache Arrow format is not supported",
-            )
-            if not any(e in str(err) for e in arrow_not_supported):
-                raise
-
-        return None
-
-    def _from_rows(
-        self,
-        *,
-        batch_size: int | None,
-        iter_batches: bool,
-        schema_overrides: SchemaDict | None,
-        infer_schema_length: int | None,
-    ) -> DataFrame | Iterator[DataFrame] | None:
-        """Return resultset data row-wise for frame init."""
-        from polars import DataFrame
-
-        if iter_batches and not batch_size:
-            msg = (
-                "Cannot set `iter_batches` without also setting a non-zero `batch_size`"
-            )
-            raise ValueError(msg)
-
-        if is_async := isinstance(original_result := self.result, Coroutine):
-            self.result = _run_async(self.result)
-        try:
-            if hasattr(self.result, "fetchall"):
-                if is_alchemy := (self.driver_name == "sqlalchemy"):
-                    if hasattr(self.result, "cursor"):
-                        cursor_desc = [
-                            (d[0], d[1:]) for d in self.result.cursor.description
-                        ]
-                    elif hasattr(self.result, "_metadata"):
-                        cursor_desc = [(k, None) for k in self.result._metadata.keys]
-                    else:
-                        msg = f"Unable to determine metadata from query result; {self.result!r}"
-                        raise ValueError(msg)
-
-                elif hasattr(self.result, "description"):
-                    cursor_desc = [(d[0], d[1:]) for d in self.result.description]
-                else:
-                    cursor_desc = []
-
-                schema_overrides = self._inject_type_overrides(
-                    description=cursor_desc,
-                    schema_overrides=(schema_overrides or {}),
-                )
-                result_columns = [nm for nm, _ in cursor_desc]
-                frames = (
-                    DataFrame(
-                        data=rows,
-                        schema=result_columns or None,
-                        schema_overrides=schema_overrides,
-                        infer_schema_length=infer_schema_length,
-                        orient="row",
-                    )
-                    for rows in (
-                        self._fetchmany_rows(
-                            self.result,
-                            batch_size=batch_size,
-                            is_alchemy=is_alchemy,
-                        )
-                        if iter_batches
-                        else [self._fetchall_rows(self.result, is_alchemy=is_alchemy)]  # type: ignore[list-item]
-                    )
-                )
-                return frames if iter_batches else next(frames)  # type: ignore[arg-type]
-            return None
-        finally:
-            if is_async:
-                original_result.close()
-
-    def _inject_type_overrides(
-        self,
-        description: list[tuple[str, Any]],
-        schema_overrides: SchemaDict,
-    ) -> SchemaDict:
-        """
-        Attempt basic dtype inference from a cursor description.
-
-        Notes
-        -----
-        This is limited; the `type_code` description attr may contain almost anything,
-        from strings or python types to driver-specific codes, classes, enums, etc.
-        We currently only do the additional inference from string/python type values.
-        (Further refinement will require per-driver module knowledge and lookups).
-        """
-        dupe_check = set()
-        for nm, desc in description:
-            if nm in dupe_check:
-                msg = f"column {nm!r} appears more than once in the query/result cursor"
-                raise DuplicateError(msg)
-            elif desc is not None and nm not in schema_overrides:
-                dtype = dtype_from_cursor_description(self.cursor, desc)
-                if dtype is not None:
-                    schema_overrides[nm] = dtype  # type: ignore[index]
-            dupe_check.add(nm)
-
-        return schema_overrides
-
-    @staticmethod
-    def _is_alchemy_async(conn: Any) -> bool:
-        """Check if the given connection is SQLALchemy async."""
-        try:
-            from sqlalchemy.ext.asyncio import (
-                AsyncConnection,
-                AsyncSession,
-                async_sessionmaker,
-            )
-
-            return isinstance(conn, (AsyncConnection, AsyncSession, async_sessionmaker))
-        except ImportError:
-            return False
-
-    @staticmethod
-    def _is_alchemy_engine(conn: Any) -> bool:
-        """Check if the given connection is a SQLAlchemy Engine."""
-        from sqlalchemy.engine import Engine
-
-        if isinstance(conn, Engine):
-            return True
-        try:
-            from sqlalchemy.ext.asyncio import AsyncEngine
-
-            return isinstance(conn, AsyncEngine)
-        except ImportError:
-            return False
-
-    @staticmethod
-    def _is_alchemy_object(conn: Any) -> bool:
-        """Check if the given connection is a SQLAlchemy object (of any kind)."""
-        return type(conn).__module__.split(".", 1)[0] == "sqlalchemy"
-
-    @staticmethod
-    def _is_alchemy_session(conn: Any) -> bool:
-        """Check if the given connection is a SQLAlchemy Session object."""
-        from sqlalchemy.ext.asyncio import AsyncSession
-        from sqlalchemy.orm import Session, sessionmaker
-
-        if isinstance(conn, (AsyncSession, Session, sessionmaker)):
-            return True
-
-        try:
-            from sqlalchemy.ext.asyncio import async_sessionmaker
-
-            return isinstance(conn, async_sessionmaker)
-        except ImportError:
-            return False
-
-    @staticmethod
-    def _is_alchemy_result(result: Any) -> bool:
-        """Check if the given result is a SQLAlchemy Result object."""
-        try:
-            from sqlalchemy.engine import CursorResult
-
-            if isinstance(result, CursorResult):
-                return True
-
-            from sqlalchemy.ext.asyncio import AsyncResult
-
-            return isinstance(result, AsyncResult)
-        except ImportError:
-            return False
-
-    def _normalise_cursor(self, conn: Any) -> Cursor:
-        """Normalise a connection object such that we have the query executor."""
-        if self.driver_name == "sqlalchemy":
-            if self._is_alchemy_session(conn):
-                return conn
-            else:
-                # where possible, use the raw connection to access arrow integration
-                if conn.engine.driver == "databricks-sql-python":
-                    self.driver_name = "databricks"
-                    return conn.engine.raw_connection().cursor()
-                elif conn.engine.driver == "duckdb_engine":
-                    self.driver_name = "duckdb"
-                    return conn
-                elif self._is_alchemy_engine(conn):
-                    # note: if we create it, we can close it
-                    self.can_close_cursor = True
-                    return conn.connect()
-                else:
-                    return conn
-
-        elif hasattr(conn, "cursor"):
-            # connection has a dedicated cursor; prefer over direct execute
-            cursor = cursor() if callable(cursor := conn.cursor) else cursor
-            self.can_close_cursor = True
-            return cursor
-
-        elif hasattr(conn, "execute"):
-            # can execute directly (given cursor, sqlalchemy connection, etc)
-            return conn
-
-        msg = f"""Unrecognised connection type "{conn!r}"; no 'execute' or 'cursor' method"""
-        raise TypeError(msg)
-
-    async def _sqlalchemy_async_execute(self, query: TextClause, **options: Any) -> Any:
-        """Execute a query using an async SQLAlchemy connection."""
-        is_session = self._is_alchemy_session(self.cursor)
-        cursor = self.cursor.begin() if is_session else self.cursor  # type: ignore[attr-defined]
-        async with cursor as conn:  # type: ignore[union-attr]
-            if is_session and not hasattr(conn, "execute"):
-                conn = conn.session
-            result = await conn.execute(query, **options)
-            return result
-
-    def _sqlalchemy_setup(
-        self, query: str | TextClause | Selectable, options: dict[str, Any]
-    ) -> tuple[Any, dict[str, Any], str | TextClause | Selectable]:
-        """Prepare a query for execution using a SQLAlchemy connection."""
-        from sqlalchemy.orm import Session
-        from sqlalchemy.sql import text
-        from sqlalchemy.sql.elements import TextClause
-
-        param_key = "parameters"
-        cursor_execute = None
-        if (
-            isinstance(self.cursor, Session)
-            and "parameters" in options
-            and "params" not in options
-        ):
-            options = options.copy()
-            options["params"] = options.pop("parameters")
-            param_key = "params"
-
-        params = options.get(param_key)
-        is_async = self._is_alchemy_async(self.cursor)
-        if (
-            not is_async
-            and isinstance(params, Sequence)
-            and hasattr(self.cursor, "exec_driver_sql")
-        ):
-            cursor_execute = self.cursor.exec_driver_sql
-            if isinstance(query, TextClause):
-                query = str(query)
-            if isinstance(params, list) and not all(
-                isinstance(p, (dict, tuple)) for p in params
-            ):
-                options[param_key] = tuple(params)
-
-        elif isinstance(query, str):
-            query = text(query)
-
-        if cursor_execute is None:
-            cursor_execute = (
-                self._sqlalchemy_async_execute if is_async else self.cursor.execute
-            )
-        return cursor_execute, options, query
-
-    def execute(
-        self,
-        query: str | TextClause | Selectable,
-        *,
-        options: dict[str, Any] | None = None,
-        select_queries_only: bool = True,
-    ) -> Self:
-        """Execute a query and reference the result set."""
-        if select_queries_only and isinstance(query, str):
-            q = re.search(r"\w{3,}", re.sub(r"/\*(.|[\r\n])*?\*/", "", query))
-            if (query_type := "" if not q else q.group(0)) in _INVALID_QUERY_TYPES:
-                msg = f"{query_type} statements are not valid 'read' queries"
-                raise UnsuitableSQLError(msg)
-
-        options = options or {}
-
-        if self._is_alchemy_object(self.cursor):
-            cursor_execute, options, query = self._sqlalchemy_setup(query, options)
-        else:
-            cursor_execute = self.cursor.execute
-
-        # note: some cursor execute methods (eg: sqlite3) only take positional
-        # params, hence the slightly convoluted resolution of the 'options' dict
-        try:
-            params = signature(cursor_execute).parameters
-        except ValueError:
-            params = {}  # type: ignore[assignment]
-
-        if not options or any(
-            p.kind in (Parameter.KEYWORD_ONLY, Parameter.POSITIONAL_OR_KEYWORD)
-            for p in params.values()
-        ):
-            result = cursor_execute(query, **options)
-        else:
-            positional_options = (
-                options[o] for o in (params or options) if (not options or o in options)
-            )
-            result = cursor_execute(query, *positional_options)
-
-        # note: some cursors execute in-place, some access results via a property
-        result = self.cursor if (result is None or result is True) else result
-        if self.driver_name == "duckdb" and self._is_alchemy_result(result):
-            result = result.cursor
-
-        self.result = result
-        return self
-
-    def to_polars(
-        self,
-        *,
-        iter_batches: bool = False,
-        batch_size: int | None = None,
-        schema_overrides: SchemaDict | None = None,
-        infer_schema_length: int | None = N_INFER_DEFAULT,
-    ) -> DataFrame | Iterator[DataFrame]:
-        """
-        Convert the result set to a DataFrame.
-
-        Wherever possible we try to return arrow-native data directly; only
-        fall back to initialising with row-level data if no other option.
-        """
-        if self.result is None:
-            msg = "cannot return a frame before executing a query"
-            raise RuntimeError(msg)
-
-        can_close = self.can_close_cursor
-
-        if defer_cursor_close := (iter_batches and can_close):
-            self.can_close_cursor = False
-
-        for frame_init in (
-            self._from_arrow,  # init from arrow-native data (where support exists)
-            self._from_rows,  # row-wise fallback (sqlalchemy, dbapi2, pyodbc, etc)
-        ):
-            frame = frame_init(
-                batch_size=batch_size,
-                iter_batches=iter_batches,
-                schema_overrides=schema_overrides,
-                infer_schema_length=infer_schema_length,
-            )
-            if frame is not None:
-                if defer_cursor_close:
-                    frame = (
-                        df
-                        for df in CloseAfterFrameIter(
-                            frame,
-                            cursor=self.result,
-                        )
-                    )
-                return frame
-
-        msg = (
-            f"Currently no support for {self.driver_name!r} connection {self.cursor!r}"
-        )
-        raise NotImplementedError(msg)