FIX: Preserve Arrow reader fetch exceptions (#712)

mssql_python/cursor.py

@@ -3042,54 +3042,52 @@ def batch_generator():
                 # body.  This is the single canonical cleanup site.
                 cur = cursor_ref[0]
                 cursor_ref[0] = None
-                if cur is None or cur.closed or cur.hstmt is None:
-                    return
-
-                # 1) Drain diagnostics produced by the (possibly cancelled)
-                #    fetch *before* SQL_CLOSE so we don't lose them.
-                try:
-                    cur.messages.extend(ddbc_bindings.DDBCSQLGetAllDiagRecords(cur.hstmt))
-                except Exception as e:  # pylint: disable=broad-exception-caught
-                    logger.debug("arrow_reader cleanup: pre-close diag drain failed: %s", e)
-
-                # 2) Release the server-side cursor & locks while keeping the
-                #    HSTMT and prepared plan intact, so the parent Cursor can
-                #    be re-executed.
-                try:
-                    cur.hstmt._close_cursor()  # pylint: disable=protected-access
-                except Exception as e:  # pylint: disable=broad-exception-caught
-                    # Elevated to WARNING: unlike the diag-drain failures
-                    # (which only cost us some warning text), a failed
-                    # SQLFreeStmt(SQL_CLOSE) leaves the server-side cursor
-                    # and its locks/tempdb resources open on SQL Server
-                    # until this parent Cursor is closed or re-executed.
-                    # DEBUG is typically disabled in production, so that
-                    # leak would be invisible; WARNING makes it visible.
-                    logger.warning(
-                        "arrow_reader cleanup: _close_cursor failed (%s); "
-                        "server-side cursor may remain open until this "
-                        "Cursor is closed or re-executed",
-                        e,
-                    )
-
-                # 3) Drain diagnostics produced by SQL_CLOSE itself.  This
-                #    runs unconditionally because SQL_CLOSE can return
-                #    SQL_SUCCESS_WITH_INFO (a *success* code) and still leave
-                #    warning records on the HSTMT diag stack; the previous
-                #    "only on failure" path would silently drop those.
-                try:
-                    cur.messages.extend(ddbc_bindings.DDBCSQLGetAllDiagRecords(cur.hstmt))
-                except Exception as e:  # pylint: disable=broad-exception-caught
-                    logger.debug("arrow_reader cleanup: post-close diag drain failed: %s", e)
+                if cur is not None and not cur.closed and cur.hstmt is not None:
+                    # 1) Drain diagnostics produced by the (possibly cancelled)
+                    #    fetch *before* SQL_CLOSE so we don't lose them.
+                    try:
+                        cur.messages.extend(ddbc_bindings.DDBCSQLGetAllDiagRecords(cur.hstmt))
+                    except Exception as e:  # pylint: disable=broad-exception-caught
+                        logger.debug("arrow_reader cleanup: pre-close diag drain failed: %s", e)
+
+                    # 2) Release the server-side cursor & locks while keeping the
+                    #    HSTMT and prepared plan intact, so the parent Cursor can
+                    #    be re-executed.
+                    try:
+                        cur.hstmt._close_cursor()  # pylint: disable=protected-access
+                    except Exception as e:  # pylint: disable=broad-exception-caught
+                        # Elevated to WARNING: unlike the diag-drain failures
+                        # (which only cost us some warning text), a failed
+                        # SQLFreeStmt(SQL_CLOSE) leaves the server-side cursor
+                        # and its locks/tempdb resources open on SQL Server
+                        # until this parent Cursor is closed or re-executed.
+                        # DEBUG is typically disabled in production, so that
+                        # leak would be invisible; WARNING makes it visible.
+                        logger.warning(
+                            "arrow_reader cleanup: _close_cursor failed (%s); "
+                            "server-side cursor may remain open until this "
+                            "Cursor is closed or re-executed",
+                            e,
+                        )
 
-                # 4) Reset cursor bookkeeping to a clean "no result set"
-                #    state.  rowcount becomes -1 to signal that the prior
-                #    result is no longer meaningful.
-                try:
-                    cur._clear_rownumber()  # pylint: disable=protected-access
-                    cur.rowcount = -1
-                except Exception as e:  # pylint: disable=broad-exception-caught
-                    logger.debug("arrow_reader cleanup: bookkeeping reset failed: %s", e)
+                    # 3) Drain diagnostics produced by SQL_CLOSE itself.  This
+                    #    runs unconditionally because SQL_CLOSE can return
+                    #    SQL_SUCCESS_WITH_INFO (a *success* code) and still leave
+                    #    warning records on the HSTMT diag stack; the previous
+                    #    "only on failure" path would silently drop those.
+                    try:
+                        cur.messages.extend(ddbc_bindings.DDBCSQLGetAllDiagRecords(cur.hstmt))
+                    except Exception as e:  # pylint: disable=broad-exception-caught
+                        logger.debug("arrow_reader cleanup: post-close diag drain failed: %s", e)
+
+                    # 4) Reset cursor bookkeeping to a clean "no result set"
+                    #    state.  rowcount becomes -1 to signal that the prior
+                    #    result is no longer meaningful.
+                    try:
+                        cur._clear_rownumber()  # pylint: disable=protected-access
+                        cur.rowcount = -1
+                    except Exception as e:  # pylint: disable=broad-exception-caught
+                        logger.debug("arrow_reader cleanup: bookkeeping reset failed: %s", e)
 
         gen = batch_generator()
         inner = pyarrow.RecordBatchReader.from_batches(schema, gen)

tests/test_004_cursor.py

@@ -10,7 +10,10 @@
 
 import pytest
 import os
+import subprocess
+import sys
 from datetime import datetime, date, time, timedelta, timezone
+from pathlib import Path
 import time as time_module
 import decimal
 from contextlib import closing
@@ -108,6 +111,21 @@
 ]
 
 
+@pytest.mark.skipif(sys.version_info < (3, 14), reason="PEP 765 warnings begin in Python 3.14")
+def test_cursor_compiles_with_warnings_as_errors():
+    """The driver source must compile when SyntaxWarning is promoted to an error."""
+    cursor_source = Path(__file__).parents[1] / "mssql_python" / "cursor.py"
+
+    result = subprocess.run(
+        [sys.executable, "-B", "-W", "error", "-m", "py_compile", str(cursor_source)],
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+
+    assert result.returncode == 0, result.stderr
+
+
 def drop_table_if_exists(cursor, table_name):
     """Drop the table if it exists"""
     try:

tests/test_004_cursor_arrow.py

@@ -688,6 +688,95 @@ def close(self):
 #     leak the server-side cursor or crash close()
 
 
+@pytest.mark.parametrize(
+    ("closed", "has_hstmt"),
+    [(True, True), (False, False)],
+    ids=["closed-cursor", "missing-hstmt"],
+)
+def test_arrow_reader_propagates_fetch_error_when_cleanup_is_skipped(closed, has_hstmt):
+    """A defensive cleanup guard must not turn a fetch error into end-of-stream."""
+
+    class FakeCursor:
+        def __init__(self):
+            self.closed = False
+            self.hstmt = object()
+            self.calls = 0
+
+        def _check_closed(self):
+            pass
+
+        def _ensure_pyarrow(self):
+            return pa
+
+        def arrow_batch(self, _batch_size):
+            self.calls += 1
+            if self.calls == 1:
+                return pa.record_batch([pa.array([], type=pa.int64())], names=["value"])
+
+            self.closed = closed
+            self.hstmt = object() if has_hstmt else None
+            raise RuntimeError("fetch failed")
+
+    fake_cursor = FakeCursor()
+    reader = mssql_python.Cursor.arrow_reader(fake_cursor, batch_size=1)
+    try:
+        with pytest.raises(RuntimeError, match="fetch failed"):
+            reader.read_next_batch()
+    finally:
+        reader.close()
+
+
+def test_arrow_reader_propagates_fetch_error_after_cleanup(monkeypatch):
+    """Fetch errors must survive the normal cleanup path, which must still run."""
+    from mssql_python import cursor as cursor_mod
+
+    class FakeHstmt:
+        def __init__(self):
+            self.close_calls = 0
+
+        def _cancel(self):
+            pass
+
+        def _close_cursor(self):
+            self.close_calls += 1
+
+    class FakeCursor:
+        def __init__(self):
+            self.closed = False
+            self.hstmt = FakeHstmt()
+            self.messages = []
+            self.rowcount = 1
+            self.calls = 0
+            self.rownumber_cleared = False
+
+        def _check_closed(self):
+            pass
+
+        def _ensure_pyarrow(self):
+            return pa
+
+        def _clear_rownumber(self):
+            self.rownumber_cleared = True
+
+        def arrow_batch(self, _batch_size):
+            self.calls += 1
+            if self.calls == 1:
+                return pa.record_batch([pa.array([], type=pa.int64())], names=["value"])
+            raise RuntimeError("fetch failed")
+
+    monkeypatch.setattr(cursor_mod.ddbc_bindings, "DDBCSQLGetAllDiagRecords", lambda _h: [])
+    fake_cursor = FakeCursor()
+    reader = mssql_python.Cursor.arrow_reader(fake_cursor, batch_size=1)
+    try:
+        with pytest.raises(RuntimeError, match="fetch failed"):
+            reader.read_next_batch()
+        assert fake_cursor.hstmt.close_calls == 1
+        assert fake_cursor.rownumber_cleared is True
+        assert fake_cursor.rowcount == -1
+    finally:
+        reader.close()
+
+
 def test_arrow_reader_getattr_refuses_private_names(cursor: mssql_python.Cursor):
     """__getattr__ refuses leading-underscore names so a partially-constructed
     instance during __del__ cannot recurse forever trying to resolve its own