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 86%

🎯 Overall Coverage 81%

📈 Total Lines Covered: 7158 out of 8788 📁 Project: mssql-python

---

Diff Coverage

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

• mssql_python/cursor.py (87.1%): Missing lines 140,172,189,207,213-214,252,2994,3000,3008,3018,3027
• mssql_python/pybind/ddbc_bindings.cpp (84.2%): Missing lines 1449,1624-1625

Summary

• Total: 112 lines
• Missing: 15 lines
• Coverage: 86%

mssql_python/cursor.py

Lines 136-144

  136     @property
  137     def schema(self):
  138         """Schema of the record batches produced by this reader."""
  139         if self._closed:
! 140             raise self._arrow_invalid("Reader is closed")
  141         return self._inner.schema
  142 
  143     def read_next_batch(self):
  144         if self._closed:

---

Lines 168-176

  168         # versions do not implement the protocol.  Fail explicitly rather
  169         # than silently returning something invalid.
  170         inner_export = getattr(self._inner, "__arrow_c_stream__", None)
  171         if inner_export is None:
! 172             raise self._arrow_invalid(
  173                 "Arrow PyCapsule Protocol requires pyarrow>=14; "
  174                 "the installed pyarrow version does not expose "
  175                 "RecordBatchReader.__arrow_c_stream__."
  176             )

---

Lines 185-193

  185         return self._inner.read_next_batch()
  186 
  187     def __enter__(self):
  188         if self._closed:
! 189             raise self._arrow_invalid("Reader is closed")
  190         return self
  191 
  192     def __exit__(self, exc_type, exc_val, exc_tb):
  193         self.close()

---

Lines 203-211

  203         try:
  204             import sys as _sys
  205 
  206             if _sys.is_finalizing():
! 207                 return
  208             # Retry whenever the generator is still referenced — covers both
  209             # "user never called close()" and "earlier close() raised before
  210             # the generator was released".
  211             if getattr(self, "_generator", None) is not None:

---

Lines 209-218

  209             # "user never called close()" and "earlier close() raised before
  210             # the generator was released".
  211             if getattr(self, "_generator", None) is not None:
  212                 self.close()
! 213         except Exception:  # pylint: disable=broad-exception-caught
! 214             pass
  215 
  216     # ── Close implementation ──────────────────────────────────────────────
  217 
  218     def close(self) -> None:

---

Lines 248-256

  248         cursor = self._cursor
  249         if cursor is not None and not cursor.closed and cursor.hstmt is not None:
  250             try:
  251                 cursor.hstmt._cancel()  # pylint: disable=protected-access
! 252             except Exception as e:  # pylint: disable=broad-exception-caught
  253                 logger.debug("arrow_reader.close: SQLCancel raised: %s", e)
  254 
  255         # Close the generator — this raises GeneratorExit inside it, which
  256         # runs the try/finally cleanup block (SQLFreeStmt + diag drain +

---

Lines 2990-2998

  2990                 # body.  This is the single canonical cleanup site.
  2991                 cur = cursor_ref[0]
  2992                 cursor_ref[0] = None
  2993                 if cur is None or cur.closed or cur.hstmt is None:
! 2994                     return
  2995 
  2996                 # 1) Drain diagnostics produced by the (possibly cancelled)
  2997                 #    fetch *before* SQL_CLOSE so we don't lose them.
  2998                 try:

---

Lines 2996-3004

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

---

Lines 3004-3012

  3004                 #    HSTMT and prepared plan intact, so the parent Cursor can
  3005                 #    be re-executed.
  3006                 try:
  3007                     cur.hstmt._close_cursor()  # pylint: disable=protected-access
! 3008                 except Exception as e:  # pylint: disable=broad-exception-caught
  3009                     logger.debug("arrow_reader cleanup: _close_cursor failed: %s", e)
  3010 
  3011                 # 3) Drain diagnostics produced by SQL_CLOSE itself.  This
  3012                 #    runs unconditionally because SQL_CLOSE can return

---

Lines 3014-3022

  3014                 #    warning records on the HSTMT diag stack; the previous
  3015                 #    "only on failure" path would silently drop those.
  3016                 try:
  3017                     cur.messages.extend(ddbc_bindings.DDBCSQLGetAllDiagRecords(cur.hstmt))
! 3018                 except Exception as e:  # pylint: disable=broad-exception-caught
  3019                     logger.debug("arrow_reader cleanup: post-close diag drain failed: %s", e)
  3020 
  3021                 # 4) Reset cursor bookkeeping to a clean "no result set"
  3022                 #    state.  rowcount becomes -1 to signal that the prior

---

Lines 3023-3031

  3023                 #    result is no longer meaningful.
  3024                 try:
  3025                     cur._clear_rownumber()  # pylint: disable=protected-access
  3026                     cur.rowcount = -1
! 3027                 except Exception as e:  # pylint: disable=broad-exception-caught
  3028                     logger.debug("arrow_reader cleanup: bookkeeping reset failed: %s", e)
  3029 
  3030         gen = batch_generator()
  3031         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 1620-1629

  1620     if (!SQLCancel_ptr) {
  1621         return;
  1622     }
  1623     SQLHANDLE h = _handle;
! 1624     SQLRETURN ret;
! 1625     {
  1626         py::gil_scoped_release release;
  1627         ret = SQLCancel_ptr(h);
  1628     }
  1629     // 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]

_ArrowReader has no internal lock, so two threads calling close() concurrently (or close() racing with __del__) can both pass the self._generator is not Noneguard and both callgen.close()/_cancel(). The individual native calls are guarded/idempotent so this is unlikely to corrupt state, but do you think we need to consider a lightweight threading.Lock` around the close body to make the documented cross-thread semantics fully robust?

[sumitmsft]

We are allowing cancel() to be called from a different thread than the one owning the statement (that's the cross-thread close() feature). But SqlHandle itself has no mutex.

Because of that, cancel() checks _handle / _implicitly_freed and then calls SQLCancel_ptr(h) with no lock. If another thread runs free() / close_cursor() in between, we can hit a use-after-free (calling SQLCancel on a handle that was just released), which may crash.

Since this PR is what makes concurrent access possible, could we protect the check + SQLCancel call with a lock, and have free() / close_cursor() take the same lock?

Alternatively, if we're confident free() can never race with cancel() in practice, a short comment explaining why it's safe would resolve the concern.

[sumitmsft]

In the cleanup path, each step is wrapped in a broad except ... logger.debug(...). That's fine for the diagnostic-drain steps - if those fail we just lose some warning text.

But _close_cursor() is different: if it fails, the server-side cursor (and its locks) leak on SQL Server, and over time that can cause real production issues. Right now that failure is logged at DEBUG, which is typically disabled in production (INFO/WARN), so the leak would be completely invisible - there'd be no signal that anything went wrong.

I suggest we raise just the _close_cursor() failure log to WARNING (a failed cursor release is operationally significant and should show up in normal logs). The diagnostic-drain failures can stay at DEBUG.

[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.