FIX: Implementing RecordBatchReader.Close() functionality for Arrow

[subrata-ms]

Work Item / Issue Reference

AB#45978

GitHub Issue: #643

---

Summary

This pull request introduces a new _ArrowReader wrapper to enhance the behavior of the arrow_reader method in the Cursor class. The new wrapper ensures that server-side resources are properly released when the reader is closed, supporting robust cleanup and allowing the parent cursor to remain usable. The tests are updated and extended to verify the improved semantics, including close behavior and context manager support.

Enhancements to Arrow batch reading and resource management:

• Added the _ArrowReader class to cursor.py, which wraps a pyarrow.RecordBatchReader and implements an 8-step close sequence to properly release server-side resources, reset cursor state, and support idempotent and context-manager-based cleanup. The parent Cursor remains usable after closing the reader.
• Modified the Cursor.arrow_reader method to return an instance of _ArrowReader instead of a raw pyarrow.RecordBatchReader, ensuring that closing the reader stops fetching, releases the server-side cursor, and resets cursor state. [1] [2]

Test improvements for Arrow reader behavior:

• Updated the test_arrow_reader test to check for duck-typed compatibility with pyarrow.RecordBatchReader, reflecting the new wrapper class.
• Added new tests: test_arrow_reader_close_semantics to verify that .close() stops fetching, marks the reader as closed, is idempotent, and leaves the parent cursor usable; and test_arrow_reader_context_manager to verify that the reader is closed on context manager exit and the cursor remains usable.

[github-actions]

📊 Code Coverage Report

🔥 Diff Coverage 84%

🎯 Overall Coverage 81%

📈 Total Lines Covered: 7170 out of 8798 📁 Project: mssql-python

---

Diff Coverage

Diff: main...HEAD, staged and unstaged changes

• mssql_python/cursor.py (86.7%): Missing lines 167,196,213,231,237-238,276,3043,3049,3057,3079,3088
• mssql_python/pybind/ddbc_bindings.cpp (79.4%): Missing lines 1449,1612-1613,1623-1624,1646-1647

Summary

• Total: 124 lines
• Missing: 19 lines
• Coverage: 84%

mssql_python/cursor.py

Lines 163-171

  163         methods — a reader that has been marked closed must not delegate
  164         even if a retry-pending state still holds ``self._inner``.
  165         """
  166         if name.startswith("_"):
! 167             raise AttributeError(name)
  168         if self._closed:
  169             raise self._arrow_invalid("Reader is closed")
  170         return getattr(self._inner, name)

---

Lines 192-200

  192         # versions do not implement the protocol.  Fail explicitly rather
  193         # than silently returning something invalid.
  194         inner_export = getattr(self._inner, "__arrow_c_stream__", None)
  195         if inner_export is None:
! 196             raise self._arrow_invalid(
  197                 "Arrow PyCapsule Protocol requires pyarrow>=14; "
  198                 "the installed pyarrow version does not expose "
  199                 "RecordBatchReader.__arrow_c_stream__."
  200             )

---

Lines 209-217

  209         return self._inner.read_next_batch()
  210 
  211     def __enter__(self):
  212         if self._closed:
! 213             raise self._arrow_invalid("Reader is closed")
  214         return self
  215 
  216     def __exit__(self, exc_type, exc_val, exc_tb):
  217         self.close()

---

Lines 227-235

  227         try:
  228             import sys as _sys
  229 
  230             if _sys.is_finalizing():
! 231                 return
  232             # Retry whenever the generator is still referenced — covers both
  233             # "user never called close()" and "earlier close() raised before
  234             # the generator was released".
  235             if getattr(self, "_generator", None) is not None:

---

Lines 233-242

  233             # "user never called close()" and "earlier close() raised before
  234             # the generator was released".
  235             if getattr(self, "_generator", None) is not None:
  236                 self.close()
! 237         except Exception:  # pylint: disable=broad-exception-caught
! 238             pass
  239 
  240     # ── Close implementation ──────────────────────────────────────────────
  241 
  242     def close(self) -> None:

---

Lines 272-280

  272         cursor = self._cursor
  273         if cursor is not None and not cursor.closed and cursor.hstmt is not None:
  274             try:
  275                 cursor.hstmt._cancel()  # pylint: disable=protected-access
! 276             except Exception as e:  # pylint: disable=broad-exception-caught
  277                 logger.debug("arrow_reader.close: SQLCancel raised: %s", e)
  278 
  279         # Close the generator — this raises GeneratorExit inside it, which
  280         # runs the try/finally cleanup block (SQLFreeStmt + diag drain +

---

Lines 3039-3047

  3039                 # body.  This is the single canonical cleanup site.
  3040                 cur = cursor_ref[0]
  3041                 cursor_ref[0] = None
  3042                 if cur is None or cur.closed or cur.hstmt is None:
! 3043                     return
  3044 
  3045                 # 1) Drain diagnostics produced by the (possibly cancelled)
  3046                 #    fetch *before* SQL_CLOSE so we don't lose them.
  3047                 try:

---

Lines 3045-3053

  3045                 # 1) Drain diagnostics produced by the (possibly cancelled)
  3046                 #    fetch *before* SQL_CLOSE so we don't lose them.
  3047                 try:
  3048                     cur.messages.extend(ddbc_bindings.DDBCSQLGetAllDiagRecords(cur.hstmt))
! 3049                 except Exception as e:  # pylint: disable=broad-exception-caught
  3050                     logger.debug("arrow_reader cleanup: pre-close diag drain failed: %s", e)
  3051 
  3052                 # 2) Release the server-side cursor & locks while keeping the
  3053                 #    HSTMT and prepared plan intact, so the parent Cursor can

---

Lines 3053-3061

  3053                 #    HSTMT and prepared plan intact, so the parent Cursor can
  3054                 #    be re-executed.
  3055                 try:
  3056                     cur.hstmt._close_cursor()  # pylint: disable=protected-access
! 3057                 except Exception as e:  # pylint: disable=broad-exception-caught
  3058                     # Elevated to WARNING: unlike the diag-drain failures
  3059                     # (which only cost us some warning text), a failed
  3060                     # SQLFreeStmt(SQL_CLOSE) leaves the server-side cursor
  3061                     # and its locks/tempdb resources open on SQL Server

---

Lines 3075-3083

  3075                 #    warning records on the HSTMT diag stack; the previous
  3076                 #    "only on failure" path would silently drop those.
  3077                 try:
  3078                     cur.messages.extend(ddbc_bindings.DDBCSQLGetAllDiagRecords(cur.hstmt))
! 3079                 except Exception as e:  # pylint: disable=broad-exception-caught
  3080                     logger.debug("arrow_reader cleanup: post-close diag drain failed: %s", e)
  3081 
  3082                 # 4) Reset cursor bookkeeping to a clean "no result set"
  3083                 #    state.  rowcount becomes -1 to signal that the prior

---

Lines 3084-3092

  3084                 #    result is no longer meaningful.
  3085                 try:
  3086                     cur._clear_rownumber()  # pylint: disable=protected-access
  3087                     cur.rowcount = -1
! 3088                 except Exception as e:  # pylint: disable=broad-exception-caught
  3089                     logger.debug("arrow_reader cleanup: bookkeeping reset failed: %s", e)
  3090 
  3091         gen = batch_generator()
  3092         inner = pyarrow.RecordBatchReader.from_batches(schema, gen)

---

mssql_python/pybind/ddbc_bindings.cpp

Lines 1445-1453

  1445     SQLEndTran_ptr = GetFunctionPointer<SQLEndTranFunc>(handle, "SQLEndTran");
  1446     SQLDisconnect_ptr = GetFunctionPointer<SQLDisconnectFunc>(handle, "SQLDisconnect");
  1447     SQLFreeHandle_ptr = GetFunctionPointer<SQLFreeHandleFunc>(handle, "SQLFreeHandle");
  1448     SQLFreeStmt_ptr = GetFunctionPointer<SQLFreeStmtFunc>(handle, "SQLFreeStmt");
! 1449     SQLCancel_ptr = GetFunctionPointer<SQLCancelFunc>(handle, "SQLCancel");
  1450 
  1451     SQLGetDiagRec_ptr = GetFunctionPointer<SQLGetDiagRecFunc>(handle, "SQLGetDiagRecW");
  1452 
  1453     SQLParamData_ptr = GetFunctionPointer<SQLParamDataFunc>(handle, "SQLParamData");

---

Lines 1608-1617

  1608     }
  1609 }
  1610 
  1611 void SqlHandle::cancel() {
! 1612     // SQLCancel is intentionally lenient: it is a no-op on non-STMT handles,
! 1613     // already-freed handles, or if the driver does not expose it. This lets
  1614     // _ArrowReader.close() call it unconditionally without coordinating with
  1615     // the fetch thread. The GIL is released so a blocked fetch thread can
  1616     // observe the cancel and return.
  1617     //

---

Lines 1619-1628

  1619     //   The only cross-thread pattern this driver blesses is exactly the one
  1620     //   ODBC blesses: cancel() may be called from a thread *other than* the
  1621     //   fetch thread to unblock an in-flight SQLFetch/SQLExecute on the same
  1622     //   HSTMT. Per the ODBC spec, SQLCancel (with the SQLGetDiagRec/Field
! 1623     //   family) is the only entry point safe to call across threads on the
! 1624     //   same statement handle. All other operations on a Cursor/SqlHandle
  1625     //   are single-owner: per DB API 2.0 and the Cursor thread-safety TODO
  1626     //   in cursor.py, callers must not share a Cursor for its lifecycle
  1627     //   operations (execute/fetch/close/free) across threads. Under that
  1628     //   contract, free() / close_cursor() / SQLFreeHandle can never be in

---

Lines 1642-1651

  1642     if (!SQLCancel_ptr) {
  1643         return;
  1644     }
  1645     SQLHANDLE h = _handle;
! 1646     SQLRETURN ret;
! 1647     {
  1648         py::gil_scoped_release release;
  1649         ret = SQLCancel_ptr(h);
  1650     }
  1651     // SQLCancel may return SQL_SUCCESS_WITH_INFO when there was nothing to

---

---

📋 Files Needing Attention

📉 Files with overall lowest coverage (click to expand)

mssql_python.pybind.logger_bridge.cpp: 59.2%
mssql_python.pybind.ddbc_bindings.h: 59.9%
mssql_python.pybind.logger_bridge.hpp: 70.8%
mssql_python.pybind.ddbc_bindings.cpp: 76.2%
mssql_python.__init__.py: 77.3%
mssql_python.row.py: 77.6%
mssql_python.ddbc_bindings.py: 79.6%
mssql_python.pybind.connection.connection_pool.cpp: 81.4%
mssql_python.pybind.connection.connection.cpp: 83.7%
mssql_python.connection.py: 84.7%

---

🔗 Quick Links

⚙️ Build Summary	📋 Coverage Details
View Azure DevOps Build	Browse Full Coverage Report

[Copilot]

Pull request overview

This PR improves Arrow streaming ergonomics by making Cursor.arrow_reader() return a RecordBatchReader-compatible wrapper whose .close() actually cancels in-flight fetches and releases server-side ODBC cursor resources, keeping the parent Cursor reusable.

Changes:

• Introduces _ArrowReader in cursor.py and updates Cursor.arrow_reader() to return it instead of a raw pyarrow.RecordBatchReader.
• Adds an ODBC SQLCancel binding and exposes it via SqlHandle._cancel() to support cross-thread cancellation during .close().
• Extends Arrow reader tests to validate close semantics, context-manager behavior, GC cleanup, and cross-thread cancellation.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 5 comments.

File	Description
mssql_python/cursor.py	Adds _ArrowReader and reworks arrow_reader() to drive robust cleanup/cancellation and cursor state reset.
mssql_python/pybind/ddbc_bindings.h	Declares the SQLCancel function pointer type and adds SqlHandle::cancel() API surface.
mssql_python/pybind/ddbc_bindings.cpp	Loads SQLCancel, implements SqlHandle::cancel() (releasing the GIL), and exposes _cancel to Python.
tests/test_004_cursor_arrow.py	Updates/expands tests to cover new reader wrapper behavior and resource release semantics.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[ffelixg]

Huh, I didn't know that arrow readers supported a close method! I think that _ArrowReader should be subclassing pyarrow.RecordBatchReader. That way, we won't have to worry about duck typing. It also might save on some implementation overhead. @subrata-ms

[subrata-ms]

Huh, I didn't know that arrow readers supported a close method! I think that _ArrowReader should be subclassing pyarrow.RecordBatchReader. That way, we won't have to worry about duck typing. It also might save on some implementation overhead. @subrata-ms

Thanks for the suggestion — I looked into it and unfortunately subclassing doesn't work cleanly here because pyarrow.RecordBatchReader is a Cython extension type, not a normal Python class. Specifically:

RecordBatchReader.from_batches(...) is a Cython factory that hard-codes the return type to the base class; Sub.from_batches(...) returns a plain RecordBatchReader, so overridden close() / read_next_batch() are never called.
class reassignment is rejected on Cython types (TypeError: class assignment only supported for mutable types or ModuleType subclasses), so downcasting a factory-produced reader isn't an option either.
A Cython instance also rejects arbitrary attributes, so self._cursor, self._generator, etc. can't live on the base instance directly.
The bare Sub.new(Sub) path leaves the underlying C stream unset — super().close() segfaults immediately (I reproduced it locally on pyarrow 24.0.0).
So there's no path where the object is simultaneously (a) an isinstance of pa.RecordBatchReader, (b) able to intercept close() for our SQLCancel + generator teardown, and (c) able to carry our wrapper state.

The current composition/duck-typing shape is intentional for that reason. Happy to add arrow_c_stream to the wrapper so consumers that follow the Arrow PyCapsule Protocol (polars, duckdb, RecordBatchReader.from_stream, etc.) can consume it without any isinstance check — that gives us the "no duck-typing worries" property without the subclass constraints.

[sumitmsft]

Previously arrow_reader() returned a real pyarrow.RecordBatchReader; it now returns _ArrowReader, which (for good reasons documented in the code) does not subclass pyarrow.RecordBatchReader. This is a genuine behavioral change and can silently break existing users in two ways:

• isinstance checks fail: Any consumer doing isinstance(reader, pyarrow.RecordBatchReader) will now get False. The new __arrow_c_stream__ protocol nicely covers polars/duckdb/from_stream, but a direct isinstance check has no way to benefit from it.

• The type hint is now inaccurate: it still advertises -> "pyarrow.RecordBatchReader" while returning _ArrowReader, which misleads users and type checkers. Pls check above.

So,here is that I was thinking as suggestions:

1. Update the return annotation to reflect reality
   def arrow_reader(self, batch_size: int = 8192) -> "_ArrowReader":
2. Add a changelog entry and a line in the docstring explicitly calling out: "this is a pyarrow-compatible reader, not a pyarrow.RecordBatchReader; consume it via the Arrow C-stream protocol rather than an isinstance check."
   Please confirm no internal callers rely on the exact pyarrow.RecordBatchReader type.
   Keeping the wrapper is the right call - this is just about making the contract change visible so users aren't caught off guard.

[subrata-ms]

fixed. Added a change-log

[bewithgaurav]

the return type hint might've gotten missed, added a separate comment to highlight
else lgtm

[bewithgaurav]

suggesting a better impl of the wrapper

[bewithgaurav]

need to change the return type annotation here to _ArrowReader

[bewithgaurav]

the return type hint might've gotten missed, added a separate comment to highlight
else lgtm