arrowsql 0.1.0__cp313-cp313-win_amd64.whl

arrowsql/async_query/async_connection.py

@@ -0,0 +1,160 @@
+"""Public asynchronous connection backed directly by mssql-py-core.
+
+Warning:
+    Async query execution APIs are under active development and are not intended
+    for production use. Their signatures, behavior, error handling, and compatibility
+    may change without notice.
+"""
+
+from typing import Any, Optional
+
+from ..logging import logger
+from ._native import load_py_core
+from ._connection_context import build_async_connection_context
+from .async_cursor import _AsyncCursor  # pyright: ignore[reportPrivateUsage]
+from .exception_translator import (
+    DataError,
+    DatabaseError,
+    Error,
+    IntegrityError,
+    InterfaceError,
+    InternalError,
+    NotSupportedError,
+    OperationalError,
+    ProgrammingError,
+    Warning,
+    translate_py_core_exceptions,
+)
+
+
+class _AsyncConnection:
+    """Thin Python wrapper over ``mssql_py_core.PyAsyncConnection``.
+
+    Warning:
+        This preview API is under active development and is not intended for production use.
+        Its signatures, behavior, error handling, and compatibility may change without notice.
+    """
+
+    Warning = Warning
+    Error = Error
+    InterfaceError = InterfaceError
+    DatabaseError = DatabaseError
+    DataError = DataError
+    OperationalError = OperationalError
+    IntegrityError = IntegrityError
+    InternalError = InternalError
+    ProgrammingError = ProgrammingError
+    NotSupportedError = NotSupportedError
+
+    def __init__(self, py_core_async_connection: Any) -> None:
+        self._py_core_async_connection = py_core_async_connection
+
+    @classmethod
+    async def connect(
+        cls,
+        connection_str: str = "",
+        autocommit: bool = False,
+        timeout: int = 0,
+        python_logger: Optional[Any] = None,
+    ) -> "_AsyncConnection":
+        """Establish an asynchronous connection from an ODBC connection string."""
+        logger_bridge = python_logger
+        if logger_bridge is None and logger.is_debug_enabled:
+            logger_bridge = logger
+        logger.debug(
+            "AsyncConnection.connect: starting; autocommit=%s; logger_source=%s",
+            autocommit,
+            (
+                "custom"
+                if python_logger is not None
+                else "mssql_python" if logger_bridge is not None else "disabled"
+            ),
+        )
+        with translate_py_core_exceptions():
+            client_context_dict = build_async_connection_context(connection_str, timeout)
+            py_core = load_py_core()
+            py_core_async_connection = await py_core.PyAsyncConnection.connect(
+                client_context_dict,
+                python_logger=logger_bridge,
+                autocommit=autocommit,
+            )
+        logger.debug("AsyncConnection.connect: connected")
+        return cls(py_core_async_connection)
+
+    def cursor(self) -> _AsyncCursor:
+        """Create a public asynchronous cursor sharing this connection."""
+        with translate_py_core_exceptions():
+            py_core_async_cursor = self._py_core_async_connection.cursor()
+        logger.debug("AsyncConnection.cursor: cursor created")
+        return _AsyncCursor(py_core_async_cursor, self)
+
+    async def commit(self) -> None:
+        """Commit the active transaction, if any."""
+        logger.debug("AsyncConnection.commit: starting")
+        with translate_py_core_exceptions():
+            await self._py_core_async_connection.commit()
+        logger.debug("AsyncConnection.commit: completed")
+
+    async def rollback(self) -> None:
+        """Roll back the active transaction, if any."""
+        logger.debug("AsyncConnection.rollback: starting")
+        with translate_py_core_exceptions():
+            await self._py_core_async_connection.rollback()
+        logger.debug("AsyncConnection.rollback: completed")
+
+    async def close(self) -> None:
+        """Close the py-core async connection."""
+        logger.debug("AsyncConnection.close: starting")
+        with translate_py_core_exceptions():
+            await self._py_core_async_connection.close()
+        logger.debug("AsyncConnection.close: completed")
+
+    async def __aenter__(self) -> "_AsyncConnection":
+        logger.debug("AsyncConnection.__aenter__: entering context")
+        with translate_py_core_exceptions():
+            await self._py_core_async_connection.__aenter__()
+        logger.debug("AsyncConnection.__aenter__: context entered")
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback) -> Any:
+        logger.debug(
+            "AsyncConnection.__aexit__: exiting context; block_error=%s",
+            exc_type is not None,
+        )
+        with translate_py_core_exceptions():
+            result = await self._py_core_async_connection.__aexit__(exc_type, exc_value, traceback)
+        logger.debug("AsyncConnection.__aexit__: context exited")
+        return result
+
+    @property
+    def timeout(self) -> int:
+        """Default query timeout inherited by subsequently created cursors."""
+        with translate_py_core_exceptions():
+            return self._py_core_async_connection.timeout
+
+    @timeout.setter
+    def timeout(self, value: int) -> None:
+        with translate_py_core_exceptions():
+            self._py_core_async_connection.timeout = value
+        logger.debug("AsyncConnection.timeout: updated")
+
+    @property
+    def autocommit(self) -> bool:
+        """Whether the connection was opened in autocommit mode."""
+        with translate_py_core_exceptions():
+            return self._py_core_async_connection.autocommit
+
+    @property
+    def closed(self) -> bool:
+        """Whether close has been initiated on the py-core async connection."""
+        with translate_py_core_exceptions():
+            return self._py_core_async_connection.closed
+
+    def is_connected(self) -> bool:
+        """Return whether the py-core async connection remains open."""
+        with translate_py_core_exceptions():
+            return self._py_core_async_connection.is_connected()
+
+    def __repr__(self) -> str:
+        state = "closed" if self.closed else "connected"
+        return f"_AsyncConnection({state})"

arrowsql/async_query/async_cursor.py

@@ -0,0 +1,225 @@
+"""Internal asynchronous cursor backed directly by mssql-py-core.
+
+Warning:
+    Async query execution APIs are under active development and are not intended
+    for production use. Their signatures, behavior, error handling, and compatibility
+    may change without notice.
+"""
+
+import asyncio
+from collections.abc import Mapping, Sequence
+from contextlib import asynccontextmanager
+from typing import Any, Optional
+import uuid
+
+from ..helpers import get_settings
+from ..logging import logger
+from ..row import Row
+from . import async_execute, async_fetch
+from .exception_translator import is_connection_busy, translate_py_core_exceptions
+
+
+class _AsyncCursor:
+    """Internal Python wrapper over ``mssql_py_core.PyAsyncCursor``.
+
+    Warning:
+        This preview API is under active development and is not intended for production use.
+        Its signatures, behavior, error handling, and compatibility may change without notice.
+    """
+
+    def __init__(self, py_core_async_cursor: Any, connection: Any = None) -> None:
+        self._py_core_async_cursor = py_core_async_cursor
+        self._connection = connection
+        self._closed = False
+        self._result_transition_lock = asyncio.Lock()
+        self._result_ready = asyncio.Event()
+        self._result_ready.set()
+        self._result_generation = 0
+        self._fetched_row_count = 0
+        self._fetch_rowcount: int | None = None
+        self._description: list[tuple[Any, ...]] | None = None
+        self._column_map: dict[str, int] = {}
+        self._column_map_lower: dict[str, int] | None = None
+        self._column_names: tuple[str, ...] | None = None
+        self._uuid_str_indices: tuple[int, ...] | None = None
+
+    def _clear_result_metadata(self) -> None:
+        self._result_generation += 1
+        self._description = None
+        self._column_map = {}
+        self._column_map_lower = None
+        self._column_names = None
+        self._uuid_str_indices = None
+
+    def _initialize_result_metadata(self) -> None:
+        with translate_py_core_exceptions():
+            description = self._py_core_async_cursor.description
+        if description is None:
+            self._clear_result_metadata()
+            return
+
+        settings = get_settings()
+        self._description = [
+            ((column[0].lower() if settings.lowercase else column[0]), *column[1:])
+            for column in description
+        ]
+        self._column_names = tuple(column[0] for column in self._description)
+        self._column_map = {column[0]: index for index, column in enumerate(self._description)}
+        self._column_map_lower = (
+            {name.lower(): index for name, index in self._column_map.items()}
+            if settings.lowercase
+            else None
+        )
+        self._uuid_str_indices = (
+            tuple(index for index, column in enumerate(self._description) if column[1] is uuid.UUID)
+            if not settings.native_uuid
+            else None
+        )
+
+    def _reset_fetch_tracking(self) -> None:
+        self._fetched_row_count = 0
+        self._fetch_rowcount = None
+
+    @asynccontextmanager
+    async def _result_transition(self):
+        async with self._result_transition_lock:
+            self._result_ready.clear()
+            try:
+                yield
+            finally:
+                self._result_ready.set()
+
+    async def _wait_for_result_publication(self) -> None:
+        while True:
+            await self._result_ready.wait()
+            if self._result_ready.is_set():
+                return
+
+    def _reconcile_failed_result_operation(self, operation: str, error: BaseException) -> None:
+        if is_connection_busy(error):
+            return
+        self._reset_fetch_tracking()
+        self._clear_result_metadata()
+        try:
+            self._initialize_result_metadata()
+        except Exception as error:
+            logger.debug("AsyncCursor.%s: metadata recovery failed: %s", operation, error)
+
+    def _check_closed(self) -> None:
+        if self._closed or (self._connection is not None and self._connection.closed):
+            message = "Cursor is closed" if self._closed else "Connection is closed"
+            with translate_py_core_exceptions():
+                raise RuntimeError(message)
+
+    def _record_fetch(self, generation: int, count: int, exhausted: bool) -> None:
+        if generation != self._result_generation:
+            return
+        if count:
+            self._fetched_row_count += count
+            self._fetch_rowcount = self._fetched_row_count
+        elif exhausted and self._fetched_row_count == 0:
+            self._fetch_rowcount = 0
+
+    async def execute(
+        self,
+        operation: str,
+        *parameters: Any,
+        use_prepare: bool = True,
+        reset_cursor: bool = True,
+    ) -> "_AsyncCursor":
+        async with self._result_transition():
+            return await async_execute.execute(
+                self,
+                operation,
+                *parameters,
+                use_prepare=use_prepare,
+                reset_cursor=reset_cursor,
+            )
+
+    async def executemany(
+        self,
+        operation: str,
+        seq_of_parameters: Sequence[Sequence[Any]] | Sequence[Mapping[str, Any]],
+        *,
+        use_prepare: bool = True,
+    ) -> None:
+        async with self._result_transition():
+            await async_execute.executemany(
+                self,
+                operation,
+                seq_of_parameters,
+                use_prepare=use_prepare,
+            )
+
+    async def fetchone(self) -> Row | None:
+        return await async_fetch.fetchone(self)
+
+    async def fetchmany(self, size: Optional[int] = None) -> list[Row]:
+        return await async_fetch.fetchmany(self, size)
+
+    async def fetchall(self) -> list[Row]:
+        return await async_fetch.fetchall(self)
+
+    async def nextset(self) -> bool:
+        async with self._result_transition():
+            try:
+                with translate_py_core_exceptions():
+                    has_next = await self._py_core_async_cursor.nextset()
+            except (Exception, asyncio.CancelledError) as error:
+                self._reconcile_failed_result_operation("nextset", error)
+                raise
+            self._reset_fetch_tracking()
+            self._clear_result_metadata()
+            if has_next:
+                self._initialize_result_metadata()
+            return has_next
+
+    async def close(self) -> None:
+        logger.debug("AsyncCursor.close: starting")
+        async with self._result_transition():
+            with translate_py_core_exceptions():
+                close_awaitable = self._py_core_async_cursor.close()
+            close_accepted = True
+            try:
+                with translate_py_core_exceptions():
+                    await close_awaitable
+            except (Exception, asyncio.CancelledError) as error:
+                if is_connection_busy(error):
+                    close_accepted = False
+                raise
+            finally:
+                if close_accepted:
+                    self._closed = True
+                    self._reset_fetch_tracking()
+                    self._clear_result_metadata()
+        logger.debug("AsyncCursor.close: completed")
+
+    def setinputsizes(self, sizes: Any) -> None:
+        with translate_py_core_exceptions():
+            self._py_core_async_cursor.setinputsizes(sizes)
+
+    @property
+    def timeout(self) -> int:
+        with translate_py_core_exceptions():
+            return self._py_core_async_cursor.timeout
+
+    @property
+    def description(self) -> Any:
+        return self._description
+
+    @property
+    def rowcount(self) -> int:
+        if self._fetch_rowcount is not None:
+            return self._fetch_rowcount
+        with translate_py_core_exceptions():
+            return self._py_core_async_cursor.rowcount
+
+    @property
+    def arraysize(self) -> int:
+        with translate_py_core_exceptions():
+            return self._py_core_async_cursor.arraysize
+
+    @arraysize.setter
+    def arraysize(self, value: int) -> None:
+        with translate_py_core_exceptions():
+            self._py_core_async_cursor.arraysize = value

arrowsql/async_query/async_execute.py

@@ -0,0 +1,113 @@
+"""Asynchronous statement execution through mssql-py-core."""
+
+import asyncio
+from collections.abc import Mapping, Sequence
+from typing import TYPE_CHECKING, Any
+
+from ..exceptions import ProgrammingError
+from ..logging import logger
+from ..row import Row
+from .exception_translator import is_connection_busy, translate_py_core_exceptions
+
+if TYPE_CHECKING:
+    from .async_cursor import _AsyncCursor  # pyright: ignore[reportPrivateUsage]
+
+
+def _get_py_core_async_cursor(cursor: "_AsyncCursor") -> Any:
+    return cursor._py_core_async_cursor  # pyright: ignore[reportPrivateUsage]
+
+
+def _is_non_mutating_rejection(error: BaseException) -> bool:
+    """Return whether py-core guarantees rejection before result-state mutation."""
+    if is_connection_busy(error):
+        return True
+    return isinstance(error, (TypeError, KeyError)) or (
+        isinstance(error, ProgrammingError) and isinstance(error.__cause__, TypeError)
+    )
+
+
+def _reconcile_failed_execution(cursor: "_AsyncCursor", error: BaseException) -> None:
+    if _is_non_mutating_rejection(error):
+        return
+    cursor._reset_fetch_tracking()  # pyright: ignore[reportPrivateUsage]
+    cursor._clear_result_metadata()  # pyright: ignore[reportPrivateUsage]
+    try:
+        cursor._initialize_result_metadata()  # pyright: ignore[reportPrivateUsage]
+    except Exception as error:
+        logger.debug("Async execution metadata recovery failed: %s", error)
+
+
+async def execute(
+    cursor: "_AsyncCursor",
+    operation: str,
+    *parameters: Any,
+    use_prepare: bool = True,
+    reset_cursor: bool = True,
+) -> "_AsyncCursor":
+    """Execute a statement using the py-core async cursor."""
+    if len(parameters) == 1 and isinstance(parameters[0], (tuple, list, Row)):
+        parameters = tuple(parameters[0])
+
+    logger.debug(
+        "AsyncCursor.execute: starting; param_count=%d; use_prepare=%s; reset_cursor=%s",
+        len(parameters),
+        use_prepare,
+        reset_cursor,
+    )
+    with translate_py_core_exceptions():
+        execute_awaitable = _get_py_core_async_cursor(cursor).execute(
+            operation,
+            *parameters,
+            use_prepare=use_prepare,
+            reset_cursor=reset_cursor,
+        )
+    try:
+        with translate_py_core_exceptions():
+            await execute_awaitable
+    except (Exception, asyncio.CancelledError) as error:
+        _reconcile_failed_execution(cursor, error)
+        raise
+    cursor._reset_fetch_tracking()  # pyright: ignore[reportPrivateUsage]
+    cursor._clear_result_metadata()  # pyright: ignore[reportPrivateUsage]
+    cursor._initialize_result_metadata()  # pyright: ignore[reportPrivateUsage]
+    description = cursor.description
+    logger.debug(
+        "AsyncCursor.execute: completed; rowcount=%d; column_count=%d; has_result_set=%s",
+        cursor.rowcount,
+        len(description) if description is not None else 0,
+        description is not None,
+    )
+    return cursor
+
+
+async def executemany(
+    cursor: "_AsyncCursor",
+    operation: str,
+    seq_of_parameters: Sequence[Sequence[Any]] | Sequence[Mapping[str, Any]],
+    *,
+    use_prepare: bool = True,
+) -> None:
+    """Execute a statement for every parameter row using the py-core async cursor."""
+    cursor._check_closed()  # pyright: ignore[reportPrivateUsage]
+    batch_count = len(seq_of_parameters)
+    logger.debug(
+        "AsyncCursor.executemany: starting; batch_count=%d; use_prepare=%s",
+        batch_count,
+        use_prepare,
+    )
+    with translate_py_core_exceptions():
+        executemany_awaitable = _get_py_core_async_cursor(cursor).executemany(
+            operation,
+            seq_of_parameters,
+            use_prepare=use_prepare,
+        )
+    try:
+        with translate_py_core_exceptions():
+            await executemany_awaitable
+    except (Exception, asyncio.CancelledError) as error:
+        _reconcile_failed_execution(cursor, error)
+        raise
+    cursor._reset_fetch_tracking()  # pyright: ignore[reportPrivateUsage]
+    cursor._clear_result_metadata()  # pyright: ignore[reportPrivateUsage]
+    cursor._initialize_result_metadata()  # pyright: ignore[reportPrivateUsage]
+    logger.debug("AsyncCursor.executemany: completed; rowcount=%d", cursor.rowcount)

arrowsql/async_query/async_fetch.py

@@ -0,0 +1,153 @@
+"""Asynchronous result fetching through mssql-py-core."""
+
+import asyncio
+import logging
+
+from typing import TYPE_CHECKING, Any
+
+from ..exceptions import ProgrammingError
+from ..logging import logger
+from ..row import Row
+from .async_execute import _get_py_core_async_cursor
+from .exception_translator import (
+    is_connection_busy,
+    translate_py_core_exception,
+    translate_py_core_exceptions,
+)
+
+if TYPE_CHECKING:
+    from .async_cursor import _AsyncCursor  # pyright: ignore[reportPrivateUsage]
+
+_ResultSnapshot = tuple[
+    int,
+    dict[str, int],
+    tuple[int, ...] | None,
+    dict[str, int] | None,
+    tuple[str, ...] | None,
+]
+
+
+def _snapshot_result(cursor: "_AsyncCursor") -> _ResultSnapshot:
+    return (
+        cursor._result_generation,  # pyright: ignore[reportPrivateUsage]
+        cursor._column_map,  # pyright: ignore[reportPrivateUsage]
+        cursor._uuid_str_indices,  # pyright: ignore[reportPrivateUsage]
+        cursor._column_map_lower,  # pyright: ignore[reportPrivateUsage]
+        cursor._column_names,  # pyright: ignore[reportPrivateUsage]
+    )
+
+
+def _reconcile_failed_fetch(
+    cursor: "_AsyncCursor", generation: int, operation: str, error: BaseException
+) -> None:
+    if generation != cursor._result_generation:  # pyright: ignore[reportPrivateUsage]
+        return
+    if is_connection_busy(error):
+        return
+    logger.debug("AsyncCursor.%s: invalidating result state after fetch failure", operation)
+    cursor._reset_fetch_tracking()  # pyright: ignore[reportPrivateUsage]
+    cursor._clear_result_metadata()  # pyright: ignore[reportPrivateUsage]
+
+
+def _wrap_row(snapshot: _ResultSnapshot, values: tuple[Any, ...]) -> Row:
+    _, column_map, uuid_str_indices, column_map_lower, column_names = snapshot
+    if not uuid_str_indices:
+        return Row._fast_create(list(values), column_map, None, column_map_lower, column_names)  # pyright: ignore[reportPrivateUsage]
+    return Row(
+        list(values),
+        column_map,
+        uuid_str_indices=uuid_str_indices,
+        column_map_lower=column_map_lower,
+        column_names=column_names,
+    )
+
+
+def _wrap_rows(snapshot: _ResultSnapshot, rows: list[tuple[Any, ...]]) -> list[Row]:
+    _, column_map, uuid_str_indices, column_map_lower, column_names = snapshot
+    if uuid_str_indices:
+        return [_wrap_row(snapshot, row) for row in rows]
+    create = Row._fast_create  # pyright: ignore[reportPrivateUsage]
+    return [create(list(row), column_map, None, column_map_lower, column_names) for row in rows]
+
+
+async def fetchone(cursor: "_AsyncCursor") -> Row | None:
+    """Fetch the next row through the py-core async cursor."""
+    logger.debug("AsyncCursor.fetchone: starting")
+    if not cursor._result_ready.is_set():  # pyright: ignore[reportPrivateUsage]
+        await cursor._wait_for_result_publication()  # pyright: ignore[reportPrivateUsage]
+    snapshot = _snapshot_result(cursor)
+    try:
+        try:
+            row = await _get_py_core_async_cursor(cursor).fetchone()
+        except Exception as error:
+            translated = translate_py_core_exception(error)
+            if translated is error:
+                raise
+            raise translated from error
+    except (Exception, asyncio.CancelledError) as error:
+        _reconcile_failed_fetch(cursor, snapshot[0], "fetchone", error)
+        raise
+    cursor._record_fetch(  # pyright: ignore[reportPrivateUsage]
+        snapshot[0], row is not None, row is None
+    )
+    if logger.isEnabledFor(logging.DEBUG):
+        logger.debug(
+            "AsyncCursor.fetchone: completed; row_found=%s; rowcount=%d",
+            row is not None,
+            cursor.rowcount,
+        )
+    return None if row is None else _wrap_row(snapshot, row)
+
+
+async def fetchmany(cursor: "_AsyncCursor", size: int | None = None) -> list[Row]:
+    """Fetch up to size rows, using cursor arraysize when size is omitted."""
+    if not cursor._result_ready.is_set():  # pyright: ignore[reportPrivateUsage]
+        await cursor._wait_for_result_publication()  # pyright: ignore[reportPrivateUsage]
+    cursor._check_closed()  # pyright: ignore[reportPrivateUsage]
+    requested_size = cursor.arraysize if size is None else size
+    logger.debug("AsyncCursor.fetchmany: starting; requested_size=%s", requested_size)
+    if requested_size <= 0:
+        if cursor.description is None:
+            raise ProgrammingError("Async operation failed", "No active result set")
+        logger.debug("AsyncCursor.fetchmany: completed; row_count=0; rowcount=%d", cursor.rowcount)
+        return []
+    snapshot = _snapshot_result(cursor)
+    with translate_py_core_exceptions():
+        if size is None:
+            fetch_awaitable = _get_py_core_async_cursor(cursor).fetchmany()
+        else:
+            fetch_awaitable = _get_py_core_async_cursor(cursor).fetchmany(size)
+    try:
+        with translate_py_core_exceptions():
+            rows = await fetch_awaitable
+    except (Exception, asyncio.CancelledError) as error:
+        _reconcile_failed_fetch(cursor, snapshot[0], "fetchmany", error)
+        raise
+    cursor._record_fetch(snapshot[0], len(rows), not rows)  # pyright: ignore[reportPrivateUsage]
+    logger.debug(
+        "AsyncCursor.fetchmany: completed; row_count=%d; rowcount=%d",
+        len(rows),
+        cursor.rowcount,
+    )
+    return _wrap_rows(snapshot, rows)
+
+
+async def fetchall(cursor: "_AsyncCursor") -> list[Row]:
+    """Fetch all remaining rows through the py-core async cursor."""
+    logger.debug("AsyncCursor.fetchall: starting")
+    if not cursor._result_ready.is_set():  # pyright: ignore[reportPrivateUsage]
+        await cursor._wait_for_result_publication()  # pyright: ignore[reportPrivateUsage]
+    snapshot = _snapshot_result(cursor)
+    try:
+        with translate_py_core_exceptions():
+            rows = await _get_py_core_async_cursor(cursor).fetchall()
+    except (Exception, asyncio.CancelledError) as error:
+        _reconcile_failed_fetch(cursor, snapshot[0], "fetchall", error)
+        raise
+    cursor._record_fetch(snapshot[0], len(rows), not rows)  # pyright: ignore[reportPrivateUsage]
+    logger.debug(
+        "AsyncCursor.fetchall: completed; row_count=%d; rowcount=%d",
+        len(rows),
+        cursor.rowcount,
+    )
+    return _wrap_rows(snapshot, rows)