Merge branch 'main' into jahnvi/fix-fetch-diagnostics

Author: jahnvi480

mssql_python/async_query/async_connection.py

@@ -46,8 +46,8 @@ class AsyncConnection:
     ProgrammingError = ProgrammingError
     NotSupportedError = NotSupportedError
 
-    def __init__(self, native_connection: Any) -> None:
-        self._native_connection = native_connection
+    def __init__(self, py_core_async_connection: Any) -> None:
+        self._py_core_async_connection = py_core_async_connection
 
     @classmethod
     async def connect(
@@ -58,54 +58,61 @@ async def connect(
         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; custom_logger=%s",
+            "AsyncConnection.connect: starting; autocommit=%s; logger_source=%s",
             autocommit,
-            python_logger is not None,
+            (
+                "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()
-            native_connection = await py_core.PyAsyncConnection.connect(
+            py_core_async_connection = await py_core.PyAsyncConnection.connect(
                 client_context_dict,
-                python_logger=python_logger,
+                python_logger=logger_bridge,
                 autocommit=autocommit,
             )
         logger.debug("AsyncConnection.connect: connected")
-        return cls(native_connection)
+        return cls(py_core_async_connection)
 
     def cursor(self) -> AsyncCursor:
         """Create a public asynchronous cursor sharing this connection."""
         with translate_py_core_exceptions():
-            native_cursor = self._native_connection.cursor()
+            py_core_async_cursor = self._py_core_async_connection.cursor()
         logger.debug("AsyncConnection.cursor: cursor created")
-        return AsyncCursor(native_cursor)
+        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._native_connection.commit()
+            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._native_connection.rollback()
+            await self._py_core_async_connection.rollback()
         logger.debug("AsyncConnection.rollback: completed")
 
     async def close(self) -> None:
-        """Close the native connection."""
+        """Close the py-core async connection."""
         logger.debug("AsyncConnection.close: starting")
         with translate_py_core_exceptions():
-            await self._native_connection.close()
+            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._native_connection.__aenter__()
+            await self._py_core_async_connection.__aenter__()
         logger.debug("AsyncConnection.__aenter__: context entered")
         return self
 
@@ -115,38 +122,38 @@ async def __aexit__(self, exc_type, exc_value, traceback) -> Any:
             exc_type is not None,
         )
         with translate_py_core_exceptions():
-            result = await self._native_connection.__aexit__(exc_type, exc_value, traceback)
+            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._native_connection.timeout
+            return self._py_core_async_connection.timeout
 
     @timeout.setter
     def timeout(self, value: int) -> None:
         with translate_py_core_exceptions():
-            self._native_connection.timeout = value
+            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._native_connection.autocommit
+            return self._py_core_async_connection.autocommit
 
     @property
     def closed(self) -> bool:
-        """Whether close has been initiated on the native connection."""
+        """Whether close has been initiated on the py-core async connection."""
         with translate_py_core_exceptions():
-            return self._native_connection.closed
+            return self._py_core_async_connection.closed
 
     def is_connected(self) -> bool:
-        """Return whether the native connection remains open."""
+        """Return whether the py-core async connection remains open."""
         with translate_py_core_exceptions():
-            return self._native_connection.is_connected()
+            return self._py_core_async_connection.is_connected()
 
     def __repr__(self) -> str:
         state = "closed" if self.closed else "connected"

mssql_python/async_query/async_cursor.py

@@ -6,9 +6,17 @@
     may change without notice.
 """
 
+import asyncio
+from collections.abc import Mapping, Sequence
+from contextlib import asynccontextmanager
 from typing import Any, Optional
+import uuid
 
+from ..exceptions import OperationalError
+from ..helpers import get_settings
 from ..logging import logger
+from ..row import Row
+from . import async_execute, async_fetch
 from .exception_translator import translate_py_core_exceptions
 
 
@@ -20,8 +28,100 @@ class AsyncCursor:
         Its signatures, behavior, error handling, and compatibility may change without notice.
     """
 
-    def __init__(self, native_cursor: Any) -> None:
-        self._native_cursor = native_cursor
+    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 isinstance(error, OperationalError) and str(error.__cause__).startswith(
+            "Connection is busy"
+        ):
+            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,
@@ -30,86 +130,101 @@ async def execute(
         use_prepare: bool = True,
         reset_cursor: bool = True,
     ) -> "AsyncCursor":
-        if len(parameters) == 1 and isinstance(parameters[0], (tuple, list)):
-            parameters = tuple(parameters[0])
-
-        logger.debug("AsyncCursor.execute: starting")
-        with translate_py_core_exceptions():
-            await self._native_cursor.execute(
+        async with self._result_transition():
+            return await async_execute.execute(
+                self,
                 operation,
                 *parameters,
                 use_prepare=use_prepare,
                 reset_cursor=reset_cursor,
             )
-        logger.debug("AsyncCursor.execute: completed")
-        return self
 
     async def executemany(
         self,
         operation: str,
-        seq_of_parameters: Any,
+        seq_of_parameters: Sequence[Sequence[Any]] | Sequence[Mapping[str, Any]],
         *,
         use_prepare: bool = True,
-    ) -> "AsyncCursor":
-        logger.debug("AsyncCursor.executemany: starting")
-        with translate_py_core_exceptions():
-            await self._native_cursor.executemany(
+    ) -> None:
+        async with self._result_transition():
+            await async_execute.executemany(
+                self,
                 operation,
                 seq_of_parameters,
                 use_prepare=use_prepare,
             )
-        logger.debug("AsyncCursor.executemany: completed")
-        return self
 
-    async def fetchone(self) -> Any:
-        with translate_py_core_exceptions():
-            return await self._native_cursor.fetchone()
+    async def fetchone(self) -> Row | None:
+        return await async_fetch.fetchone(self)
 
-    async def fetchmany(self, size: Optional[int] = None) -> Any:
-        with translate_py_core_exceptions():
-            if size is None:
-                return await self._native_cursor.fetchmany()
-            return await self._native_cursor.fetchmany(size)
+    async def fetchmany(self, size: Optional[int] = None) -> list[Row]:
+        return await async_fetch.fetchmany(self, size)
 
-    async def fetchall(self) -> Any:
-        with translate_py_core_exceptions():
-            return await self._native_cursor.fetchall()
+    async def fetchall(self) -> list[Row]:
+        return await async_fetch.fetchall(self)
 
     async def nextset(self) -> bool:
-        with translate_py_core_exceptions():
-            return await self._native_cursor.nextset()
+        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")
-        with translate_py_core_exceptions():
-            await self._native_cursor.close()
+        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 isinstance(error, OperationalError) and str(error.__cause__).startswith(
+                    "Connection is busy"
+                ):
+                    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._native_cursor.setinputsizes(sizes)
+            self._py_core_async_cursor.setinputsizes(sizes)
 
     @property
     def timeout(self) -> int:
         with translate_py_core_exceptions():
-            return self._native_cursor.timeout
+            return self._py_core_async_cursor.timeout
 
     @property
     def description(self) -> Any:
-        with translate_py_core_exceptions():
-            return self._native_cursor.description
+        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._native_cursor.rowcount
+            return self._py_core_async_cursor.rowcount
 
     @property
     def arraysize(self) -> int:
         with translate_py_core_exceptions():
-            return self._native_cursor.arraysize
+            return self._py_core_async_cursor.arraysize
 
     @arraysize.setter
     def arraysize(self, value: int) -> None:
         with translate_py_core_exceptions():
-            self._native_cursor.arraysize = value
+            self._py_core_async_cursor.arraysize = value