From 1758590f2c12ba46fc66e7ca75438256141b9351 Mon Sep 17 00:00:00 2001 From: Jahnvi Thakkar Date: Wed, 23 Sep 2026 14:52:26 +0530 Subject: [PATCH 1/6] PERF: Avoid redundant fetch diagnostic scans Capture fetch diagnostics before subsequent ODBC calls replace them, preserve data-read errors and Arrow cleanup behavior, and update existing regression coverage only. The tree is identical to tested revision c1808128312c956dbf108028a9c6fb750d209a51; broader warning/platform coverage and the faster-than-pyodbc goal remain unqualified. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- mssql_python/cursor.py | 63 +++-- mssql_python/pybind/ddbc_bindings.cpp | 373 +++++++++++++++++--------- mssql_python/pybind/ddbc_bindings.h | 30 ++- tests/test_004_cursor.py | 84 ++++-- tests/test_004_cursor_arrow.py | 7 + tests/test_fetch_settings_cache.py | 16 +- 6 files changed, 388 insertions(+), 185 deletions(-) diff --git a/mssql_python/cursor.py b/mssql_python/cursor.py index ad5eab798..37551d1e0 100644 --- a/mssql_python/cursor.py +++ b/mssql_python/cursor.py @@ -120,9 +120,9 @@ class _ArrowReader: the single ODBC entry point (with the diag-record functions) that the spec marks as safe to call from a different thread than the one owning the statement. - * Diagnostics are drained *before* the cursor is closed, so records - produced by a cancelled fetch are not lost; a second drain after - close picks up anything ``SQL_CLOSE`` itself emits. + * Cancellation/error diagnostics are drained before closing the cursor. + Natural exhaustion uses the native fetch's captured diagnostics; a + drain after close picks up anything ``SQL_CLOSE`` itself emits. * Cached ``pyarrow.ArrowInvalid`` avoids per-read imports on the post-close error path. * ``__del__`` is guarded against interpreter finalization. @@ -138,7 +138,14 @@ class _ArrowReader: The parent ``Cursor`` is **not** closed; it remains fully usable. """ - __slots__ = ("_cursor", "_inner", "_generator", "_closed", "_arrow_invalid") + __slots__ = ( + "_cursor", + "_inner", + "_generator", + "_closed", + "_arrow_invalid", + "_close_requested", + ) def __init__( self, @@ -146,11 +153,13 @@ def __init__( inner: "pyarrow.RecordBatchReader", generator, arrow_invalid_exc: type, + close_requested: list[bool], ) -> None: self._cursor = cursor self._inner = inner self._generator = generator self._closed = False + self._close_requested = close_requested # Cache the exception class so post-close reads in a hot loop don't # re-import pyarrow. self._arrow_invalid = arrow_invalid_exc @@ -285,6 +294,7 @@ def close(self) -> None: # Mark closed first so any racing read raises immediately, even if # the cleanup steps below fail and we end up retried later. self._closed = True + self._close_requested[0] = True # SQLCancel (cross-thread safe) — unblocks a fetch running on another # thread so that the generator's finally clause can then run @@ -1137,6 +1147,9 @@ def _capture_diagnostics(self, ret: int) -> None: driver's internal state when no records exist. SQL_ERROR is handled separately by check_error() which extracts diagnostics and raises. + + Composite native fetches capture intermediate diagnostics directly + into self.messages before subsequent ODBC calls can replace them. """ if self.hstmt and ret in ( ddbc_sql_const.SQL_SUCCESS_WITH_INFO.value, @@ -2835,13 +2848,10 @@ def fetchone(self) -> Union[None, Row]: char_enc, wchar_enc, self._cached_char_ctype, + self.messages, ) check_error(ddbc_sql_const.SQL_HANDLE_STMT.value, self.hstmt, ret) - with perf_phase("py::fetchone::diag_records"): - # The native bridge's final status can mask earlier fetch warnings. - if self.hstmt: - self.messages.extend(ddbc_bindings.DDBCSQLGetAllDiagRecords(self.hstmt)) if ret == ddbc_sql_const.SQL_NO_DATA.value: # No more data available @@ -2914,12 +2924,10 @@ def fetchmany(self, size: Optional[int] = None) -> List[Row]: char_enc, wchar_enc, self._cached_char_ctype, + self.messages, ) check_error(ddbc_sql_const.SQL_HANDLE_STMT.value, self.hstmt, ret) - with perf_phase("py::fetchmany::diag_records"): - if self.hstmt: - self.messages.extend(ddbc_bindings.DDBCSQLGetAllDiagRecords(self.hstmt)) # Update rownumber for the number of rows actually fetched if rows_data and self._has_result_set: @@ -2990,15 +2998,12 @@ def fetchall(self) -> List[Row]: char_enc, wchar_enc, self._cached_char_ctype, + self.messages, ) # Check for errors check_error(ddbc_sql_const.SQL_HANDLE_STMT.value, self.hstmt, ret) - with perf_phase("py::fetchall::diag_records"): - if self.hstmt: - self.messages.extend(ddbc_bindings.DDBCSQLGetAllDiagRecords(self.hstmt)) - # Update rownumber for the number of rows actually fetched if rows_data and self._has_result_set: self._next_row_index += len(rows_data) @@ -3062,15 +3067,12 @@ def arrow_batch(self, batch_size: int = 8192) -> "pyarrow.RecordBatch": char_decoding = self._get_decoding_settings(ddbc_sql_const.SQL_CHAR.value) char_c_type = char_decoding.get("ctype", ddbc_sql_const.SQL_WCHAR.value) ret = ddbc_bindings.DDBCSQLFetchArrowBatch( - self.hstmt, capsules, max(batch_size, 0), char_c_type + self.hstmt, capsules, max(batch_size, 0), char_c_type, self.messages ) check_error(ddbc_sql_const.SQL_HANDLE_STMT.value, self.hstmt, ret) batch = pyarrow.RecordBatch._import_from_c_capsule(*capsules) - if self.hstmt: - self.messages.extend(ddbc_bindings.DDBCSQLGetAllDiagRecords(self.hstmt)) - # Update rownumber for the number of rows actually fetched num_fetched = batch.num_rows if num_fetched > 0 and self._has_result_set: @@ -3145,11 +3147,14 @@ def arrow_reader(self, batch_size: int = 8192) -> "_ArrowReader": # can null out after cleanup, so a GC'd reader does not keep the # cursor pinned. cursor_ref = [self] + close_requested = [False] def batch_generator(): + exhausted = False try: while (batch := cursor_ref[0].arrow_batch(batch_size)).num_rows > 0: yield batch + exhausted = True finally: # Symmetric server-side teardown — runs on exhaustion, # GeneratorExit (from close()), or an exception inside the @@ -3157,12 +3162,13 @@ def batch_generator(): cur = cursor_ref[0] cursor_ref[0] = None if 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) + # Natural EOF was captured natively. A close/cancel request + # (including a racing one) or an error still needs the drain. + if not exhausted or close_requested[0]: + 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 @@ -3205,7 +3211,7 @@ def batch_generator(): gen = batch_generator() inner = pyarrow.RecordBatchReader.from_batches(schema, gen) - return _ArrowReader(self, inner, gen, pyarrow.ArrowInvalid) + return _ArrowReader(self, inner, gen, pyarrow.ArrowInvalid, close_requested) def nextset(self) -> Optional[bool]: """ @@ -4071,7 +4077,10 @@ def _execute_tables( # pylint: disable=too-many-arguments,too-many-positional-a check_error(ddbc_sql_const.SQL_HANDLE_STMT.value, stmt_handle, retcode) # Capture any diagnostic messages - if stmt_handle: + if stmt_handle and retcode in ( + ddbc_sql_const.SQL_SUCCESS_WITH_INFO.value, + ddbc_sql_const.SQL_NO_DATA.value, + ): self.messages.extend(ddbc_bindings.DDBCSQLGetAllDiagRecords(stmt_handle)) def tables( diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index 3cb568b4b..94fb28139 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -1896,21 +1896,8 @@ ErrorInfo SQLReadError(SQLSMALLINT handleType, SQLHANDLE rawHandle, SQLRETURN re return errorInfo; } -py::list SQLGetAllDiagRecords(SqlHandlePtr handle) { - PERF_TIMER("SQLGetAllDiagRecords"); - LOG("SQLGetAllDiagRecords: Retrieving all diagnostic records for handle " - "%p, handleType=%d", - (void*)handle->get(), handle->type()); - if (!SQLGetDiagRec_ptr) { - LOG("SQLGetAllDiagRecords: SQLGetDiagRec function pointer not " - "initialized, loading driver"); - DriverLoader::getInstance().loadDriver(); - } - - py::list records; - SQLHANDLE rawHandle = handle->get(); - SQLSMALLINT handleType = handle->type(); - +static void AppendDiagRecords(SQLHANDLE rawHandle, SQLSMALLINT handleType, py::handle records, + bool internalTruncation = false) { // Iterate through all available diagnostic records for (SQLSMALLINT recNumber = 1;; recNumber++) { SQLWCHAR sqlState[6] = {0}; @@ -1918,14 +1905,20 @@ py::list SQLGetAllDiagRecords(SqlHandlePtr handle) { SQLINTEGER nativeError = 0; SQLSMALLINT messageLen = 0; - SQLRETURN diagReturn = - SQLGetDiagRec_ptr(handleType, rawHandle, recNumber, sqlState, &nativeError, message, - SQL_MAX_MESSAGE_LENGTH_SQLSERVER, &messageLen); + SQLRETURN diagReturn; + { + PERF_TIMER("AppendDiagRecords::SQLGetDiagRec_call"); + diagReturn = SQLGetDiagRec_ptr(handleType, rawHandle, recNumber, sqlState, &nativeError, + message, SQL_MAX_MESSAGE_LENGTH_SQLSERVER, &messageLen); + } if (diagReturn == SQL_NO_DATA || !SQL_SUCCEEDED(diagReturn)) break; std::u16string sqlStateUtf16 = dupeSqlWCharAsUtf16Le(sqlState, 5); + // A continuation/probe can also carry unrelated warnings; filter each record. + if (internalTruncation && sqlStateUtf16 == u"01004") + continue; std::u16string messageUtf16 = dupeSqlWCharAsUtf16Le( message, std::min(static_cast(messageLen), static_cast(SQL_MAX_MESSAGE_LENGTH_SQLSERVER - 1))); @@ -1937,12 +1930,40 @@ py::list SQLGetAllDiagRecords(SqlHandlePtr handle) { std::string stateWithError = "[" + stateStr + "] (" + std::to_string(nativeError) + ")"; // Create the tuple with converted strings - records.append(py::make_tuple(py::str(stateWithError), py::str(msgStr))); + py::tuple record = py::make_tuple(py::str(stateWithError), py::str(msgStr)); + if (PyList_Append(records.ptr(), record.ptr()) < 0) + throw py::error_already_set(); } +} +py::list SQLGetAllDiagRecords(SqlHandlePtr handle) { + PERF_TIMER("SQLGetAllDiagRecords"); + LOG("SQLGetAllDiagRecords: Retrieving all diagnostic records for handle " + "%p, handleType=%d", + (void*)handle->get(), handle->type()); + if (!SQLGetDiagRec_ptr) { + LOG("SQLGetAllDiagRecords: SQLGetDiagRec function pointer not " + "initialized, loading driver"); + DriverLoader::getInstance().loadDriver(); + } + py::list records; + AppendDiagRecords(handle->get(), handle->type(), records); return records; } +// Called only with the GIL held, immediately after the originating ODBC call. +static void CaptureFetchDiagnostics(SQLHSTMT hStmt, SQLRETURN ret, py::handle messages, + bool internalTruncation = false) { + if ((ret == SQL_SUCCESS_WITH_INFO || ret == SQL_NO_DATA) && messages && !messages.is_none()) + AppendDiagRecords(hStmt, SQL_HANDLE_STMT, messages, internalTruncation); +} + +static void CheckFetchError(const SqlHandlePtr& handle, SQLRETURN ret) { + if (ret < 0) + py::module_::import("mssql_python.helpers") + .attr("check_error")(SQL_HANDLE_STMT, handle, ret); +} + // Wrap SQLExecDirect SQLRETURN SQLExecDirect_wrap(SqlHandlePtr StatementHandle, const std::u16string& Query) { PERF_TIMER("SQLExecDirect_wrap"); @@ -2991,7 +3012,7 @@ SQLRETURN SQLExecuteMany_wrap(const SqlHandlePtr statementHandle, const std::u16 } // Wrap SQLNumResultCols -SQLSMALLINT SQLNumResultCols_wrap(SqlHandlePtr statementHandle) { +SQLSMALLINT SQLNumResultCols_wrap(SqlHandlePtr statementHandle, py::handle messages = {}) { PERF_TIMER("SQLNumResultCols_wrap"); LOG("SQLNumResultCols: Getting number of columns in result set for " "statement_handle=%p", @@ -3004,12 +3025,14 @@ SQLSMALLINT SQLNumResultCols_wrap(SqlHandlePtr statementHandle) { SQLSMALLINT columnCount; // TODO: Handle the return code - SQLNumResultCols_ptr(statementHandle->get(), &columnCount); + SQLRETURN ret = SQLNumResultCols_ptr(statementHandle->get(), &columnCount); + CaptureFetchDiagnostics(statementHandle->get(), ret, messages); return columnCount; } // Wrap SQLDescribeCol -SQLRETURN SQLDescribeCol_wrap(SqlHandlePtr StatementHandle, py::list& ColumnMetadata) { +SQLRETURN SQLDescribeCol_wrap(SqlHandlePtr StatementHandle, py::list& ColumnMetadata, + py::handle messages = {}) { PERF_TIMER("SQLDescribeCol_wrap"); LOG("SQLDescribeCol: Getting column descriptions for statement_handle=%p", (void*)StatementHandle->get()); @@ -3020,6 +3043,7 @@ SQLRETURN SQLDescribeCol_wrap(SqlHandlePtr StatementHandle, py::list& ColumnMeta SQLSMALLINT ColumnCount; SQLRETURN retcode = SQLNumResultCols_ptr(StatementHandle->get(), &ColumnCount); + CaptureFetchDiagnostics(StatementHandle->get(), retcode, messages); if (!SQL_SUCCEEDED(retcode)) { LOG("SQLDescribeCol: Failed to get number of columns - SQLRETURN=%d", retcode); return retcode; @@ -3036,6 +3060,7 @@ SQLRETURN SQLDescribeCol_wrap(SqlHandlePtr StatementHandle, py::list& ColumnMeta retcode = SQLDescribeCol_ptr(StatementHandle->get(), i, ColumnName, sizeof(ColumnName) / sizeof(SQLWCHAR), &NameLength, &DataType, &ColumnSize, &DecimalDigits, &Nullable); + CaptureFetchDiagnostics(StatementHandle->get(), retcode, messages); if (SQL_SUCCEEDED(retcode)) { // Append a named py::dict to ColumnMetadata @@ -3092,7 +3117,8 @@ SQLRETURN SQLFetch_wrap(SqlHandlePtr StatementHandle) { // Non-static so it can be called from inline functions in header py::object FetchLobColumnData(SQLHSTMT hStmt, SQLUSMALLINT colIndex, SQLSMALLINT cType, - bool isWideChar, bool isBinary, const std::string& charEncoding) { + bool isWideChar, bool isBinary, const std::string& charEncoding, + py::handle messages) { PERF_TIMER("FetchLobColumnData"); std::vector buffer; SQLRETURN ret = SQL_SUCCESS_WITH_INFO; @@ -3107,6 +3133,7 @@ py::object FetchLobColumnData(SQLHSTMT hStmt, SQLUSMALLINT colIndex, SQLSMALLINT py::gil_scoped_release release; ret = SQLGetData_ptr(hStmt, colIndex, cType, chunk.data(), DAE_CHUNK_SIZE, &actualRead); } + CaptureFetchDiagnostics(hStmt, ret, messages, true); if (ret == SQL_ERROR || !SQL_SUCCEEDED(ret) && ret != SQL_SUCCESS_WITH_INFO) { std::ostringstream oss; @@ -3274,7 +3301,7 @@ static inline bool IsLobOrVariantColumn(SQLSMALLINT dataType, SQLULEN columnSize SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, py::list& row, const std::string& charEncoding = "utf-16le", const std::string& wcharEncoding = "utf-16le", - int charCtype = SQL_C_WCHAR) { + int charCtype = SQL_C_WCHAR, py::handle messages = {}) { PERF_TIMER("SQLGetData_wrap"); // Note: wcharEncoding parameter is reserved for future use // Currently WCHAR data always uses UTF-16LE for Windows compatibility @@ -3302,6 +3329,7 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p ret = SQLDescribeCol_ptr(hStmt, i, columnName, sizeof(columnName) / sizeof(SQLWCHAR), &columnNameLen, &dataType, &columnSize, &decimalDigits, &nullable); + CaptureFetchDiagnostics(hStmt, ret, messages); if (!SQL_SUCCEEDED(ret)) { LOG("SQLGetData: Error retrieving metadata for column %d - " "SQLDescribeCol SQLRETURN=%d", @@ -3321,6 +3349,7 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p // Without this probe call, SQLColAttribute returns incorrect type codes. SQLLEN indicator; ret = SQLGetData_ptr(hStmt, i, SQL_C_BINARY, NULL, 0, &indicator); + CaptureFetchDiagnostics(hStmt, ret, messages, true); if (!SQL_SUCCEEDED(ret)) { LOG_ERROR("SQLGetData: Failed to probe sql_variant column %d - SQLRETURN=%d", i, ret); @@ -3335,6 +3364,7 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p SQLLEN variantCType = 0; ret = SQLColAttribute_ptr(hStmt, i, SQL_CA_SS_VARIANT_TYPE, NULL, 0, NULL, &variantCType); + CaptureFetchDiagnostics(hStmt, ret, messages); if (!SQL_SUCCEEDED(ret)) { LOG_ERROR("SQLGetData: Failed to get sql_variant underlying type for column %d", i); row.append(py::none()); @@ -3380,11 +3410,11 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p "- columnSize=%lu", i, useWideChar ? "SQL_C_WCHAR" : "SQL_C_CHAR", (unsigned long)columnSize); if (useWideChar) { - row.append( - FetchLobColumnData(hStmt, i, SQL_C_WCHAR, true, false, "utf-16le")); + row.append(FetchLobColumnData(hStmt, i, SQL_C_WCHAR, true, false, + "utf-16le", messages)); } else { row.append(FetchLobColumnData(hStmt, i, SQL_C_CHAR, false, false, - effectiveCharEnc)); + effectiveCharEnc, messages)); } } else if (useWideChar) { // Wide-char path: fetch VARCHAR data as SQL_C_WCHAR @@ -3394,6 +3424,10 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p SQLLEN dataLen; ret = SQLGetData_ptr(hStmt, i, SQL_C_WCHAR, dataBuffer.data(), fetchBufferSize, &dataLen); + CaptureFetchDiagnostics( + hStmt, ret, messages, + ret == SQL_SUCCESS_WITH_INFO && + (dataLen == SQL_NO_TOTAL || dataLen >= fetchBufferSize)); if (SQL_SUCCEEDED(ret)) { if (dataLen > 0) { uint64_t numCharsInData = dataLen / sizeof(SQLWCHAR); @@ -3415,7 +3449,7 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p "truncated, using streaming LOB", i); row.append(FetchLobColumnData(hStmt, i, SQL_C_WCHAR, true, false, - "utf-16le")); + "utf-16le", messages)); } } else if (dataLen == SQL_NULL_DATA) { LOG("SQLGetData: Column %d is NULL (CHAR via WCHAR)", i); @@ -3430,8 +3464,8 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p LOG("SQLGetData: SQL_NO_TOTAL for column %d (CHAR via WCHAR), " "streaming via FetchLobColumnData", i); - row.append( - FetchLobColumnData(hStmt, i, SQL_C_WCHAR, true, false, "utf-16le")); + row.append(FetchLobColumnData(hStmt, i, SQL_C_WCHAR, true, false, + "utf-16le", messages)); } else if (dataLen < 0) { LOG("SQLGetData: Unexpected negative data length " "for column %d - dataType=%d, dataLen=%ld", @@ -3472,6 +3506,10 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p SQLLEN dataLen; ret = SQLGetData_ptr(hStmt, i, SQL_C_CHAR, dataBuffer.data(), dataBuffer.size(), &dataLen); + CaptureFetchDiagnostics( + hStmt, ret, messages, + ret == SQL_SUCCESS_WITH_INFO && + (dataLen == SQL_NO_TOTAL || dataLen >= dataBuffer.size())); if (SQL_SUCCEEDED(ret)) { // columnSize is in chars, dataLen is in bytes if (dataLen > 0) { @@ -3504,7 +3542,7 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p "(buffer_size=%zu), using streaming LOB", i, dataBuffer.size()); row.append(FetchLobColumnData(hStmt, i, SQL_C_CHAR, false, false, - effectiveCharEnc)); + effectiveCharEnc, messages)); } } else if (dataLen == SQL_NULL_DATA) { LOG("SQLGetData: Column %d is NULL (CHAR)", i); @@ -3520,7 +3558,7 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p "streaming via FetchLobColumnData", i); row.append(FetchLobColumnData(hStmt, i, SQL_C_CHAR, false, false, - effectiveCharEnc)); + effectiveCharEnc, messages)); } else if (dataLen < 0) { LOG("SQLGetData: Unexpected negative data length " "for column %d - dataType=%d, dataLen=%ld", @@ -3543,7 +3581,8 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p } case SQL_SS_XML: { LOG("SQLGetData: Streaming XML for column %d", i); - row.append(FetchLobColumnData(hStmt, i, SQL_C_WCHAR, true, false, "utf-16le")); + row.append( + FetchLobColumnData(hStmt, i, SQL_C_WCHAR, true, false, "utf-16le", messages)); break; } case SQL_WCHAR: @@ -3553,7 +3592,8 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p LOG("SQLGetData: Streaming LOB for column %d (SQL_C_WCHAR) " "- columnSize=%lu", i, (unsigned long)columnSize); - row.append(FetchLobColumnData(hStmt, i, SQL_C_WCHAR, true, false, "utf-16le")); + row.append(FetchLobColumnData(hStmt, i, SQL_C_WCHAR, true, false, "utf-16le", + messages)); } else { uint64_t fetchBufferSize = (columnSize + 1) * sizeof(SQLWCHAR); // +1 for null terminator @@ -3561,6 +3601,10 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p SQLLEN dataLen; ret = SQLGetData_ptr(hStmt, i, SQL_C_WCHAR, dataBuffer.data(), fetchBufferSize, &dataLen); + CaptureFetchDiagnostics( + hStmt, ret, messages, + ret == SQL_SUCCESS_WITH_INFO && + (dataLen == SQL_NO_TOTAL || dataLen >= fetchBufferSize)); if (SQL_SUCCEEDED(ret)) { if (dataLen > 0) { uint64_t numCharsInData = dataLen / sizeof(SQLWCHAR); @@ -3582,7 +3626,7 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p "truncated, using streaming LOB", i); row.append(FetchLobColumnData(hStmt, i, SQL_C_WCHAR, true, false, - "utf-16le")); + "utf-16le", messages)); } } else if (dataLen == SQL_NULL_DATA) { LOG("SQLGetData: Column %d is NULL (NVARCHAR)", i); @@ -3597,8 +3641,8 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p LOG("SQLGetData: SQL_NO_TOTAL for column %d (NVARCHAR), " "streaming via FetchLobColumnData", i); - row.append( - FetchLobColumnData(hStmt, i, SQL_C_WCHAR, true, false, "utf-16le")); + row.append(FetchLobColumnData(hStmt, i, SQL_C_WCHAR, true, false, + "utf-16le", messages)); } else if (dataLen < 0) { LOG("SQLGetData: Unexpected negative data length " "for column %d (NVARCHAR) - dataLen=%ld", @@ -3623,6 +3667,7 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p SQLINTEGER intValue; SQLLEN indicator = 0; ret = SQLGetData_ptr(hStmt, i, SQL_C_LONG, &intValue, 0, &indicator); + CaptureFetchDiagnostics(hStmt, ret, messages); if (SQL_SUCCEEDED(ret) && indicator != SQL_NULL_DATA) { row.append(static_cast(intValue)); } else { @@ -3634,6 +3679,7 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p SQLSMALLINT smallIntValue; SQLLEN indicator = 0; ret = SQLGetData_ptr(hStmt, i, SQL_C_SHORT, &smallIntValue, 0, &indicator); + CaptureFetchDiagnostics(hStmt, ret, messages); if (SQL_SUCCEEDED(ret) && indicator == SQL_NULL_DATA) { row.append(py::none()); break; @@ -3652,6 +3698,7 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p SQLREAL realValue; SQLLEN indicator = 0; ret = SQLGetData_ptr(hStmt, i, SQL_C_FLOAT, &realValue, 0, &indicator); + CaptureFetchDiagnostics(hStmt, ret, messages); if (SQL_SUCCEEDED(ret) && indicator == SQL_NULL_DATA) { row.append(py::none()); break; @@ -3673,6 +3720,7 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p ret = SQLGetData_ptr(hStmt, i, SQL_C_CHAR, numericStr, sizeof(numericStr), &indicator); + CaptureFetchDiagnostics(hStmt, ret, messages); if (SQL_SUCCEEDED(ret)) { try { @@ -3730,6 +3778,7 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p SQLDOUBLE doubleValue; SQLLEN indicator = 0; ret = SQLGetData_ptr(hStmt, i, SQL_C_DOUBLE, &doubleValue, 0, &indicator); + CaptureFetchDiagnostics(hStmt, ret, messages); if (SQL_SUCCEEDED(ret) && indicator == SQL_NULL_DATA) { row.append(py::none()); break; @@ -3748,6 +3797,7 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p SQLBIGINT bigintValue; SQLLEN indicator = 0; ret = SQLGetData_ptr(hStmt, i, SQL_C_SBIGINT, &bigintValue, 0, &indicator); + CaptureFetchDiagnostics(hStmt, ret, messages); if (SQL_SUCCEEDED(ret) && indicator == SQL_NULL_DATA) { row.append(py::none()); break; @@ -3767,6 +3817,7 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p SQLLEN indicator = 0; ret = SQLGetData_ptr(hStmt, i, SQL_C_TYPE_DATE, &dateValue, sizeof(dateValue), &indicator); + CaptureFetchDiagnostics(hStmt, ret, messages); if (SQL_SUCCEEDED(ret) && indicator != SQL_NULL_DATA) { row.append( FetchTemporal::date(dateValue.year, dateValue.month, dateValue.day)); @@ -3780,6 +3831,7 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p SQL_SS_TIME2_STRUCT t2 = {}; SQLLEN indicator = 0; ret = SQLGetData_ptr(hStmt, i, SQL_C_SS_TIME2, &t2, sizeof(t2), &indicator); + CaptureFetchDiagnostics(hStmt, ret, messages); if (SQL_SUCCEEDED(ret) && indicator != SQL_NULL_DATA) { row.append(FetchTemporal::time( t2.hour, t2.minute, t2.second, t2.fraction / 1000)); // ns to µs @@ -3800,6 +3852,7 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p SQLLEN indicator = 0; ret = SQLGetData_ptr(hStmt, i, SQL_C_TYPE_TIMESTAMP, ×tampValue, sizeof(timestampValue), &indicator); + CaptureFetchDiagnostics(hStmt, ret, messages); if (SQL_SUCCEEDED(ret) && indicator == SQL_NULL_DATA) { row.append(py::none()); break; @@ -3823,6 +3876,7 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p SQLLEN indicator; ret = SQLGetData_ptr(hStmt, i, SQL_C_SS_TIMESTAMPOFFSET, &dtoValue, sizeof(dtoValue), &indicator); + CaptureFetchDiagnostics(hStmt, ret, messages); if (SQL_SUCCEEDED(ret) && indicator != SQL_NULL_DATA) { LOG("SQLGetData: Retrieved DATETIMEOFFSET for column %d - " "%d-%d-%d %d:%d:%d, fraction_ns=%u, tz_hour=%d, " @@ -3867,13 +3921,17 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p LOG("SQLGetData: Streaming LOB for column %d " "(SQL_C_BINARY) - columnSize=%lu", i, (unsigned long)columnSize); - row.append(FetchLobColumnData(hStmt, i, SQL_C_BINARY, false, true, "")); + row.append( + FetchLobColumnData(hStmt, i, SQL_C_BINARY, false, true, "", messages)); } else { // Small VARBINARY, fetch directly std::vector dataBuffer(columnSize); SQLLEN dataLen; ret = SQLGetData_ptr(hStmt, i, SQL_C_BINARY, dataBuffer.data(), columnSize, &dataLen); + CaptureFetchDiagnostics(hStmt, ret, messages, + ret == SQL_SUCCESS_WITH_INFO && + (dataLen == SQL_NO_TOTAL || dataLen > columnSize)); if (SQL_SUCCEEDED(ret)) { if (dataLen > 0) { @@ -3881,8 +3939,8 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p row.append(py::bytes( reinterpret_cast(dataBuffer.data()), dataLen)); } else { - row.append( - FetchLobColumnData(hStmt, i, SQL_C_BINARY, false, true, "")); + row.append(FetchLobColumnData(hStmt, i, SQL_C_BINARY, false, true, + "", messages)); } } else if (dataLen == SQL_NULL_DATA) { row.append(py::none()); @@ -3909,6 +3967,7 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p SQLCHAR tinyIntValue; SQLLEN indicator = 0; ret = SQLGetData_ptr(hStmt, i, SQL_C_TINYINT, &tinyIntValue, 0, &indicator); + CaptureFetchDiagnostics(hStmt, ret, messages); if (SQL_SUCCEEDED(ret) && indicator == SQL_NULL_DATA) { row.append(py::none()); break; @@ -3927,6 +3986,7 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p SQLCHAR bitValue; SQLLEN indicator = 0; ret = SQLGetData_ptr(hStmt, i, SQL_C_BIT, &bitValue, 0, &indicator); + CaptureFetchDiagnostics(hStmt, ret, messages); if (SQL_SUCCEEDED(ret) && indicator == SQL_NULL_DATA) { row.append(py::none()); break; @@ -3947,6 +4007,7 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p SQLLEN indicator; ret = SQLGetData_ptr(hStmt, i, SQL_C_GUID, &guidValue, sizeof(guidValue), &indicator); + CaptureFetchDiagnostics(hStmt, ret, messages); if (SQL_SUCCEEDED(ret) && indicator != SQL_NULL_DATA) { std::vector guid_bytes(16); @@ -3983,6 +4044,8 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p ThrowStdException(errorString.str()); break; } + if (ret < 0) + return ret; } return ret; } @@ -4025,7 +4088,8 @@ SQLRETURN SQLFetchScroll_wrap(SqlHandlePtr StatementHandle, SQLSMALLINT FetchOri // For column in the result set, binds a buffer to retrieve column data // TODO: Move to anonymous namespace, since it is not used outside this file SQLRETURN SQLBindColums(SQLHSTMT hStmt, ColumnBuffers& buffers, py::list& columnNames, - SQLUSMALLINT numCols, int fetchSize, int charCtype = SQL_C_WCHAR) { + SQLUSMALLINT numCols, int fetchSize, int charCtype = SQL_C_WCHAR, + py::handle messages = {}) { PERF_TIMER("SQLBindColums"); SQLRETURN ret = SQL_SUCCESS; const bool useWideChar = (charCtype == SQL_C_WCHAR); @@ -4182,6 +4246,7 @@ SQLRETURN SQLBindColums(SQLHSTMT hStmt, ColumnBuffers& buffers, py::list& column ThrowStdException(errorString.str()); return ret; } + CaptureFetchDiagnostics(hStmt, ret, messages); } return ret; } @@ -4191,8 +4256,8 @@ SQLRETURN SQLBindColums(SQLHSTMT hStmt, ColumnBuffers& buffers, py::list& column SQLRETURN FetchBatchData(SQLHSTMT hStmt, ColumnBuffers& buffers, py::list& columnNames, py::list& rows, SQLUSMALLINT numCols, SQLULEN& numRowsFetched, const std::vector& lobColumns, - const std::string& charEncoding = "utf-16le", - int charCtype = SQL_C_WCHAR) { + const std::string& charEncoding = "utf-16le", int charCtype = SQL_C_WCHAR, + py::handle messages = {}) { PERF_TIMER("FetchBatchData"); LOG("FetchBatchData: Fetching data in batches"); SQLRETURN ret; @@ -4202,6 +4267,7 @@ SQLRETURN FetchBatchData(SQLHSTMT hStmt, ColumnBuffers& buffers, py::list& colum PERF_TIMER("FetchBatchData::SQLFetchScroll_call"); ret = SQLFetchScroll_ptr(hStmt, SQL_FETCH_NEXT, 0); } + CaptureFetchDiagnostics(hStmt, ret, messages); if (ret == SQL_NO_DATA) { LOG("FetchBatchData: No data to fetch"); return ret; @@ -4279,6 +4345,7 @@ SQLRETURN FetchBatchData(SQLHSTMT hStmt, ColumnBuffers& buffers, py::list& colum columnInfosExt[col].fetchBufferSize = columnInfos[col].fetchBufferSize; columnInfosExt[col].isLob = columnInfos[col].isLob; columnInfosExt[col].charEncoding = effectiveCharEnc; + columnInfosExt[col].messages = messages.ptr(); columnInfosExt[col].isUtf8 = (effectiveCharEnc == "utf-8"); // Set useWideChar for SQL_CHAR/VARCHAR columns when charCtype is SQL_C_WCHAR SQLSMALLINT dt = columnInfos[col].dataType; @@ -4635,6 +4702,62 @@ size_t calculateRowSize(py::list& columnNames, SQLUSMALLINT numCols) { return rowSize; } +struct FetchStateGuard { + SqlHandlePtr handle; + PyObject* messages; // Borrowed from the enclosing native fetch call. + int cleanupStep = 0; + + FetchStateGuard(SqlHandlePtr handle, py::handle messages) + : handle(std::move(handle)), messages(messages.ptr()) {} + + void observe(SQLRETURN ret) const { + CaptureFetchDiagnostics(handle->get(), ret, py::handle(messages)); + CheckFetchError(handle, ret); + } + + void setRowArraySize(SQLULEN rowArraySize) const { + observe(SQLSetStmtAttr_ptr(handle->get(), SQL_ATTR_ROW_ARRAY_SIZE, + (SQLPOINTER)(intptr_t)rowArraySize, 0)); + } + + void configure(SQLULEN* numRowsFetched, SQLULEN rowArraySize) const { + setRowArraySize(rowArraySize); + observe(SQLSetStmtAttr_ptr(handle->get(), SQL_ATTR_ROWS_FETCHED_PTR, numRowsFetched, 0)); + } + + void close() { + while (cleanupStep < 3) { + SQLRETURN ret; + switch (cleanupStep++) { + case 0: + ret = SQLSetStmtAttr_ptr(handle->get(), SQL_ATTR_ROW_ARRAY_SIZE, (SQLPOINTER)1, + 0); + break; + case 1: + ret = SQLSetStmtAttr_ptr(handle->get(), SQL_ATTR_ROWS_FETCHED_PTR, nullptr, 0); + break; + default: + ret = SQLFreeStmt_ptr(handle->get(), SQL_UNBIND); + break; + } + observe(ret); + } + } + + ~FetchStateGuard() { + // Explicit close handles normal completion; unwinding must keep the primary exception. + while (cleanupStep < 3) { + try { + close(); + } catch (py::error_already_set& error) { + error.discard_as_unraisable("fetch state cleanup"); + } catch (const std::exception& error) { + LOG_ERROR("Fetch state cleanup failed: %s", error.what()); + } + } + } +}; + // FetchMany_wrap - Fetches multiple rows of data from the result set. // // @param StatementHandle: Handle to the statement from which data is to be @@ -4654,8 +4777,8 @@ size_t calculateRowSize(py::list& columnNames, SQLUSMALLINT numCols) { // during fetching, it throws a runtime error. SQLRETURN FetchMany_wrap(SqlHandlePtr StatementHandle, py::list& rows, int fetchSize, const std::string& charEncoding = "utf-16le", - const std::string& wcharEncoding = "utf-16le", - int charCtype = SQL_C_WCHAR) { + const std::string& wcharEncoding = "utf-16le", int charCtype = SQL_C_WCHAR, + py::handle messages = {}) { PERF_TIMER("FetchMany_wrap"); // Issue #531: upgrade SQL_C_CHAR + utf-8 to SQL_C_WCHAR on Windows so the // driver does lossless UTF-16 conversion instead of returning ACP bytes. @@ -4663,11 +4786,11 @@ SQLRETURN FetchMany_wrap(SqlHandlePtr StatementHandle, py::list& rows, int fetch SQLRETURN ret; SQLHSTMT hStmt = StatementHandle->get(); // Retrieve column count - SQLSMALLINT numCols = SQLNumResultCols_wrap(StatementHandle); + SQLSMALLINT numCols = SQLNumResultCols_wrap(StatementHandle, messages); // Retrieve column metadata py::list columnNames; - ret = SQLDescribeCol_wrap(StatementHandle, columnNames); + ret = SQLDescribeCol_wrap(StatementHandle, columnNames, messages); if (!SQL_SUCCEEDED(ret)) { LOG("FetchMany_wrap: Failed to get column descriptions - SQLRETURN=%d", ret); return ret; @@ -4697,14 +4820,16 @@ SQLRETURN FetchMany_wrap(SqlHandlePtr StatementHandle, py::list& rows, int fetch py::gil_scoped_release release; ret = SQLFetch_ptr(hStmt); } + CaptureFetchDiagnostics(hStmt, ret, messages); if (ret == SQL_NO_DATA) break; if (!SQL_SUCCEEDED(ret)) return ret; py::list row; - SQLGetData_wrap(StatementHandle, numCols, row, charEncoding, wcharEncoding, - charCtype); // <-- streams LOBs correctly + ret = SQLGetData_wrap(StatementHandle, numCols, row, charEncoding, wcharEncoding, + charCtype, messages); + CheckFetchError(StatementHandle, ret); rows.append(row); numRowsFetched++; } @@ -4713,30 +4838,26 @@ SQLRETURN FetchMany_wrap(SqlHandlePtr StatementHandle, py::list& rows, int fetch // Initialize column buffers ColumnBuffers buffers(numCols, fetchSize); + FetchStateGuard fetchStateGuard(StatementHandle, messages); // Bind columns - ret = SQLBindColums(hStmt, buffers, columnNames, numCols, fetchSize, charCtype); + ret = SQLBindColums(hStmt, buffers, columnNames, numCols, fetchSize, charCtype, messages); if (!SQL_SUCCEEDED(ret)) { LOG("FetchMany_wrap: Error when binding columns - SQLRETURN=%d", ret); return ret; } - SQLSetStmtAttr_ptr(hStmt, SQL_ATTR_ROW_ARRAY_SIZE, (SQLPOINTER)(intptr_t)fetchSize, 0); - SQLSetStmtAttr_ptr(hStmt, SQL_ATTR_ROWS_FETCHED_PTR, &numRowsFetched, 0); + fetchStateGuard.configure(&numRowsFetched, fetchSize); ret = FetchBatchData(hStmt, buffers, columnNames, rows, numCols, numRowsFetched, lobColumns, - charEncoding, charCtype); + charEncoding, charCtype, messages); + CheckFetchError(StatementHandle, ret); if (!SQL_SUCCEEDED(ret) && ret != SQL_NO_DATA) { LOG("FetchMany_wrap: Error when fetching data - SQLRETURN=%d", ret); return ret; } - // Reset attributes before returning to avoid using stack pointers later - SQLSetStmtAttr_ptr(hStmt, SQL_ATTR_ROW_ARRAY_SIZE, (SQLPOINTER)1, 0); - SQLSetStmtAttr_ptr(hStmt, SQL_ATTR_ROWS_FETCHED_PTR, NULL, 0); - - // Unbind columns to allow subsequent fetchone() calls to use SQLGetData - SQLFreeStmt_ptr(hStmt, SQL_UNBIND); + fetchStateGuard.close(); return ret; } @@ -4755,7 +4876,7 @@ SQLRETURN FetchMany_wrap(SqlHandlePtr StatementHandle, py::list& rows, int fetch // @return SQLRETURN: SQL_SUCCESS on success, or error code on failure template SQLRETURN GetDataVar(SQLHSTMT hStmt, SQLUSMALLINT colNumber, SQLSMALLINT cType, - std::vector& dataVec, SQLLEN* indicator) { + std::vector& dataVec, SQLLEN* indicator, py::handle messages) { size_t start = 0; size_t end = 0; @@ -4784,6 +4905,7 @@ SQLRETURN GetDataVar(SQLHSTMT hStmt, SQLUSMALLINT colNumber, SQLSMALLINT cType, hStmt, colNumber, cType, reinterpret_cast(dataVec.data() + start), sizeof(T) * (dataVec.size() - start), // Available buffer size from start position &localInd); + CaptureFetchDiagnostics(hStmt, ret, messages, true); // Handle NULL data if (localInd == SQL_NULL_DATA) { @@ -4832,26 +4954,6 @@ SQLRETURN GetDataVar(SQLHSTMT hStmt, SQLUSMALLINT colNumber, SQLSMALLINT cType, return SQL_SUCCESS; } -struct FetchStateGuard { - SQLHSTMT hStmt; - - FetchStateGuard(SQLHSTMT stmtHandle, SQLULEN* numRowsFetched, SQLULEN rowArraySize) - : hStmt(stmtHandle) { - SQLSetStmtAttr_ptr(hStmt, SQL_ATTR_ROW_ARRAY_SIZE, (SQLPOINTER)(intptr_t)rowArraySize, 0); - SQLSetStmtAttr_ptr(hStmt, SQL_ATTR_ROWS_FETCHED_PTR, numRowsFetched, 0); - } - - ~FetchStateGuard() { - SQLSetStmtAttr_ptr(hStmt, SQL_ATTR_ROW_ARRAY_SIZE, (SQLPOINTER)1, 0); - SQLSetStmtAttr_ptr(hStmt, SQL_ATTR_ROWS_FETCHED_PTR, NULL, 0); - SQLFreeStmt_ptr(hStmt, SQL_UNBIND); - } - - void setRowArraySize(SQLULEN rowArraySize) const { - SQLSetStmtAttr_ptr(hStmt, SQL_ATTR_ROW_ARRAY_SIZE, (SQLPOINTER)(intptr_t)rowArraySize, 0); - } -}; - int32_t days_from_civil(int y, int m, int d) { // Implements the "days_from_civil" algorithm by Howard Hinnant // Returns number of days since Unix epoch (1970-01-01) @@ -4863,9 +4965,8 @@ int32_t days_from_civil(int y, int m, int d) { return era * 146097 + static_cast(doe) - 719468; } -SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, - int arrowBatchSize, - int charCtype) { +SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, int arrowBatchSize, + int charCtype, py::handle messages = {}) { PERF_TIMER("FetchArrowBatch_wrap"); // Fetch narrow char data as SQL_C_CHAR if on Linux/macOS and configured by the user charCtype = EffectiveCharCtypeForFetch(charCtype, "utf-8"); @@ -4876,14 +4977,14 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, SQLRETURN ret; SQLHSTMT hStmt = StatementHandle->get(); // Retrieve column count - SQLSMALLINT numCols = SQLNumResultCols_wrap(StatementHandle); + SQLSMALLINT numCols = SQLNumResultCols_wrap(StatementHandle, messages); if (numCols <= 0) { ThrowStdException("No active result set. Cannot fetch Arrow batch."); } // Retrieve column metadata py::list columnNames; - ret = SQLDescribeCol_wrap(StatementHandle, columnNames); + ret = SQLDescribeCol_wrap(StatementHandle, columnNames, messages); if (!SQL_SUCCEEDED(ret)) { LOG("Failed to get column descriptions"); return ret; @@ -5056,17 +5157,18 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, // Initialize column buffers ColumnBuffers buffers(numCols, fetchSize); + SQLULEN numRowsFetched = 0; + FetchStateGuard fetchStateGuard(StatementHandle, messages); if (!hasLobColumns && fetchSize > 0) { - ret = SQLBindColums(hStmt, buffers, columnNames, numCols, fetchSize, charCtype); + ret = SQLBindColums(hStmt, buffers, columnNames, numCols, fetchSize, charCtype, messages); if (!SQL_SUCCEEDED(ret)) { LOG("Error when binding columns"); return ret; } } - SQLULEN numRowsFetched = 0; - FetchStateGuard fetchStateGuard(hStmt, &numRowsFetched, fetchSize); + fetchStateGuard.configure(&numRowsFetched, fetchSize); int idxRowArrow = 0; @@ -5081,6 +5183,7 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, py::gil_scoped_release release; ret = SQLFetch_ptr(hStmt); } + fetchStateGuard.observe(ret); if (ret == SQL_NO_DATA) { ret = SQL_SUCCESS; // Normal completion break; @@ -5108,7 +5211,8 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, case SQL_LONGVARBINARY: { ret = GetDataVar(hStmt, idxCol + 1, SQL_C_BINARY, buffers.charBuffers[idxCol], - buffers.indicators[idxCol].data()); + buffers.indicators[idxCol].data(), messages); + CheckFetchError(StatementHandle, ret); if (!SQL_SUCCEEDED(ret)) { LOG("Error fetching BINARY LOB for column %d", idxCol + 1); return ret; @@ -5121,7 +5225,8 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, if (charCtype == SQL_C_CHAR) { ret = GetDataVar(hStmt, idxCol + 1, SQL_C_CHAR, buffers.charBuffers[idxCol], - buffers.indicators[idxCol].data()); + buffers.indicators[idxCol].data(), messages); + CheckFetchError(StatementHandle, ret); if (!SQL_SUCCEEDED(ret)) { LOG("Error fetching CHAR LOB data for column %d", idxCol + 1); return ret; @@ -5136,7 +5241,8 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, case SQL_WLONGVARCHAR: { ret = GetDataVar(hStmt, idxCol + 1, SQL_C_WCHAR, buffers.wcharBuffers[idxCol], - buffers.indicators[idxCol].data()); + buffers.indicators[idxCol].data(), messages); + CheckFetchError(StatementHandle, ret); if (!SQL_SUCCEEDED(ret)) { LOG("Error fetching WCHAR LOB data for column %d", idxCol + 1); return ret; @@ -5148,6 +5254,7 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, ret = SQLGetData_ptr( hStmt, idxCol + 1, SQL_C_SLONG, buffers.intBuffers[idxCol].data(), sizeof(SQLINTEGER), buffers.indicators[idxCol].data()); + fetchStateGuard.observe(ret); if (!SQL_SUCCEEDED(ret)) { LOG("Error fetching SLONG data for column %d", idxCol + 1); return ret; @@ -5160,6 +5267,7 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, buffers.smallIntBuffers[idxCol].data(), sizeof(SQLSMALLINT), buffers.indicators[idxCol].data()); + fetchStateGuard.observe(ret); if (!SQL_SUCCEEDED(ret)) { LOG("Error fetching SSHORT data for column %d", idxCol + 1); return ret; @@ -5172,6 +5280,7 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, SQLGetData_ptr(hStmt, idxCol + 1, SQL_C_TINYINT, buffers.charBuffers[idxCol].data(), sizeof(SQLCHAR), buffers.indicators[idxCol].data()); + fetchStateGuard.observe(ret); if (!SQL_SUCCEEDED(ret)) { LOG("Error fetching TINYINT data for column %d", idxCol + 1); return ret; @@ -5183,6 +5292,7 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, ret = SQLGetData_ptr( hStmt, idxCol + 1, SQL_C_BIT, buffers.charBuffers[idxCol].data(), sizeof(SQLCHAR), buffers.indicators[idxCol].data()); + fetchStateGuard.observe(ret); if (!SQL_SUCCEEDED(ret)) { LOG("Error fetching BIT data for column %d", idxCol + 1); return ret; @@ -5194,6 +5304,7 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, ret = SQLGetData_ptr( hStmt, idxCol + 1, SQL_C_FLOAT, buffers.realBuffers[idxCol].data(), sizeof(SQLREAL), buffers.indicators[idxCol].data()); + fetchStateGuard.observe(ret); if (!SQL_SUCCEEDED(ret)) { LOG("Error fetching FLOAT data for column %d", idxCol + 1); return ret; @@ -5207,6 +5318,7 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, buffers.charBuffers[idxCol].data(), MAX_DIGITS_IN_NUMERIC * sizeof(SQLCHAR), buffers.indicators[idxCol].data()); + fetchStateGuard.observe(ret); if (!SQL_SUCCEEDED(ret)) { LOG("Error fetching CHAR data for column %d", idxCol + 1); return ret; @@ -5220,6 +5332,7 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, buffers.doubleBuffers[idxCol].data(), sizeof(SQLDOUBLE), buffers.indicators[idxCol].data()); + fetchStateGuard.observe(ret); if (!SQL_SUCCEEDED(ret)) { LOG("Error fetching DOUBLE data for column %d", idxCol + 1); return ret; @@ -5234,6 +5347,7 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, buffers.timestampBuffers[idxCol].data(), sizeof(SQL_TIMESTAMP_STRUCT), buffers.indicators[idxCol].data()); + fetchStateGuard.observe(ret); if (!SQL_SUCCEEDED(ret)) { LOG("Error fetching TYPE_TIMESTAMP data for column %d", idxCol + 1); return ret; @@ -5246,6 +5360,7 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, buffers.bigIntBuffers[idxCol].data(), sizeof(SQLBIGINT), buffers.indicators[idxCol].data()); + fetchStateGuard.observe(ret); if (!SQL_SUCCEEDED(ret)) { LOG("Error fetching SBIGINT data for column %d", idxCol + 1); return ret; @@ -5258,6 +5373,7 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, buffers.dateBuffers[idxCol].data(), sizeof(SQL_DATE_STRUCT), buffers.indicators[idxCol].data()); + fetchStateGuard.observe(ret); if (!SQL_SUCCEEDED(ret)) { LOG("Error fetching TYPE_DATE data for column %d", idxCol + 1); return ret; @@ -5270,6 +5386,7 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, buffers.timeBuffers[idxCol].data(), sizeof(SQL_SS_TIME2_STRUCT), buffers.indicators[idxCol].data()); + fetchStateGuard.observe(ret); if (!SQL_SUCCEEDED(ret)) { LOG("Error fetching TYPE_TIME data for column %d", idxCol + 1); return ret; @@ -5281,6 +5398,7 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, ret = SQLGetData_ptr( hStmt, idxCol + 1, SQL_C_GUID, buffers.guidBuffers[idxCol].data(), sizeof(SQLGUID), buffers.indicators[idxCol].data()); + fetchStateGuard.observe(ret); if (!SQL_SUCCEEDED(ret)) { LOG("Error fetching GUID data for column %d", idxCol + 1); return ret; @@ -5293,6 +5411,7 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, buffers.datetimeoffsetBuffers[idxCol].data(), sizeof(DateTimeOffset), buffers.indicators[idxCol].data()); + fetchStateGuard.observe(ret); if (!SQL_SUCCEEDED(ret)) { LOG("Error fetching SS_TIMESTAMPOFFSET data for column %d", idxCol + 1); @@ -5571,6 +5690,8 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, } } + fetchStateGuard.close(); + // Transfer ownership of buffers to batch ArrowSchema // First, allocate memory for the necessary structures auto arrowSchemaBatch = std::make_unique(); @@ -5775,8 +5896,8 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, // throws a runtime error. SQLRETURN FetchAll_wrap(SqlHandlePtr StatementHandle, py::list& rows, const std::string& charEncoding = "utf-16le", - const std::string& wcharEncoding = "utf-16le", - int charCtype = SQL_C_WCHAR) { + const std::string& wcharEncoding = "utf-16le", int charCtype = SQL_C_WCHAR, + py::handle messages = {}) { PERF_TIMER("FetchAll_wrap"); // Issue #531: upgrade SQL_C_CHAR + utf-8 to SQL_C_WCHAR on Windows so the // driver does lossless UTF-16 conversion instead of returning ACP bytes. @@ -5784,11 +5905,11 @@ SQLRETURN FetchAll_wrap(SqlHandlePtr StatementHandle, py::list& rows, SQLRETURN ret; SQLHSTMT hStmt = StatementHandle->get(); // Retrieve column count - SQLSMALLINT numCols = SQLNumResultCols_wrap(StatementHandle); + SQLSMALLINT numCols = SQLNumResultCols_wrap(StatementHandle, messages); // Retrieve column metadata py::list columnNames; - ret = SQLDescribeCol_wrap(StatementHandle, columnNames); + ret = SQLDescribeCol_wrap(StatementHandle, columnNames, messages); if (!SQL_SUCCEEDED(ret)) { LOG("FetchAll_wrap: Failed to get column descriptions - SQLRETURN=%d", ret); return ret; @@ -5818,14 +5939,16 @@ SQLRETURN FetchAll_wrap(SqlHandlePtr StatementHandle, py::list& rows, py::gil_scoped_release release; ret = SQLFetch_ptr(hStmt); } + CaptureFetchDiagnostics(hStmt, ret, messages); if (ret == SQL_NO_DATA) break; if (!SQL_SUCCEEDED(ret)) return ret; py::list row; - SQLGetData_wrap(StatementHandle, numCols, row, charEncoding, wcharEncoding, - charCtype); // <-- streams LOBs correctly + ret = SQLGetData_wrap(StatementHandle, numCols, row, charEncoding, wcharEncoding, + charCtype, messages); + CheckFetchError(StatementHandle, ret); rows.append(row); } return SQL_SUCCESS; @@ -5873,33 +5996,29 @@ SQLRETURN FetchAll_wrap(SqlHandlePtr StatementHandle, py::list& rows, LOG("FetchAll_wrap: Fetching data in batch sizes of %d", fetchSize); ColumnBuffers buffers(numCols, fetchSize); + SQLULEN numRowsFetched = 0; + FetchStateGuard fetchStateGuard(StatementHandle, messages); // Bind columns - ret = SQLBindColums(hStmt, buffers, columnNames, numCols, fetchSize, charCtype); + ret = SQLBindColums(hStmt, buffers, columnNames, numCols, fetchSize, charCtype, messages); if (!SQL_SUCCEEDED(ret)) { LOG("FetchAll_wrap: Error when binding columns - SQLRETURN=%d", ret); return ret; } - SQLULEN numRowsFetched; - SQLSetStmtAttr_ptr(hStmt, SQL_ATTR_ROW_ARRAY_SIZE, (SQLPOINTER)(intptr_t)fetchSize, 0); - SQLSetStmtAttr_ptr(hStmt, SQL_ATTR_ROWS_FETCHED_PTR, &numRowsFetched, 0); + fetchStateGuard.configure(&numRowsFetched, fetchSize); while (ret != SQL_NO_DATA) { ret = FetchBatchData(hStmt, buffers, columnNames, rows, numCols, numRowsFetched, lobColumns, - charEncoding, charCtype); + charEncoding, charCtype, messages); + CheckFetchError(StatementHandle, ret); if (!SQL_SUCCEEDED(ret) && ret != SQL_NO_DATA) { LOG("FetchAll_wrap: Error when fetching data - SQLRETURN=%d", ret); return ret; } } - // Reset attributes before returning to avoid using stack pointers later - SQLSetStmtAttr_ptr(hStmt, SQL_ATTR_ROW_ARRAY_SIZE, (SQLPOINTER)1, 0); - SQLSetStmtAttr_ptr(hStmt, SQL_ATTR_ROWS_FETCHED_PTR, NULL, 0); - - // Unbind columns to allow subsequent fetchone() calls to use SQLGetData - SQLFreeStmt_ptr(hStmt, SQL_UNBIND); + fetchStateGuard.close(); return ret; } @@ -5922,8 +6041,8 @@ SQLRETURN FetchAll_wrap(SqlHandlePtr StatementHandle, py::list& rows, // fetching, it throws a runtime error. SQLRETURN FetchOne_wrap(SqlHandlePtr StatementHandle, py::list& row, const std::string& charEncoding = "utf-16le", - const std::string& wcharEncoding = "utf-16le", - int charCtype = SQL_C_WCHAR) { + const std::string& wcharEncoding = "utf-16le", int charCtype = SQL_C_WCHAR, + py::handle messages = {}) { PERF_TIMER("FetchOne_wrap"); // Issue #531: upgrade SQL_C_CHAR + utf-8 to SQL_C_WCHAR on Windows so the // driver does lossless UTF-16 conversion instead of returning ACP bytes. @@ -5934,7 +6053,10 @@ SQLRETURN FetchOne_wrap(SqlHandlePtr StatementHandle, py::list& row, // Unbind any columns from previous fetch operations (e.g., fetchmany) // to avoid conflicts with SQLGetData. SQLGetData cannot be used on // columns that are already bound. - SQLFreeStmt_ptr(hStmt, SQL_UNBIND); + ret = SQLFreeStmt_ptr(hStmt, SQL_UNBIND); + CaptureFetchDiagnostics(hStmt, ret, messages); + if (!SQL_SUCCEEDED(ret)) + return ret; // Assume hStmt is already allocated and a query has been executed { @@ -5942,11 +6064,12 @@ SQLRETURN FetchOne_wrap(SqlHandlePtr StatementHandle, py::list& row, py::gil_scoped_release release; ret = SQLFetch_ptr(hStmt); } + CaptureFetchDiagnostics(hStmt, ret, messages); if (SQL_SUCCEEDED(ret)) { // Retrieve column count - SQLSMALLINT colCount = SQLNumResultCols_wrap(StatementHandle); - ret = - SQLGetData_wrap(StatementHandle, colCount, row, charEncoding, wcharEncoding, charCtype); + SQLSMALLINT colCount = SQLNumResultCols_wrap(StatementHandle, messages); + ret = SQLGetData_wrap(StatementHandle, colCount, row, charEncoding, wcharEncoding, + charCtype, messages); if (!SQL_SUCCEEDED(ret)) { LOG("FetchOne_wrap: Error retrieving data with SQLGetData - SQLRETURN=%d", ret); return ret; @@ -6142,23 +6265,31 @@ PYBIND11_MODULE(ddbc_bindings, m) { "Get the number of rows affected by the last statement"); m.def("DDBCSQLFetch", &SQLFetch_wrap, "Fetch the next row from the result set"); m.def("DDBCSQLNumResultCols", &SQLNumResultCols_wrap, - "Get the number of columns in the result set"); + "Get the number of columns in the result set", py::arg("statementHandle"), + py::arg("messages") = py::none()); m.def("DDBCSQLDescribeCol", &SQLDescribeCol_wrap, - "Get information about a column in the result set"); - m.def("DDBCSQLGetData", &SQLGetData_wrap, "Retrieve data from the result set"); + "Get information about a column in the result set", py::arg("StatementHandle"), + py::arg("ColumnMetadata"), py::arg("messages") = py::none()); + m.def("DDBCSQLGetData", &SQLGetData_wrap, "Retrieve data from the result set", + py::arg("StatementHandle"), py::arg("colCount"), py::arg("row"), py::arg("charEncoding"), + py::arg("wcharEncoding"), py::arg("charCtype"), py::arg("messages") = py::none()); m.def("DDBCSQLMoreResults", &SQLMoreResults_wrap, "Check for more results in the result set"); m.def("DDBCSQLFetchOne", &FetchOne_wrap, "Fetch one row from the result set", py::arg("StatementHandle"), py::arg("row"), py::arg("charEncoding") = "utf-16le", - py::arg("wcharEncoding") = "utf-16le", py::arg("charCtype") = SQL_C_WCHAR); + py::arg("wcharEncoding") = "utf-16le", py::arg("charCtype") = SQL_C_WCHAR, + py::arg("messages") = py::none()); m.def("DDBCSQLFetchMany", &FetchMany_wrap, py::arg("StatementHandle"), py::arg("rows"), py::arg("fetchSize"), py::arg("charEncoding") = "utf-16le", py::arg("wcharEncoding") = "utf-16le", py::arg("charCtype") = SQL_C_WCHAR, - "Fetch many rows from the result set"); + py::arg("messages") = py::none(), "Fetch many rows from the result set"); m.def("DDBCSQLFetchAll", &FetchAll_wrap, "Fetch all rows from the result set", py::arg("StatementHandle"), py::arg("rows"), py::arg("charEncoding") = "utf-16le", - py::arg("wcharEncoding") = "utf-16le", py::arg("charCtype") = SQL_C_WCHAR); + py::arg("wcharEncoding") = "utf-16le", py::arg("charCtype") = SQL_C_WCHAR, + py::arg("messages") = py::none()); m.def("DDBCSQLFetchArrowBatch", &FetchArrowBatch_wrap, - "Fetch an arrow batch of given length from the result set"); + "Fetch an arrow batch of given length from the result set", py::arg("StatementHandle"), + py::arg("capsules"), py::arg("arrowBatchSize"), py::arg("charCtype"), + py::arg("messages") = py::none()); m.def("DDBCSQLFreeHandle", &SQLFreeHandle_wrap, "Free a handle"); m.def("DDBCSQLResetStmt", &SQLResetStmt_wrap, "Close cursor and unbind params without freeing HSTMT"); diff --git a/mssql_python/pybind/ddbc_bindings.h b/mssql_python/pybind/ddbc_bindings.h index 11c33d8d2..0d6bdc609 100644 --- a/mssql_python/pybind/ddbc_bindings.h +++ b/mssql_python/pybind/ddbc_bindings.h @@ -456,12 +456,14 @@ struct ColumnInfoExt { bool isUtf8; // Pre-computed from charEncoding (avoids string compare per cell) bool useWideChar; // True when charCtype == SQL_C_WCHAR (VARCHAR fetched as UTF-16) std::string charEncoding; // Effective decoding encoding for SQL_C_CHAR data + PyObject* messages = nullptr; // Borrowed from the enclosing native fetch call. }; // Forward declare FetchLobColumnData (defined in ddbc_bindings.cpp) - MUST be // outside namespace py::object FetchLobColumnData(SQLHSTMT hStmt, SQLUSMALLINT col, SQLSMALLINT cType, bool isWideChar, - bool isBinary, const std::string& charEncoding = "utf-8"); + bool isBinary, const std::string& charEncoding = "utf-8", + py::handle messages = {}); // Specialized column processors for each data type (eliminates switch in hot // loop) @@ -633,7 +635,8 @@ inline void ProcessChar(PyObject* row, ColumnBuffers& buffers, const void* colIn } else { // LOB / truncated: stream with SQL_C_WCHAR PyList_SET_ITEM(row, col - 1, - FetchLobColumnData(hStmt, col, SQL_C_WCHAR, true, false, "utf-16le") + FetchLobColumnData(hStmt, col, SQL_C_WCHAR, true, false, "utf-16le", + py::handle(colInfo->messages)) .release() .ptr()); } @@ -683,11 +686,11 @@ inline void ProcessChar(PyObject* row, ColumnBuffers& buffers, const void* colIn } } else { // Slow path: LOB data requires separate fetch call - PyList_SET_ITEM( - row, col - 1, - FetchLobColumnData(hStmt, col, SQL_C_CHAR, false, false, colInfo->charEncoding) - .release() - .ptr()); + PyList_SET_ITEM(row, col - 1, + FetchLobColumnData(hStmt, col, SQL_C_CHAR, false, false, + colInfo->charEncoding, py::handle(colInfo->messages)) + .release() + .ptr()); } } @@ -751,7 +754,10 @@ inline void ProcessWChar(PyObject* row, ColumnBuffers& buffers, const void* colI } else { // Slow path: LOB data requires separate fetch call PyList_SET_ITEM(row, col - 1, - FetchLobColumnData(hStmt, col, SQL_C_WCHAR, true, false).release().ptr()); + FetchLobColumnData(hStmt, col, SQL_C_WCHAR, true, false, "utf-8", + py::handle(colInfo->messages)) + .release() + .ptr()); } } @@ -790,9 +796,11 @@ inline void ProcessBinary(PyObject* row, ColumnBuffers& buffers, const void* col } } else { // Slow path: LOB data requires separate fetch call - PyList_SET_ITEM( - row, col - 1, - FetchLobColumnData(hStmt, col, SQL_C_BINARY, false, true, "").release().ptr()); + PyList_SET_ITEM(row, col - 1, + FetchLobColumnData(hStmt, col, SQL_C_BINARY, false, true, "", + py::handle(colInfo->messages)) + .release() + .ptr()); } } diff --git a/tests/test_004_cursor.py b/tests/test_004_cursor.py index 40ea4d09d..b49e49cc4 100644 --- a/tests/test_004_cursor.py +++ b/tests/test_004_cursor.py @@ -7425,29 +7425,71 @@ def test_cursor_messages_format(cursor): def test_cursor_messages_with_warnings(cursor, db_connection): - """Test that warning messages are captured correctly""" - try: - # Create a test case that might generate a warning - cursor.execute("CREATE TABLE #test_messages_warnings (id INT, value DECIMAL(5,2))") - db_connection.commit() - - # Clear messages - del cursor.messages[:] - - # Try to insert a value that might cause truncation warning - cursor.execute("INSERT INTO #test_messages_warnings VALUES (1, 123.456)") - - # Check if any warning was captured - # Note: This might be implementation-dependent - # Some drivers might not report this as a warning - if len(cursor.messages) > 0: + """Fetch warnings, including EOF diagnostics, survive batch cleanup exactly once.""" + original_ansi_warnings = cursor.execute("SELECT SESSIONPROPERTY('ANSI_WARNINGS')").fetchval() + try: + cursor.execute("SET ANSI_WARNINGS ON") + expected_messages = None + for method in ("fetchone", "fetchmany", "fetchall"): + # Stream beyond the driver's initial packet before evaluating the aggregate. + cursor.execute(""" + SELECT TOP (8192) 1 AS id, CAST(REPLICATE('x', 128) AS VARCHAR(128)) AS txt + FROM sys.all_objects AS a CROSS JOIN sys.all_objects AS b + UNION ALL + SELECT SUM(v), CAST('aggregate' AS VARCHAR(128)) + FROM (VALUES (2), (NULL)) AS warning_source(v) + UNION ALL + SELECT 3, CAST('after warning' AS VARCHAR(128)) + OPTION (MAXDOP 1) + """) + fetched = [] + warning_snapshot = list(cursor.messages) if cursor.messages else None + first_warning_phase = "execute" if warning_snapshot is not None else None + while True: + if method == "fetchone": + row = cursor.fetchone() + rows = [] if row is None else [row] + elif method == "fetchmany": + rows = cursor.fetchmany(1) + else: + rows = cursor.fetchall() + if cursor.messages: + assert len(cursor.messages) == 1, "Warning must not be duplicated" + assert cursor.messages[0][0] == "[01003] (8153)" + if warning_snapshot is None: + warning_snapshot = list(cursor.messages) + first_warning_phase = "fetch returned rows" if rows else "exhaustion" + if warning_snapshot is not None: + assert cursor.messages == warning_snapshot + if not rows: + break + fetched.extend(tuple(row) for row in rows) + + # SQL Server can defer 8153 until EOF; this verifies that terminal + # warning, not an intermediate SQLFetch/SQLGetData warning. + print(f"{method}: warning first observed at {first_warning_phase}") + assert len(fetched) == 8194 + assert fetched.count((1, "x" * 128)) == 8192 + assert fetched.count((2, "aggregate")) == 1 + assert fetched.count((3, "after warning")) == 1 + assert len(cursor.messages) == 1 + assert cursor.messages[0][0] == "[01003] (8153)" assert ( - "truncat" in cursor.messages[0][1].lower() - or "convert" in cursor.messages[0][1].lower() - ), "Warning message should mention truncation or conversion" - + "null value is eliminated by an aggregate or other set operation" + in cursor.messages[0][1].lower() + ) + if expected_messages is None: + expected_messages = list(cursor.messages) + assert cursor.messages == expected_messages, "Warning identity/order must be stable" + if method == "fetchone": + assert cursor.fetchone() is None + elif method == "fetchmany": + assert cursor.fetchmany(1) == [] + else: + assert cursor.fetchall() == [] + assert cursor.messages == expected_messages, "Repeated EOF must not replay the warning" finally: - cursor.execute("DROP TABLE IF EXISTS #test_messages_warnings") + cursor.execute("SET ANSI_WARNINGS " + ("ON" if original_ansi_warnings else "OFF")) db_connection.commit() diff --git a/tests/test_004_cursor_arrow.py b/tests/test_004_cursor_arrow.py index 6a7d7b5d9..37a336789 100644 --- a/tests/test_004_cursor_arrow.py +++ b/tests/test_004_cursor_arrow.py @@ -612,6 +612,13 @@ def fake_drain(hstmt): "post-close drain was skipped on the SQL_CLOSE success path " "(SQL_SUCCESS_WITH_INFO warnings would be lost)" ) + assert call_count["n"] == 2 + + reader = cursor.execute("SELECT 1 AS a UNION ALL SELECT 2").arrow_reader(batch_size=1) + call_count["n"] = 0 + assert sum(batch.num_rows for batch in reader) == 2 + assert call_count["n"] == 1, "Natural exhaustion must only drain SQL_CLOSE diagnostics" + assert cursor.messages == [("01000", "synthetic warning #1")] def test_arrow_reader_close_retries_after_failed_attempt(cursor: mssql_python.Cursor): diff --git a/tests/test_fetch_settings_cache.py b/tests/test_fetch_settings_cache.py index 58a1423f3..d200d64b4 100644 --- a/tests/test_fetch_settings_cache.py +++ b/tests/test_fetch_settings_cache.py @@ -175,7 +175,9 @@ def test_wchar_decoding_forwarded_to_live_fetch_bridge(connection, method, bridg expected_reads += 2 assert fetch_rows(cursor, method)[0].txt == "\u00e9" assert fetch.call_count == index - assert fetch.call_args.args[-3:] == ("utf-16le", encoding, mssql_python.SQL_WCHAR) + assert fetch.call_args.args[-4:-1] == ("utf-16le", encoding, mssql_python.SQL_WCHAR) + assert fetch.call_args.args[-1] is cursor.messages + assert fetch.call_args.kwargs == {} assert reads.call_count == expected_reads previous_encoding = encoding assert [call.args[0] for call in reads.call_args_list] == [ @@ -726,7 +728,9 @@ def test_char_decoding_ctype_refresh(connection, method, bridge_name): cursor.execute("SELECT CONVERT(VARCHAR(1), 0xE9) AS txt") connection.setdecoding(mssql_python.SQL_CHAR, encoding=encoding, ctype=ctype) assert fetch_rows(cursor, method)[0].txt == "\u00e9" - assert fetch.call_args.args[-3:] == (encoding, "utf-16le", ctype) + assert fetch.call_args.args[-4:-1] == (encoding, "utf-16le", ctype) + assert fetch.call_args.args[-1] is cursor.messages + assert fetch.call_args.kwargs == {} assert reads.call_count == 8 @@ -839,12 +843,14 @@ def test_fetch_drains_diagnostics_independent_of_final_status( connection, method, bridge_name, status ): bridge = getattr(mssql_python.ddbc_bindings, bridge_name) - warning = ("01000", 0, "injected fetch warning") + warning = ("[01000] (0)", "injected fetch warning") with connection.cursor() as cursor: cursor.execute("SELECT CAST(N'abc' AS NVARCHAR(MAX)) AS txt") def fetch_with_final_status(*args): bridge(*args) + assert args[-1] is cursor.messages + args[-1].append(warning) return status with ( @@ -856,8 +862,8 @@ def fetch_with_final_status(*args): ) as diagnostics, ): fetch_rows(cursor, method) - diagnostics.assert_called_once_with(cursor.hstmt) - assert warning in cursor.messages + diagnostics.assert_not_called() + assert cursor.messages == [warning] @pytest.mark.parametrize( From f1bdd4a9141d31000289b7485d49ae97b5927ab3 Mon Sep 17 00:00:00 2001 From: Jahnvi Thakkar Date: Thu, 24 Sep 2026 10:59:46 +0530 Subject: [PATCH 2/6] FIX: Propagate fetch metadata and variant probe errors Raise through the existing DB-API translator before returning a failed SQLNumResultCols output. Return a failed sql_variant probe immediately, before a later column can replace its diagnostics. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- mssql_python/pybind/ddbc_bindings.cpp | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index 94fb28139..2123d3680 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -3024,9 +3024,9 @@ SQLSMALLINT SQLNumResultCols_wrap(SqlHandlePtr statementHandle, py::handle messa } SQLSMALLINT columnCount; - // TODO: Handle the return code SQLRETURN ret = SQLNumResultCols_ptr(statementHandle->get(), &columnCount); CaptureFetchDiagnostics(statementHandle->get(), ret, messages); + CheckFetchError(statementHandle, ret); return columnCount; } @@ -3353,8 +3353,7 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p if (!SQL_SUCCEEDED(ret)) { LOG_ERROR("SQLGetData: Failed to probe sql_variant column %d - SQLRETURN=%d", i, ret); - row.append(py::none()); - continue; + return ret; } if (indicator == SQL_NULL_DATA) { row.append(py::none()); From bf3322f4b30868e6c25c28f0ab6129c40ba3e67d Mon Sep 17 00:00:00 2001 From: Jahnvi Thakkar Date: Thu, 24 Sep 2026 11:06:23 +0530 Subject: [PATCH 3/6] FIX: Use signed lengths in fetch diagnostic comparisons Compare ODBC SQLLEN indicators against bounded buffer sizes converted to SQLLEN. This fixes the four MSVC C4018 errors promoted by /WX without weakening compiler settings, and preserves SQL_NO_TOTAL handling. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- mssql_python/pybind/ddbc_bindings.cpp | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index 2123d3680..d5ad5fc32 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -3426,7 +3426,8 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p CaptureFetchDiagnostics( hStmt, ret, messages, ret == SQL_SUCCESS_WITH_INFO && - (dataLen == SQL_NO_TOTAL || dataLen >= fetchBufferSize)); + (dataLen == SQL_NO_TOTAL || + dataLen >= static_cast(fetchBufferSize))); if (SQL_SUCCEEDED(ret)) { if (dataLen > 0) { uint64_t numCharsInData = dataLen / sizeof(SQLWCHAR); @@ -3508,7 +3509,8 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p CaptureFetchDiagnostics( hStmt, ret, messages, ret == SQL_SUCCESS_WITH_INFO && - (dataLen == SQL_NO_TOTAL || dataLen >= dataBuffer.size())); + (dataLen == SQL_NO_TOTAL || + dataLen >= static_cast(dataBuffer.size()))); if (SQL_SUCCEEDED(ret)) { // columnSize is in chars, dataLen is in bytes if (dataLen > 0) { @@ -3603,7 +3605,8 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p CaptureFetchDiagnostics( hStmt, ret, messages, ret == SQL_SUCCESS_WITH_INFO && - (dataLen == SQL_NO_TOTAL || dataLen >= fetchBufferSize)); + (dataLen == SQL_NO_TOTAL || + dataLen >= static_cast(fetchBufferSize))); if (SQL_SUCCEEDED(ret)) { if (dataLen > 0) { uint64_t numCharsInData = dataLen / sizeof(SQLWCHAR); @@ -3930,7 +3933,8 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p &dataLen); CaptureFetchDiagnostics(hStmt, ret, messages, ret == SQL_SUCCESS_WITH_INFO && - (dataLen == SQL_NO_TOTAL || dataLen > columnSize)); + (dataLen == SQL_NO_TOTAL || + dataLen > static_cast(columnSize))); if (SQL_SUCCEEDED(ret)) { if (dataLen > 0) { From 88aee07efe29d311209d1dc135a347dc39924e74 Mon Sep 17 00:00:00 2001 From: Jahnvi Thakkar Date: Thu, 24 Sep 2026 13:50:44 +0530 Subject: [PATCH 4/6] FIX: Skip message retrieval for internal LOB truncation Read SQLSTATE with SQLGetDiagFieldW before materializing internal continuation diagnostics. Skip only confirmed 01004 records, retrieve every unrelated record unchanged, and retain full-record fallback if the field API is unavailable or fails. Track actual field calls with the existing profiler. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- mssql_python/pybind/ddbc_bindings.cpp | 22 ++++++++++++++++++++++ mssql_python/pybind/ddbc_bindings.h | 3 +++ 2 files changed, 25 insertions(+) diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index d5ad5fc32..c30e13082 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -260,6 +260,7 @@ SQLCancelFunc SQLCancel_ptr = nullptr; // Diagnostic APIs SQLGetDiagRecFunc SQLGetDiagRec_ptr = nullptr; +SQLGetDiagFieldFunc SQLGetDiagField_ptr = nullptr; // DAE APIs SQLParamDataFunc SQLParamData_ptr = nullptr; @@ -1473,6 +1474,9 @@ DriverHandle LoadDriverOrThrowException() { SQLCancel_ptr = GetFunctionPointer(handle, "SQLCancel"); SQLGetDiagRec_ptr = GetFunctionPointer(handle, "SQLGetDiagRecW"); + SQLGetDiagField_ptr = GetFunctionPointer(handle, "SQLGetDiagFieldW"); + if (!SQLGetDiagField_ptr) + LOG("SQLGetDiagFieldW unavailable; using full diagnostic records"); SQLParamData_ptr = GetFunctionPointer(handle, "SQLParamData"); SQLPutData_ptr = GetFunctionPointer(handle, "SQLPutData"); @@ -1901,6 +1905,24 @@ static void AppendDiagRecords(SQLHANDLE rawHandle, SQLSMALLINT handleType, py::h // Iterate through all available diagnostic records for (SQLSMALLINT recNumber = 1;; recNumber++) { SQLWCHAR sqlState[6] = {0}; + if (internalTruncation && SQLGetDiagField_ptr) { + SQLRETURN stateReturn; + { + PERF_TIMER("AppendDiagRecords::SQLGetDiagField_call"); + stateReturn = SQLGetDiagField_ptr(handleType, rawHandle, recNumber, + SQL_DIAG_SQLSTATE, sqlState, + static_cast(sizeof(sqlState)), + nullptr); + } + if (stateReturn == SQL_NO_DATA) + break; + // Skip only this continuation record, without retrieving its message text. + if (stateReturn == SQL_SUCCESS && std::equal(sqlState, sqlState + 6, u"01004")) + continue; + if (stateReturn != SQL_SUCCESS) + LOG("AppendDiagRecords: SQLSTATE lookup returned %d; reading full record %d", + stateReturn, recNumber); + } SQLWCHAR message[SQL_MAX_MESSAGE_LENGTH_SQLSERVER] = {0}; SQLINTEGER nativeError = 0; SQLSMALLINT messageLen = 0; diff --git a/mssql_python/pybind/ddbc_bindings.h b/mssql_python/pybind/ddbc_bindings.h index 0d6bdc609..d8b5d1a4e 100644 --- a/mssql_python/pybind/ddbc_bindings.h +++ b/mssql_python/pybind/ddbc_bindings.h @@ -144,6 +144,8 @@ typedef SQLRETURN(SQL_API* SQLCancelFunc)(SQLHSTMT); // Diagnostic APIs typedef SQLRETURN(SQL_API* SQLGetDiagRecFunc)(SQLSMALLINT, SQLHANDLE, SQLSMALLINT, SQLWCHAR*, SQLINTEGER*, SQLWCHAR*, SQLSMALLINT, SQLSMALLINT*); +typedef SQLRETURN(SQL_API* SQLGetDiagFieldFunc)(SQLSMALLINT, SQLHANDLE, SQLSMALLINT, SQLSMALLINT, + SQLPOINTER, SQLSMALLINT, SQLSMALLINT*); typedef SQLRETURN(SQL_API* SQLDescribeParamFunc)(SQLHSTMT, SQLUSMALLINT, SQLSMALLINT*, SQLULEN*, SQLSMALLINT*, SQLSMALLINT*); @@ -202,6 +204,7 @@ extern SQLCancelFunc SQLCancel_ptr; // Diagnostic APIs extern SQLGetDiagRecFunc SQLGetDiagRec_ptr; +extern SQLGetDiagFieldFunc SQLGetDiagField_ptr; extern SQLDescribeParamFunc SQLDescribeParam_ptr; From 6adce4e049aa13c22c4181d90efb417588973cc9 Mon Sep 17 00:00:00 2001 From: Jahnvi Thakkar Date: Thu, 24 Sep 2026 14:49:38 +0530 Subject: [PATCH 5/6] TEST: Cover native mixed fetch diagnostics and fallbacks Exercise the compiled diagnostic collector in isolated subprocesses with deterministic ODBC diagnostic callbacks and a real LOB continuation. Verify warning order, internal truncation filtering, SQLSTATE fallback, full-record call counts, and the unfiltered diagnostic path. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- tests/test_fetch_settings_cache.py | 169 ++++++++++++++++++++++++++++- 1 file changed, 168 insertions(+), 1 deletion(-) diff --git a/tests/test_fetch_settings_cache.py b/tests/test_fetch_settings_cache.py index d200d64b4..1adbdd346 100644 --- a/tests/test_fetch_settings_cache.py +++ b/tests/test_fetch_settings_cache.py @@ -2,11 +2,12 @@ Copyright (c) Microsoft Corporation. Licensed under the MIT license. -Regression and operation-count tests for connection settings cached by fetch APIs. +Regression and operation-count tests for fetch settings and diagnostic preservation. All integration queries are read-only and each test owns its connection. """ import datetime +from pathlib import Path import subprocess import sys import uuid @@ -866,6 +867,172 @@ def fetch_with_final_status(*args): assert cursor.messages == [warning] +@pytest.mark.skipif( + sys.platform == "win32", + reason="Windows does not export the native diagnostic function-pointer globals", +) +@pytest.mark.parametrize("mode", ("available", "missing", "error", "info", "unterminated")) +def test_native_mixed_fetch_diagnostics(conn_str, mode): + if not conn_str: + pytest.skip("DB_CONNECTION_STRING is required") + # Driver pointers are process-global: never replace them in the pytest process. + code = ( + "import runpy, sys; " + "runpy.run_path(sys.argv[1])['_check_native_mixed_fetch_diagnostics']" + "(sys.argv[2], sys.argv[3])" + ) + result = subprocess.run( + [ + sys.executable, + "-c", + code, + str(Path(__file__).resolve()), + mode, + str(Path(mssql_python.ddbc_bindings.module.__file__).resolve()), + ], + capture_output=True, + text=True, + timeout=60, + ) + assert result.returncode == 0, (result.returncode, result.stdout, result.stderr) + + +def _check_native_mixed_fetch_diagnostics(mode, expected_native): + import ctypes + import os + + native = Path(mssql_python.ddbc_bindings.module.__file__).resolve() + assert native == Path(expected_native) + library = ctypes.CDLL(str(native)) + rec_pointer = ctypes.c_void_p.in_dll(library, "SQLGetDiagRec_ptr") + field_pointer = ctypes.c_void_p.in_dll(library, "SQLGetDiagField_ptr") + smallint = ctypes.c_short + wchar_pointer = ctypes.POINTER(ctypes.c_uint16) + smallint_pointer = ctypes.POINTER(smallint) + rec_type = ctypes.CFUNCTYPE( + smallint, + smallint, + ctypes.c_void_p, + smallint, + wchar_pointer, + ctypes.POINTER(ctypes.c_int32), + wchar_pointer, + smallint, + smallint_pointer, + ) + field_type = ctypes.CFUNCTYPE( + smallint, + smallint, + ctypes.c_void_p, + smallint, + smallint, + ctypes.c_void_p, + smallint, + smallint_pointer, + ) + success = ConstantsDDBC.SQL_SUCCESS.value + info = ConstantsDDBC.SQL_SUCCESS_WITH_INFO.value + no_data = ConstantsDDBC.SQL_NO_DATA.value + error = ConstantsDDBC.SQL_ERROR.value + records = [ + ("01000", 11, "before truncation \u00e9"), + ("01004", 0, "internal first chunk"), + ("01S07", 22, "between truncations"), + ("01004", 0, "internal second chunk"), + ("01000", 33, "after truncation"), + ] + all_records = [(f"[{state}] ({number})", message) for state, number, message in records] + wanted = [all_records[index] for index in (0, 2, 4)] + rec_calls, field_calls, callback_errors = [], [], [] + + @rec_type + def read_record(handle_type, handle, number, state, native_error, message, capacity, length): + rec_calls.append(number) + if not handle or handle_type != ConstantsDDBC.SQL_HANDLE_STMT.value or number < 1: + callback_errors.append("invalid record lookup") + return error + if number > len(records): + return no_data + sqlstate, code, text = records[number - 1] + encoded = text.encode("utf-16le") + if ( + not state + or not native_error + or not message + or not length + or capacity <= len(encoded) // 2 + ): + callback_errors.append("invalid record output buffer") + return error + ctypes.memmove(state, (sqlstate + "\0").encode("utf-16le"), 12) + ctypes.memmove(message, encoded + b"\0\0", len(encoded) + 2) + native_error[0] = code + length[0] = len(encoded) // 2 + return success + + @field_type + def read_state(handle_type, handle, number, identifier, output, capacity, length): + field_calls.append(number) + # SQL_DIAG_SQLSTATE uses bytes, including the sixth SQLWCHAR terminator. + if ( + not handle + or handle_type != ConstantsDDBC.SQL_HANDLE_STMT.value + or number < 1 + or identifier != 4 + or not output + or capacity != 12 + ): + callback_errors.append("invalid SQLSTATE lookup or byte capacity") + return error + if mode == "error": + return error + if number > len(records): + return no_data + sqlstate = records[number - 1][0] + "\0" + if mode == "info": + sqlstate = "01004\0" + elif mode == "unterminated": + sqlstate = "01004X" + ctypes.memmove(output, sqlstate.encode("utf-16le"), 12) + return info if mode == "info" else success + + try: + connection = mssql_python.connect(os.environ["DB_CONNECTION_STRING"], timeout=5) + except mssql_python.Error as failure: + raise AssertionError( + f"Connection failed: {type(failure).__name__}; connection details withheld" + ) from None + with connection, connection.cursor() as cursor: + cursor.execute("SELECT CAST(REPLICATE(CAST('x' AS VARCHAR(MAX)), 8193) AS VARBINARY(MAX))") + original_rec, original_field = rec_pointer.value, field_pointer.value + assert original_rec, "Driver diagnostic records must be available" + try: + rec_pointer.value = ctypes.cast(read_record, ctypes.c_void_p).value + field_pointer.value = ( + None if mode == "missing" else ctypes.cast(read_state, ctypes.c_void_p).value + ) + # One real SQLGetData continuation runs the compiled native mixed-record filter. + assert tuple(cursor.fetchone()) == (b"x" * 8193,) + assert not callback_errors, callback_errors + assert cursor.messages == wanted + assert field_calls == ([] if mode == "missing" else [1, 2, 3, 4, 5, 6]) + if mode == "available": + assert rec_calls == [1, 3, 5] + else: + assert rec_calls == list(range(1, 7 if mode in ("missing", "error") else 6)) + + rec_calls.clear() + field_calls.clear() + assert mssql_python.ddbc_bindings.DDBCSQLGetAllDiagRecords(cursor.hstmt) == all_records + assert not callback_errors, callback_errors + assert rec_calls == [1, 2, 3, 4, 5, 6] + assert field_calls == [] + finally: + rec_pointer.value, field_pointer.value = original_rec, original_field + cursor.execute("SELECT 42") + assert tuple(cursor.fetchone()) == (42,) + + @pytest.mark.parametrize( ("method", "bridge_name"), ( From cebec3a29ed8158e00467c622c0d8bc91947dfdb Mon Sep 17 00:00:00 2001 From: Jahnvi Thakkar Date: Thu, 24 Sep 2026 14:56:25 +0530 Subject: [PATCH 6/6] TEST: Isolate native diagnostic injection at the callback boundary Contain callback exceptions and assert them outside the C ABI. Learn the target statement through an unchanged diagnostic call, forward unrelated handles and EOF to saved driver addresses, and restore pointers before cursor cleanup. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- tests/test_fetch_settings_cache.py | 54 ++++++++++++++++++++++++++++-- 1 file changed, 51 insertions(+), 3 deletions(-) diff --git a/tests/test_fetch_settings_cache.py b/tests/test_fetch_settings_cache.py index 1adbdd346..78ad5229b 100644 --- a/tests/test_fetch_settings_cache.py +++ b/tests/test_fetch_settings_cache.py @@ -944,9 +944,32 @@ def _check_native_mixed_fetch_diagnostics(mode, expected_native): all_records = [(f"[{state}] ({number})", message) for state, number, message in records] wanted = [all_records[index] for index in (0, 2, 4)] rec_calls, field_calls, callback_errors = [], [], [] + observed_handles, delegated_records = [], [] + phase, target_handle = "observe", None - @rec_type + def guarded(callback_type): + def decorate(function): + def boundary(*args): + try: + return function(*args) + except BaseException as failure: + # ctypes otherwise prints and suppresses exceptions crossing the C boundary. + callback_errors.append(type(failure).__name__) + return error + + return callback_type(boundary) + + return decorate + + @guarded(rec_type) def read_record(handle_type, handle, number, state, native_error, message, capacity, length): + if phase == "observe": + observed_handles.append((handle_type, handle)) + if phase != "inject" or (handle_type, handle) != target_handle: + delegated_records.append((handle_type, handle, number)) + return original_read_record( + handle_type, handle, number, state, native_error, message, capacity, length + ) rec_calls.append(number) if not handle or handle_type != ConstantsDDBC.SQL_HANDLE_STMT.value or number < 1: callback_errors.append("invalid record lookup") @@ -970,8 +993,14 @@ def read_record(handle_type, handle, number, state, native_error, message, capac length[0] = len(encoded) // 2 return success - @field_type + @guarded(field_type) def read_state(handle_type, handle, number, identifier, output, capacity, length): + if phase != "inject" or (handle_type, handle) != target_handle: + if original_read_state is None: + return error + return original_read_state( + handle_type, handle, number, identifier, output, capacity, length + ) field_calls.append(number) # SQL_DIAG_SQLSTATE uses bytes, including the sixth SQLWCHAR terminator. if ( @@ -1002,15 +1031,28 @@ def read_state(handle_type, handle, number, identifier, output, capacity, length raise AssertionError( f"Connection failed: {type(failure).__name__}; connection details withheld" ) from None - with connection, connection.cursor() as cursor: + with connection, connection.cursor() as cursor, connection.cursor() as other_cursor: cursor.execute("SELECT CAST(REPLICATE(CAST('x' AS VARCHAR(MAX)), 8193) AS VARBINARY(MAX))") original_rec, original_field = rec_pointer.value, field_pointer.value assert original_rec, "Driver diagnostic records must be available" + original_read_record = rec_type(original_rec) + original_read_state = field_type(original_field) if original_field else None try: rec_pointer.value = ctypes.cast(read_record, ctypes.c_void_p).value + # Learn this cursor's raw handle while forwarding the diagnostic call unchanged. + assert mssql_python.ddbc_bindings.DDBCSQLGetAllDiagRecords(cursor.hstmt) == [] + assert not callback_errors, callback_errors + assert len(observed_handles) == 1 + target_handle = observed_handles[0] + assert target_handle[0] == ConstantsDDBC.SQL_HANDLE_STMT.value and target_handle[1] + phase = "inject" field_pointer.value = ( None if mode == "missing" else ctypes.cast(read_state, ctypes.c_void_p).value ) + delegated_records.clear() + assert mssql_python.ddbc_bindings.DDBCSQLGetAllDiagRecords(other_cursor.hstmt) == [] + assert len(delegated_records) == 1 + assert delegated_records[0][1] != target_handle[1] # One real SQLGetData continuation runs the compiled native mixed-record filter. assert tuple(cursor.fetchone()) == (b"x" * 8193,) assert not callback_errors, callback_errors @@ -1027,6 +1069,12 @@ def read_state(handle_type, handle, number, identifier, output, capacity, length assert not callback_errors, callback_errors assert rec_calls == [1, 2, 3, 4, 5, 6] assert field_calls == [] + delegated_records.clear() + phase = "delegate" + assert cursor.fetchone() is None + assert not callback_errors, callback_errors + assert delegated_records == [(*target_handle, 1)] + assert cursor.messages == wanted finally: rec_pointer.value, field_pointer.value = original_rec, original_field cursor.execute("SELECT 42")