Skip to content

FIX: Preserve Arrow reader fetch exceptions (#712) #718

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Subrata (subrata-ms) wants to merge 7 commits into
base: main
Choose a base branch
from
+152 −47
Open

FIX: Preserve Arrow reader fetch exceptions (#712) #718

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
4a7fbb9
FIX: Preserve Arrow reader fetch exceptions (#712)
Aug 13, 2026
f12dacb
Merge branch 'main' into subrata-ms/Bug712
subrata-ms Aug 13, 2026
45e9a75
STYLE: Format Bug 712 cursor test
Aug 13, 2026
25bf778
Potential fix for pull request finding
subrata-ms Aug 13, 2026
0d633b2
Merge branch 'main' into subrata-ms/Bug712
subrata-ms Aug 13, 2026
0d1de40
Merge branch 'main' into subrata-ms/Bug712
subrata-ms Aug 17, 2026
c6a3456
Merge branch 'main' into subrata-ms/Bug712
subrata-ms Aug 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
92 changes: 45 additions & 47 deletions mssql_python/cursor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3088,54 +3088,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:

@bewithgaurav Gaurav Sharma (bewithgaurav) Aug 17, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: cur is never None here, it gets read on the line right above and this block runs once per reader
so if not cur.closed and cur.hstmt is not None: would do the same, only mentioning it since the line is already changing

# 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)
Expand Down
18 changes: 18 additions & 0 deletions tests/test_004_cursor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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"

@bewithgaurav Gaurav Sharma (bewithgaurav) Aug 17, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this checks cursor.py only, but what broke was importing the package
the same return in a finally in any other file breaks it the same way, and this test would stay green

looping over the package covers all of it:

package_dir = Path(__file__).parents[1] / "mssql_python"
for source in sorted(package_dir.glob("*.py")):
    with warnings.catch_warnings():
        warnings.simplefilter("error")
        compile(source.read_text(encoding="utf-8"), str(source), "exec")

ran it on 3.14 and 3.13, green on both. so the skipif can go and it covers every leg instead of just the 3.14 ones. needs import warnings, and subprocess becomes unused


result = subprocess.run(
[sys.executable, "-B", "-W", "error", "-m", "py_compile", str(cursor_source)],
capture_output=True,
text=True,
check=False,
)
Comment thread
Copilot marked this conversation as resolved.

assert result.returncode == 0, result.stderr


def drop_table_if_exists(cursor, table_name):
"""Drop the table if it exists"""
try:
Expand Down
89 changes: 89 additions & 0 deletions tests/test_004_cursor_arrow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Comment on lines +691 to +777

@bewithgaurav Gaurav Sharma (bewithgaurav) Aug 17, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

requesting changes on this one - since these two go green even when the driver is broken.
they never open a connection, so a real regression under the fake object still shows green here

the same bug is testable through the driver: close the cursor part way through a read and assert it raises. fails on the old guard, passes here, adding a suggestion:

Suggested change
@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()
_BIG_QUERY = (
"SELECT TOP (3000) ROW_NUMBER() OVER (ORDER BY (SELECT NULL)) AS n "
"FROM sys.all_objects a CROSS JOIN sys.all_objects b"
)
def test_arrow_reader_raises_when_cursor_closes_mid_stream(db_connection):
"""A cursor closed mid-stream must raise, not report a short result set."""
cur = db_connection.cursor()
cur.execute(_BIG_QUERY)
reader = cur.arrow_reader(batch_size=500)
rows = 0
with pytest.raises(mssql_python.Error):
for batch in reader:
rows += batch.num_rows
if rows >= 1000:
cur.close()
assert 0 < rows < 3000
def test_arrow_reader_raises_when_cursor_scope_already_exited(db_connection):
"""A reader outliving its cursor's `with` block must raise, not yield nothing."""
with db_connection.cursor() as cur:
cur.execute(_BIG_QUERY)
reader = cur.arrow_reader(batch_size=500)
with pytest.raises(mssql_python.Error):
for _ in reader:
pass



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
Expand Down
Loading