From 15e4092e90620b604b14105a0162b3fae1f2f3ae Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Mon, 21 Sep 2026 11:56:17 +0530 Subject: [PATCH 01/39] FIX: Validate caller-controlled native buffer sizes --- mssql_python/cursor.py | 59 +++++++++--- mssql_python/pybind/ddbc_bindings.cpp | 126 +++++++++++++++++++++----- tests/test_004_cursor.py | 28 ++++++ tests/test_024_bulkcopy_arrow.py | 5 + 4 files changed, 181 insertions(+), 37 deletions(-) diff --git a/mssql_python/cursor.py b/mssql_python/cursor.py index db45462c8..b364f9091 100644 --- a/mssql_python/cursor.py +++ b/mssql_python/cursor.py @@ -51,6 +51,7 @@ SMALLMONEY_MAX: decimal.Decimal = decimal.Decimal("214748.3647") MONEY_MIN: decimal.Decimal = decimal.Decimal("-922337203685477.5808") MONEY_MAX: decimal.Decimal = decimal.Decimal("922337203685477.5807") +MAX_NATIVE_ROW_COUNT: int = 2_147_483_647 # SQL BIGINT is a signed 64-bit integer. Ints outside this range have no BIGINT # encoding and must be rejected at detect time on both paths (see _map_sql_type). BIGINT_MIN: int = -(2**63) @@ -385,9 +386,7 @@ def __init__(self, connection: "Connection", timeout: int = 0) -> None: ] ] = None self.rowcount: int = -1 - self.arraysize: int = ( - 1 # Default number of rows to fetch at a time is 1, user can change it - ) + self.arraysize = 1 self.buffer_length: int = 1024 # Default buffer length for string data self._result_set_empty: bool = False # Add this initialization self.last_executed_stmt: str = "" # Stores the last statement executed by this cursor @@ -1136,6 +1135,25 @@ def _ensure_pyarrow(self) -> Any: "pyarrow is required for Arrow fetch methods. Please install pyarrow." ) from e + @staticmethod + def _validate_native_row_count(value: int, name: str, allow_zero: bool) -> int: + minimum = 0 if allow_zero else 1 + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{name} must be an integer") + if value < minimum or value > MAX_NATIVE_ROW_COUNT: + raise ValueError( + f"{name} must be between {minimum} and {MAX_NATIVE_ROW_COUNT}, got {value}" + ) + return value + + @property + def arraysize(self) -> int: + return self._arraysize + + @arraysize.setter + def arraysize(self, value: int) -> None: + self._arraysize = self._validate_native_row_count(value, "arraysize", allow_zero=False) + def setinputsizes(self, sizes: List[Union[int, tuple]]) -> None: """ Sets the type information to be used for parameters in execute and executemany. @@ -1188,7 +1206,11 @@ def setinputsizes(self, sizes: List[Union[int, tuple]]) -> None: sql_type, column_size, decimal_digits = size_info # Validate SQL type - if not isinstance(sql_type, int) or sql_type not in valid_sql_types: + if ( + isinstance(sql_type, bool) + or not isinstance(sql_type, int) + or sql_type not in valid_sql_types + ): raise ValueError( f"Invalid SQL type: {sql_type}. Must be a valid SQL type constant." ) @@ -1196,12 +1218,20 @@ def setinputsizes(self, sizes: List[Union[int, tuple]]) -> None: sql_type = ODBC3_TEMPORAL_SQL_TYPES.get(sql_type, sql_type) # Validate size and precision - if not isinstance(column_size, int) or column_size < 0: + if ( + isinstance(column_size, bool) + or not isinstance(column_size, int) + or column_size < 0 + ): raise ValueError( f"Invalid column size: {column_size}. Must be a non-negative integer." ) - if not isinstance(decimal_digits, int) or decimal_digits < 0: + if ( + isinstance(decimal_digits, bool) + or not isinstance(decimal_digits, int) + or decimal_digits < 0 + ): raise ValueError( f"Invalid decimal digits: {decimal_digits}. " f"Must be a non-negative integer." @@ -1220,7 +1250,11 @@ def setinputsizes(self, sizes: List[Union[int, tuple]]) -> None: sql_type = size_info # Validate SQL type - if not isinstance(sql_type, int) or sql_type not in valid_sql_types: + if ( + isinstance(sql_type, bool) + or not isinstance(sql_type, int) + or sql_type not in valid_sql_types + ): raise ValueError( f"Invalid SQL type: {sql_type}. Must be a valid SQL type constant." ) @@ -2852,8 +2886,10 @@ def fetchmany(self, size: Optional[int] = None) -> List[Row]: if size is None: size = self.arraysize + else: + size = self._validate_native_row_count(size, "size", allow_zero=True) - if size <= 0: + if size == 0: return [] char_decoding = self._get_decoding_settings(ddbc_sql_const.SQL_CHAR.value) @@ -2989,6 +3025,7 @@ def arrow_batch(self, batch_size: int = 8192) -> "pyarrow.RecordBatch": """ self._check_closed() # Check if the cursor is closed pyarrow = self._ensure_pyarrow() + batch_size = self._validate_native_row_count(batch_size, "batch_size", allow_zero=True) if not self._has_result_set and self.description: self._reset_rownumber() @@ -2996,9 +3033,7 @@ def arrow_batch(self, batch_size: int = 8192) -> "pyarrow.RecordBatch": capsules = [] 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 - ) + ret = ddbc_bindings.DDBCSQLFetchArrowBatch(self.hstmt, capsules, batch_size, char_c_type) check_error(ddbc_sql_const.SQL_HANDLE_STMT.value, self.hstmt, ret) batch = pyarrow.RecordBatch._import_from_c_capsule(*capsules) @@ -3383,7 +3418,7 @@ def _bulkcopy_core_and_validate(table_name, batch_size, timeout): logger.error("bulkcopy: Invalid table_name parameter") raise ValueError("table_name must be a non-empty string") - if not isinstance(batch_size, int): + if not isinstance(batch_size, int) or isinstance(batch_size, bool): raise TypeError( f"batch_size must be a non-negative integer, got {type(batch_size).__name__}" ) diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index 1f235f829..8f16d46b4 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -21,6 +21,7 @@ #include // For std::memcpy #include #include +#include #include // std::forward #include // CPython datetime API (PyDateTime_IMPORT, PyDateTime_GET_*, etc.) @@ -322,12 +323,37 @@ ParamType* AllocateParamBuffer(std::vector>& paramBuffers, template ParamType* AllocateParamBufferArray(std::vector>& paramBuffers, size_t count) { + if (count > std::numeric_limits::max() / sizeof(ParamType)) { + ThrowStdException("Parameter buffer size is too large"); + } std::shared_ptr buffer(new ParamType[count], std::default_delete()); ParamType* raw = buffer.get(); paramBuffers.push_back(buffer); return raw; } +size_t CheckedAddSize(size_t left, size_t right, const char* errorMessage) { + if (left > std::numeric_limits::max() - right) { + ThrowStdException(errorMessage); + } + return left + right; +} + +size_t CheckedMultiplySize(size_t left, size_t right, const char* errorMessage) { + if (left != 0 && right > std::numeric_limits::max() / left) { + ThrowStdException(errorMessage); + } + return left * right; +} + +template +std::unique_ptr AllocateUniqueArray(size_t count, const char* errorMessage) { + if (count > std::numeric_limits::max() / sizeof(ElementType)) { + ThrowStdException(errorMessage); + } + return std::make_unique(count); +} + std::string DescribeChar(unsigned char ch) { if (ch >= 32 && ch <= 126) { return std::string("'") + static_cast(ch) + "'"; @@ -460,6 +486,10 @@ SQLRETURN BindParameters(SqlHandle& handle, SQLHANDLE hStmt, const py::list& par "with %zu parameters", (void*)hStmt, params.size()); + if (params.size() != paramInfos.size()) { + ThrowStdException("Parameter count does not match parameter metadata count"); + } + // GH-627: resolve unknown NULL param SQL types before binding any param. PreResolveUnknownNullTypes(handle, hStmt, paramInfos, ¶ms); for (int paramIndex = 0; paramIndex < params.size(); paramIndex++) { @@ -537,8 +567,7 @@ SQLRETURN BindParameters(SqlHandle& handle, SQLHANDLE hStmt, const py::list& par break; } case SQL_C_BINARY: { - if (!py::isinstance(param) && !py::isinstance(param) && - !py::isinstance(param)) { + if (!py::isinstance(param) && !py::isinstance(param)) { ThrowStdException(MakeParamMismatchErrorStr(paramInfo.paramCType, paramIndex)); } if (paramInfo.isDAE) { @@ -2194,6 +2223,9 @@ SQLRETURN BindParameterArray(SqlHandle& handle, SQLHANDLE hStmt, const py::list& LOG("BindParameterArray: Starting column-wise array binding - " "param_count=%zu, param_set_size=%zu", columnwise_params.size(), paramSetSize); + if (columnwise_params.size() != paramInfos.size()) { + ThrowStdException("Parameter count does not match parameter metadata count"); + } std::vector> tempBuffers; @@ -2261,8 +2293,18 @@ SQLRETURN BindParameterArray(SqlHandle& handle, SQLHANDLE hStmt, const py::list& LOG("BindParameterArray: Binding SQL_C_WCHAR array - " "param_index=%d, count=%zu, column_size=%zu", paramIndex, paramSetSize, info.columnSize); + const size_t elementWidth = CheckedAddSize( + info.columnSize, 1, "Wide-character parameter size is too large"); + const size_t bufferBytes = CheckedMultiplySize( + elementWidth, sizeof(SQLWCHAR), + "Wide-character parameter length is too large"); + if (bufferBytes > static_cast(std::numeric_limits::max())) { + ThrowStdException("Wide-character parameter length is too large"); + } SQLWCHAR* wcharArray = AllocateParamBufferArray( - tempBuffers, paramSetSize * (info.columnSize + 1)); + tempBuffers, + CheckedMultiplySize(paramSetSize, elementWidth, + "Wide-character parameter buffer is too large")); strLenOrIndArray = AllocateParamBufferArray(tempBuffers, paramSetSize); for (size_t i = 0; i < paramSetSize; ++i) { if (columnValues[i].is_none()) { @@ -2287,7 +2329,7 @@ SQLRETURN BindParameterArray(SqlHandle& handle, SQLHANDLE hStmt, const py::list& "param_index=%d", paramIndex); dataPtr = wcharArray; - bufferLength = (info.columnSize + 1) * sizeof(SQLWCHAR); + bufferLength = static_cast(bufferBytes); break; } case SQL_C_TINYINT: @@ -2358,8 +2400,15 @@ SQLRETURN BindParameterArray(SqlHandle& handle, SQLHANDLE hStmt, const py::list& LOG("BindParameterArray: Binding SQL_C_CHAR/BINARY array - " "param_index=%d, count=%zu, column_size=%zu, encoding='%s'", paramIndex, paramSetSize, info.columnSize, charEncoding.c_str()); + const size_t elementWidth = CheckedAddSize( + info.columnSize, 1, "Character parameter size is too large"); + if (elementWidth > static_cast(std::numeric_limits::max())) { + ThrowStdException("Character parameter length is too large"); + } char* charArray = AllocateParamBufferArray( - tempBuffers, paramSetSize * (info.columnSize + 1)); + tempBuffers, + CheckedMultiplySize(paramSetSize, elementWidth, + "Character parameter buffer is too large")); strLenOrIndArray = AllocateParamBufferArray(tempBuffers, paramSetSize); for (size_t i = 0; i < paramSetSize; ++i) { if (columnValues[i].is_none()) { @@ -2412,7 +2461,7 @@ SQLRETURN BindParameterArray(SqlHandle& handle, SQLHANDLE hStmt, const py::list& "param_index=%d", paramIndex); dataPtr = charArray; - bufferLength = info.columnSize + 1; + bufferLength = static_cast(elementWidth); break; } case SQL_C_BIT: { @@ -3177,6 +3226,9 @@ py::object FetchLobColumnData(SQLHSTMT hStmt, SQLUSMALLINT colIndex, SQLSMALLINT return py::str(""); } if (isWideChar) { + if (buffer.size() % sizeof(SQLWCHAR) != 0) { + ThrowStdException("Wide-character LOB data has an invalid byte length"); + } size_t wcharCount = buffer.size() / sizeof(SQLWCHAR); std::vector alignedBuf(wcharCount); std::memcpy(alignedBuf.data(), buffer.data(), buffer.size()); @@ -4836,6 +4888,15 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, int arrowBatchSize, int charCtype) { PERF_TIMER("FetchArrowBatch_wrap"); + if (arrowBatchSize < 0) { + ThrowStdException("Arrow batch size must be non-negative"); + } + const size_t batchSize = static_cast(arrowBatchSize); + const size_t offsetCount = CheckedAddSize(batchSize, 1, "Arrow batch size is too large"); + const size_t initialVarDataSize = + CheckedMultiplySize(batchSize, 42, "Arrow batch size is too large"); + const size_t bitmapSize = + CheckedAddSize(batchSize, 7, "Arrow batch size is too large") / 8; // Fetch narrow char data as SQL_C_CHAR if on Linux/macOS and configured by the user charCtype = EffectiveCharCtypeForFetch(charCtype, "utf-8"); @@ -4908,8 +4969,9 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, case SQL_WLONGVARCHAR: case SQL_GUID: format = "U"; - arrowColumnProducer->varVal = std::make_unique(arrowBatchSize + 1); - arrowColumnProducer->varData.resize(arrowBatchSize * 42); + arrowColumnProducer->varVal = + AllocateUniqueArray(offsetCount, "Arrow offset buffer is too large"); + arrowColumnProducer->varData.resize(initialVarDataSize); columnVarLen[i] = true; // start at offset 0 arrowColumnProducer->varVal[0] = 0; @@ -4920,8 +4982,9 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, case SQL_VARBINARY: case SQL_LONGVARBINARY: format = "Z"; - arrowColumnProducer->varVal = std::make_unique(arrowBatchSize + 1); - arrowColumnProducer->varData.resize(arrowBatchSize * 42); + arrowColumnProducer->varVal = + AllocateUniqueArray(offsetCount, "Arrow offset buffer is too large"); + arrowColumnProducer->varData.resize(initialVarDataSize); columnVarLen[i] = true; // start at offset 0 arrowColumnProducer->varVal[0] = 0; @@ -4929,33 +4992,39 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, break; case SQL_TINYINT: format = "C"; - arrowColumnProducer->uint8Val = std::make_unique(arrowBatchSize); + arrowColumnProducer->uint8Val = + AllocateUniqueArray(batchSize, "Arrow value buffer is too large"); arrowColumnProducer->ptrValueBuffer = arrowColumnProducer->uint8Val.get(); break; case SQL_SMALLINT: format = "s"; - arrowColumnProducer->int16Val = std::make_unique(arrowBatchSize); + arrowColumnProducer->int16Val = + AllocateUniqueArray(batchSize, "Arrow value buffer is too large"); arrowColumnProducer->ptrValueBuffer = arrowColumnProducer->int16Val.get(); break; case SQL_INTEGER: format = "i"; - arrowColumnProducer->int32Val = std::make_unique(arrowBatchSize); + arrowColumnProducer->int32Val = + AllocateUniqueArray(batchSize, "Arrow value buffer is too large"); arrowColumnProducer->ptrValueBuffer = arrowColumnProducer->int32Val.get(); break; case SQL_BIGINT: format = "l"; - arrowColumnProducer->int64Val = std::make_unique(arrowBatchSize); + arrowColumnProducer->int64Val = + AllocateUniqueArray(batchSize, "Arrow value buffer is too large"); arrowColumnProducer->ptrValueBuffer = arrowColumnProducer->int64Val.get(); break; case SQL_REAL: format = "f"; - arrowColumnProducer->float32Val = std::make_unique(arrowBatchSize); + arrowColumnProducer->float32Val = + AllocateUniqueArray(batchSize, "Arrow value buffer is too large"); arrowColumnProducer->ptrValueBuffer = arrowColumnProducer->float32Val.get(); break; case SQL_FLOAT: case SQL_DOUBLE: format = "g"; - arrowColumnProducer->float64Val = std::make_unique(arrowBatchSize); + arrowColumnProducer->float64Val = + AllocateUniqueArray(batchSize, "Arrow value buffer is too large"); arrowColumnProducer->ptrValueBuffer = arrowColumnProducer->float64Val.get(); break; case SQL_DECIMAL: @@ -4968,7 +5037,8 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, arrowSchemaPrivateData[i]->format = std::make_unique(formatLen); std::memcpy(arrowSchemaPrivateData[i]->format.get(), formatStr.c_str(), formatLen); format = arrowSchemaPrivateData[i]->format.get(); - arrowColumnProducer->decimalVal = std::make_unique(arrowBatchSize); + arrowColumnProducer->decimalVal = + AllocateUniqueArray(batchSize, "Arrow value buffer is too large"); arrowColumnProducer->ptrValueBuffer = arrowColumnProducer->decimalVal.get(); break; } @@ -4976,28 +5046,33 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, case SQL_TYPE_TIMESTAMP: case SQL_DATETIME: format = "tsu:"; - arrowColumnProducer->tsMicroVal = std::make_unique(arrowBatchSize); + arrowColumnProducer->tsMicroVal = + AllocateUniqueArray(batchSize, "Arrow value buffer is too large"); arrowColumnProducer->ptrValueBuffer = arrowColumnProducer->tsMicroVal.get(); break; case SQL_SS_TIMESTAMPOFFSET: format = "tsu:+00:00"; - arrowColumnProducer->tsMicroVal = std::make_unique(arrowBatchSize); + arrowColumnProducer->tsMicroVal = + AllocateUniqueArray(batchSize, "Arrow value buffer is too large"); arrowColumnProducer->ptrValueBuffer = arrowColumnProducer->tsMicroVal.get(); break; case SQL_TYPE_DATE: format = "tdD"; - arrowColumnProducer->dateVal = std::make_unique(arrowBatchSize); + arrowColumnProducer->dateVal = + AllocateUniqueArray(batchSize, "Arrow value buffer is too large"); arrowColumnProducer->ptrValueBuffer = arrowColumnProducer->dateVal.get(); break; case SQL_SS_TIME2: format = "ttn"; - arrowColumnProducer->timeNanoVal = std::make_unique(arrowBatchSize); + arrowColumnProducer->timeNanoVal = + AllocateUniqueArray(batchSize, "Arrow value buffer is too large"); arrowColumnProducer->ptrValueBuffer = arrowColumnProducer->timeNanoVal.get(); break; case SQL_BIT: format = "b"; - arrowColumnProducer->bitVal = std::make_unique((arrowBatchSize + 7) / 8); - std::memset(arrowColumnProducer->bitVal.get(), 0, (arrowBatchSize + 7) / 8); + arrowColumnProducer->bitVal = + AllocateUniqueArray(bitmapSize, "Arrow bitmap is too large"); + std::memset(arrowColumnProducer->bitVal.get(), 0, bitmapSize); arrowColumnProducer->ptrValueBuffer = arrowColumnProducer->bitVal.get(); break; default: @@ -5018,9 +5093,10 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, std::memcpy(arrowSchemaPrivateData[i]->format.get(), format.c_str(), formatLen); } - arrowColumnProducer->valid = std::make_unique((arrowBatchSize + 7) / 8); + arrowColumnProducer->valid = + AllocateUniqueArray(bitmapSize, "Arrow bitmap is too large"); // Initialize validity bitmap to all valid - std::memset(arrowColumnProducer->valid.get(), 0xFF, (arrowBatchSize + 7) / 8); + std::memset(arrowColumnProducer->valid.get(), 0xFF, bitmapSize); } // Initialize column buffers diff --git a/tests/test_004_cursor.py b/tests/test_004_cursor.py index 40ea4d09d..8dda60476 100644 --- a/tests/test_004_cursor.py +++ b/tests/test_004_cursor.py @@ -1730,6 +1730,34 @@ def test_arraysize(cursor): assert cursor.arraysize == 5, "Arraysize mismatch after change" +@pytest.mark.parametrize("value", [0, -1, 2**31]) +def test_arraysize_rejects_out_of_range_values(value): + cursor = mssql_python.Cursor.__new__(mssql_python.Cursor) + with pytest.raises(ValueError, match="arraysize"): + cursor.arraysize = value + + +@pytest.mark.parametrize("value", [True, False, 1.5, "10"]) +def test_arraysize_rejects_non_integer_values(value): + cursor = mssql_python.Cursor.__new__(mssql_python.Cursor) + with pytest.raises(TypeError, match="arraysize"): + cursor.arraysize = value + + +@pytest.mark.parametrize( + "size_info", + [ + (True,), + (mssql_python.SQL_WVARCHAR, True, 0), + (mssql_python.SQL_DECIMAL, 18, True), + ], +) +def test_setinputsizes_rejects_boolean_sizes(size_info): + cursor = mssql_python.Cursor.__new__(mssql_python.Cursor) + with pytest.raises(ValueError): + cursor.setinputsizes([size_info]) + + def test_description(cursor): """Test description""" cursor.execute("SELECT * FROM #pytest_all_data_types WHERE id = 1") diff --git a/tests/test_024_bulkcopy_arrow.py b/tests/test_024_bulkcopy_arrow.py index a840853ba..ee9858262 100644 --- a/tests/test_024_bulkcopy_arrow.py +++ b/tests/test_024_bulkcopy_arrow.py @@ -202,6 +202,11 @@ def test_batch_size_negative(self): with pytest.raises(ValueError, match="batch_size"): _bare_cursor().bulkcopy_arrow("t", pa.table({"a": [1]}), batch_size=-1) + def test_batch_size_bool_rejected(self): + for flag in (False, True): + with pytest.raises(TypeError, match="batch_size"): + _bare_cursor().bulkcopy_arrow("t", pa.table({"a": [1]}), batch_size=flag) + def test_timeout_wrong_type(self): with pytest.raises(TypeError, match="timeout"): _bare_cursor().bulkcopy_arrow("t", pa.table({"a": [1]}), timeout="30") From 39dea1b8ec01f941dbc0601f601cdae898bdd5f8 Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Mon, 21 Sep 2026 12:57:25 +0530 Subject: [PATCH 02/39] FIX: Update negative Arrow batch size test --- tests/test_004_cursor_arrow.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/test_004_cursor_arrow.py b/tests/test_004_cursor_arrow.py index 6a7d7b5d9..b815e524b 100644 --- a/tests/test_004_cursor_arrow.py +++ b/tests/test_004_cursor_arrow.py @@ -226,10 +226,9 @@ def test_arrow_empty_fetch(cursor: mssql_python.Cursor): def test_arrow_table_batchsize_negative(cursor: mssql_python.Cursor): - tbl = cursor.execute("select 1 a").arrow(batch_size=-42) - assert type(tbl) is pa.Table - assert tbl.num_rows == 0 - assert tbl.num_columns == 1 + cursor.execute("select 1 a") + with pytest.raises(ValueError, match="batch_size"): + cursor.arrow(batch_size=-42) assert cursor.fetchone()[0] == 1 From a02fdf41c37190fac0d3749bf00d4e8f64ae04f7 Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Mon, 21 Sep 2026 14:54:35 +0530 Subject: [PATCH 03/39] FIX: Address native size review findings --- mssql_python/cursor.py | 3 +- mssql_python/pybind/ddbc_bindings.cpp | 253 +++++++++++++++++++------- tests/test_004_cursor.py | 15 +- tests/test_010_pybind_functions.py | 14 ++ 4 files changed, 221 insertions(+), 64 deletions(-) diff --git a/mssql_python/cursor.py b/mssql_python/cursor.py index b364f9091..9ed94cd90 100644 --- a/mssql_python/cursor.py +++ b/mssql_python/cursor.py @@ -51,7 +51,8 @@ SMALLMONEY_MAX: decimal.Decimal = decimal.Decimal("214748.3647") MONEY_MIN: decimal.Decimal = decimal.Decimal("-922337203685477.5808") MONEY_MAX: decimal.Decimal = decimal.Decimal("922337203685477.5807") -MAX_NATIVE_ROW_COUNT: int = 2_147_483_647 +# Bound each native fetch allocation; Arrow initially reserves 42 bytes per variable-width row. +MAX_NATIVE_ROW_COUNT: int = 1_000_000 # SQL BIGINT is a signed 64-bit integer. Ints outside this range have no BIGINT # encoding and must be rejected at detect time on both paths (see _map_sql_type). BIGINT_MIN: int = -(2**63) diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index 8f16d46b4..fd7f5ca41 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -346,6 +346,17 @@ size_t CheckedMultiplySize(size_t left, size_t right, const char* errorMessage) return left * right; } +constexpr int MAX_NATIVE_ROW_COUNT = 1000000; +constexpr size_t MAX_NATIVE_FETCH_BYTES = 256ULL * 1024 * 1024; + +void ValidateNativeRowCount(int value, const char* name, bool allowZero) { + const int minimum = allowZero ? 0 : 1; + if (value < minimum || value > MAX_NATIVE_ROW_COUNT) { + ThrowStdException(std::string(name) + " must be between " + std::to_string(minimum) + + " and " + std::to_string(MAX_NATIVE_ROW_COUNT)); + } +} + template std::unique_ptr AllocateUniqueArray(size_t count, const char* errorMessage) { if (count > std::numeric_limits::max() / sizeof(ElementType)) { @@ -354,6 +365,45 @@ std::unique_ptr AllocateUniqueArray(size_t count, const char* err return std::make_unique(count); } +void ReserveNativeFetchBytes(size_t& reservedBytes, size_t count, size_t elementSize) { + const size_t allocationBytes = + CheckedMultiplySize(count, elementSize, "Native fetch buffer size is too large"); + const size_t totalBytes = + CheckedAddSize(reservedBytes, allocationBytes, "Native fetch buffer size is too large"); + if (totalBytes > MAX_NATIVE_FETCH_BYTES) { + ThrowStdException("Native fetch buffers exceed the 256 MiB allocation limit"); + } + reservedBytes = totalBytes; +} + +template +void ResizeNativeFetchBuffer(std::vector& buffer, size_t count, + size_t& reservedBytes) { + ReserveNativeFetchBytes(reservedBytes, count, sizeof(ElementType)); + buffer.resize(count); +} + +template +void EnsureNativeFetchBufferSize(std::vector& buffer, size_t requiredSize, + size_t& reservedBytes) { + if (buffer.size() >= requiredSize) { + return; + } + size_t newSize = std::max(buffer.size(), 1); + while (newSize < requiredSize) { + newSize = CheckedMultiplySize(newSize, 2, "Native fetch buffer size is too large"); + } + ReserveNativeFetchBytes(reservedBytes, newSize - buffer.size(), sizeof(ElementType)); + buffer.resize(newSize); +} + +template +std::unique_ptr AllocateArrowArray(size_t count, size_t& reservedBytes, + const char* errorMessage) { + ReserveNativeFetchBytes(reservedBytes, count, sizeof(ElementType)); + return AllocateUniqueArray(count, errorMessage); +} + std::string DescribeChar(unsigned char ch) { if (ch >= 32 && ch <= 126) { return std::string("'") + static_cast(ch) + "'"; @@ -2416,6 +2466,12 @@ SQLRETURN BindParameterArray(SqlHandle& handle, SQLHANDLE hStmt, const py::list& std::memset(charArray + i * (info.columnSize + 1), 0, info.columnSize + 1); } else { + if (info.paramCType == SQL_C_BINARY && + !py::isinstance(columnValues[i]) && + !py::isinstance(columnValues[i])) { + ThrowStdException(MakeParamMismatchErrorStr(info.paramCType, + paramIndex)); + } std::string encodedStr; if (py::isinstance(columnValues[i])) { @@ -4037,8 +4093,18 @@ 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, size_t& reservedBytes, + int charCtype = SQL_C_WCHAR) { PERF_TIMER("SQLBindColums"); + struct UnbindOnFailure { + SQLHSTMT hStmt; + bool active = true; + ~UnbindOnFailure() { + if (active) { + SQLFreeStmt_ptr(hStmt, SQL_UNBIND); + } + } + } unbindOnFailure{hStmt}; SQLRETURN ret = SQL_SUCCESS; const bool useWideChar = (charCtype == SQL_C_WCHAR); // Bind columns based on their data types @@ -4056,7 +4122,10 @@ SQLRETURN SQLBindColums(SQLHSTMT hStmt, ColumnBuffers& buffers, py::list& column // Bind VARCHAR columns as SQL_C_WCHAR so the ODBC driver // returns UTF-16 data, avoiding code-page decode issues. uint64_t fetchBufferSize = columnSize + 1 /*null-terminator*/; - buffers.wcharBuffers[col - 1].resize(fetchSize * fetchBufferSize); + ResizeNativeFetchBuffer(buffers.wcharBuffers[col - 1], + CheckedMultiplySize(fetchSize, fetchBufferSize, + "Native fetch buffer is too large"), + reservedBytes); ret = SQLBindCol_ptr( hStmt, col, SQL_C_WCHAR, buffers.wcharBuffers[col - 1].data(), fetchBufferSize * sizeof(SQLWCHAR), buffers.indicators[col - 1].data()); @@ -4067,7 +4136,10 @@ SQLRETURN SQLBindColums(SQLHSTMT hStmt, ColumnBuffers& buffers, py::list& column #else uint64_t fetchBufferSize = columnSize + 1 /*null-terminator*/; #endif - buffers.charBuffers[col - 1].resize(fetchSize * fetchBufferSize); + ResizeNativeFetchBuffer(buffers.charBuffers[col - 1], + CheckedMultiplySize(fetchSize, fetchBufferSize, + "Native fetch buffer is too large"), + reservedBytes); ret = SQLBindCol_ptr( hStmt, col, SQL_C_CHAR, buffers.charBuffers[col - 1].data(), fetchBufferSize * sizeof(SQLCHAR), buffers.indicators[col - 1].data()); @@ -4081,48 +4153,55 @@ SQLRETURN SQLBindColums(SQLHSTMT hStmt, ColumnBuffers& buffers, py::list& column // suffice HandleZeroColumnSizeAtFetch(columnSize); uint64_t fetchBufferSize = columnSize + 1 /*null-terminator*/; - buffers.wcharBuffers[col - 1].resize(fetchSize * fetchBufferSize); + ResizeNativeFetchBuffer(buffers.wcharBuffers[col - 1], + CheckedMultiplySize(fetchSize, fetchBufferSize, + "Native fetch buffer is too large"), + reservedBytes); ret = SQLBindCol_ptr(hStmt, col, SQL_C_WCHAR, buffers.wcharBuffers[col - 1].data(), fetchBufferSize * sizeof(SQLWCHAR), buffers.indicators[col - 1].data()); break; } case SQL_INTEGER: - buffers.intBuffers[col - 1].resize(fetchSize); + ResizeNativeFetchBuffer(buffers.intBuffers[col - 1], fetchSize, reservedBytes); ret = SQLBindCol_ptr(hStmt, col, SQL_C_SLONG, buffers.intBuffers[col - 1].data(), sizeof(SQLINTEGER), buffers.indicators[col - 1].data()); break; case SQL_SMALLINT: - buffers.smallIntBuffers[col - 1].resize(fetchSize); + ResizeNativeFetchBuffer(buffers.smallIntBuffers[col - 1], fetchSize, reservedBytes); ret = SQLBindCol_ptr(hStmt, col, SQL_C_SSHORT, buffers.smallIntBuffers[col - 1].data(), sizeof(SQLSMALLINT), buffers.indicators[col - 1].data()); break; case SQL_TINYINT: - buffers.charBuffers[col - 1].resize(fetchSize); + ResizeNativeFetchBuffer(buffers.charBuffers[col - 1], fetchSize, reservedBytes); ret = SQLBindCol_ptr(hStmt, col, SQL_C_TINYINT, buffers.charBuffers[col - 1].data(), sizeof(SQLCHAR), buffers.indicators[col - 1].data()); break; case SQL_BIT: - buffers.charBuffers[col - 1].resize(fetchSize); + ResizeNativeFetchBuffer(buffers.charBuffers[col - 1], fetchSize, reservedBytes); ret = SQLBindCol_ptr(hStmt, col, SQL_C_BIT, buffers.charBuffers[col - 1].data(), sizeof(SQLCHAR), buffers.indicators[col - 1].data()); break; case SQL_REAL: - buffers.realBuffers[col - 1].resize(fetchSize); + ResizeNativeFetchBuffer(buffers.realBuffers[col - 1], fetchSize, reservedBytes); ret = SQLBindCol_ptr(hStmt, col, SQL_C_FLOAT, buffers.realBuffers[col - 1].data(), sizeof(SQLREAL), buffers.indicators[col - 1].data()); break; case SQL_DECIMAL: case SQL_NUMERIC: - buffers.charBuffers[col - 1].resize(fetchSize * MAX_DIGITS_IN_NUMERIC); + ResizeNativeFetchBuffer( + buffers.charBuffers[col - 1], + CheckedMultiplySize(fetchSize, MAX_DIGITS_IN_NUMERIC, + "Native fetch buffer is too large"), + reservedBytes); ret = SQLBindCol_ptr(hStmt, col, SQL_C_CHAR, buffers.charBuffers[col - 1].data(), MAX_DIGITS_IN_NUMERIC * sizeof(SQLCHAR), buffers.indicators[col - 1].data()); break; case SQL_DOUBLE: case SQL_FLOAT: - buffers.doubleBuffers[col - 1].resize(fetchSize); + ResizeNativeFetchBuffer(buffers.doubleBuffers[col - 1], fetchSize, reservedBytes); ret = SQLBindCol_ptr(hStmt, col, SQL_C_DOUBLE, buffers.doubleBuffers[col - 1].data(), sizeof(SQLDOUBLE), buffers.indicators[col - 1].data()); @@ -4130,31 +4209,32 @@ SQLRETURN SQLBindColums(SQLHSTMT hStmt, ColumnBuffers& buffers, py::list& column case SQL_TIMESTAMP: case SQL_TYPE_TIMESTAMP: case SQL_DATETIME: - buffers.timestampBuffers[col - 1].resize(fetchSize); + ResizeNativeFetchBuffer(buffers.timestampBuffers[col - 1], fetchSize, + reservedBytes); ret = SQLBindCol_ptr( hStmt, col, SQL_C_TYPE_TIMESTAMP, buffers.timestampBuffers[col - 1].data(), sizeof(SQL_TIMESTAMP_STRUCT), buffers.indicators[col - 1].data()); break; case SQL_BIGINT: - buffers.bigIntBuffers[col - 1].resize(fetchSize); + ResizeNativeFetchBuffer(buffers.bigIntBuffers[col - 1], fetchSize, reservedBytes); ret = SQLBindCol_ptr(hStmt, col, SQL_C_SBIGINT, buffers.bigIntBuffers[col - 1].data(), sizeof(SQLBIGINT), buffers.indicators[col - 1].data()); break; case SQL_TYPE_DATE: - buffers.dateBuffers[col - 1].resize(fetchSize); + ResizeNativeFetchBuffer(buffers.dateBuffers[col - 1], fetchSize, reservedBytes); ret = SQLBindCol_ptr(hStmt, col, SQL_C_TYPE_DATE, buffers.dateBuffers[col - 1].data(), sizeof(SQL_DATE_STRUCT), buffers.indicators[col - 1].data()); break; case SQL_SS_TIME2: - buffers.timeBuffers[col - 1].resize(fetchSize); + ResizeNativeFetchBuffer(buffers.timeBuffers[col - 1], fetchSize, reservedBytes); ret = SQLBindCol_ptr(hStmt, col, SQL_C_SS_TIME2, buffers.timeBuffers[col - 1].data(), sizeof(SQL_SS_TIME2_STRUCT), buffers.indicators[col - 1].data()); break; case SQL_GUID: - buffers.guidBuffers[col - 1].resize(fetchSize); + ResizeNativeFetchBuffer(buffers.guidBuffers[col - 1], fetchSize, reservedBytes); ret = SQLBindCol_ptr(hStmt, col, SQL_C_GUID, buffers.guidBuffers[col - 1].data(), sizeof(SQLGUID), buffers.indicators[col - 1].data()); break; @@ -4165,12 +4245,16 @@ SQLRETURN SQLBindColums(SQLHSTMT hStmt, ColumnBuffers& buffers, py::list& column // TODO: handle variable length data correctly. This logic wont // suffice HandleZeroColumnSizeAtFetch(columnSize); - buffers.charBuffers[col - 1].resize(fetchSize * columnSize); + ResizeNativeFetchBuffer(buffers.charBuffers[col - 1], + CheckedMultiplySize(fetchSize, columnSize, + "Native fetch buffer is too large"), + reservedBytes); ret = SQLBindCol_ptr(hStmt, col, SQL_C_BINARY, buffers.charBuffers[col - 1].data(), columnSize, buffers.indicators[col - 1].data()); break; case SQL_SS_TIMESTAMPOFFSET: - buffers.datetimeoffsetBuffers[col - 1].resize(fetchSize); + ResizeNativeFetchBuffer(buffers.datetimeoffsetBuffers[col - 1], fetchSize, + reservedBytes); ret = SQLBindCol_ptr(hStmt, col, SQL_C_SS_TIMESTAMPOFFSET, buffers.datetimeoffsetBuffers[col - 1].data(), sizeof(DateTimeOffset) * fetchSize, @@ -4195,6 +4279,7 @@ SQLRETURN SQLBindColums(SQLHSTMT hStmt, ColumnBuffers& buffers, py::list& column return ret; } } + unbindOnFailure.active = false; return ret; } @@ -4678,6 +4763,7 @@ SQLRETURN FetchMany_wrap(SqlHandlePtr StatementHandle, py::list& rows, int fetch const std::string& wcharEncoding = "utf-16le", int charCtype = SQL_C_WCHAR) { PERF_TIMER("FetchMany_wrap"); + ValidateNativeRowCount(fetchSize, "Fetch size", false); // 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. charCtype = EffectiveCharCtypeForFetch(charCtype, charEncoding); @@ -4733,10 +4819,16 @@ SQLRETURN FetchMany_wrap(SqlHandlePtr StatementHandle, py::list& rows, int fetch } // Initialize column buffers + size_t reservedBytes = 0; + ReserveNativeFetchBytes( + reservedBytes, + CheckedMultiplySize(static_cast(numCols), static_cast(fetchSize), + "Native fetch indicator buffer is too large"), + sizeof(SQLLEN)); ColumnBuffers buffers(numCols, fetchSize); // Bind columns - ret = SQLBindColums(hStmt, buffers, columnNames, numCols, fetchSize, charCtype); + ret = SQLBindColums(hStmt, buffers, columnNames, numCols, fetchSize, reservedBytes, charCtype); if (!SQL_SUCCEEDED(ret)) { LOG("FetchMany_wrap: Error when binding columns - SQLRETURN=%d", ret); return ret; @@ -4776,7 +4868,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, size_t& reservedBytes) { size_t start = 0; size_t end = 0; @@ -4796,7 +4888,7 @@ SQLRETURN GetDataVar(SQLHSTMT hStmt, SQLUSMALLINT colNumber, SQLSMALLINT cType, // Ensure initial buffer has space for at least the null terminator if (dataVec.size() < sizeNullTerminator) { - dataVec.resize(sizeNullTerminator); + ResizeNativeFetchBuffer(dataVec, sizeNullTerminator, reservedBytes); } while (true) { @@ -4832,17 +4924,22 @@ SQLRETURN GetDataVar(SQLHSTMT hStmt, SQLUSMALLINT colNumber, SQLSMALLINT cType, // Determine how much more space we need if (localInd < 0) { // SQL_NO_TOTAL: driver doesn't know total size, double the buffer - end = dataVec.size() * 2; + end = CheckedMultiplySize(dataVec.size(), 2, + "Native fetch buffer size is too large"); } else { // Driver returned total size: allocate exactly what we need assert(localInd % sizeof(T) == 0); - end = start + static_cast(localInd) / sizeof(T) + sizeNullTerminator; + end = CheckedAddSize( + CheckedAddSize(start, static_cast(localInd) / sizeof(T), + "Native fetch buffer size is too large"), + sizeNullTerminator, "Native fetch buffer size is too large"); } // The next read starts where the null terminator would have been placed start = dataVec.size() - sizeNullTerminator; // Resize buffer for next iteration + ReserveNativeFetchBytes(reservedBytes, end - dataVec.size(), sizeof(T)); dataVec.resize(end); } else { // Unexpected return code @@ -4888,9 +4985,7 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, int arrowBatchSize, int charCtype) { PERF_TIMER("FetchArrowBatch_wrap"); - if (arrowBatchSize < 0) { - ThrowStdException("Arrow batch size must be non-negative"); - } + ValidateNativeRowCount(arrowBatchSize, "Arrow batch size", true); const size_t batchSize = static_cast(arrowBatchSize); const size_t offsetCount = CheckedAddSize(batchSize, 1, "Arrow batch size is too large"); const size_t initialVarDataSize = @@ -4926,6 +5021,7 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, std::vector columnNullable(numCols); std::vector columnVarLen(numCols, false); std::vector nullCounts(numCols, 0); + size_t reservedBytes = 0; std::vector> arrowArrayPrivateData(numCols); std::vector> arrowSchemaPrivateData(numCols); @@ -4970,8 +5066,10 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, case SQL_GUID: format = "U"; arrowColumnProducer->varVal = - AllocateUniqueArray(offsetCount, "Arrow offset buffer is too large"); - arrowColumnProducer->varData.resize(initialVarDataSize); + AllocateArrowArray(offsetCount, reservedBytes, + "Arrow offset buffer is too large"); + ResizeNativeFetchBuffer(arrowColumnProducer->varData, initialVarDataSize, + reservedBytes); columnVarLen[i] = true; // start at offset 0 arrowColumnProducer->varVal[0] = 0; @@ -4983,8 +5081,10 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, case SQL_LONGVARBINARY: format = "Z"; arrowColumnProducer->varVal = - AllocateUniqueArray(offsetCount, "Arrow offset buffer is too large"); - arrowColumnProducer->varData.resize(initialVarDataSize); + AllocateArrowArray(offsetCount, reservedBytes, + "Arrow offset buffer is too large"); + ResizeNativeFetchBuffer(arrowColumnProducer->varData, initialVarDataSize, + reservedBytes); columnVarLen[i] = true; // start at offset 0 arrowColumnProducer->varVal[0] = 0; @@ -4993,38 +5093,44 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, case SQL_TINYINT: format = "C"; arrowColumnProducer->uint8Val = - AllocateUniqueArray(batchSize, "Arrow value buffer is too large"); + AllocateArrowArray(batchSize, reservedBytes, + "Arrow value buffer is too large"); arrowColumnProducer->ptrValueBuffer = arrowColumnProducer->uint8Val.get(); break; case SQL_SMALLINT: format = "s"; arrowColumnProducer->int16Val = - AllocateUniqueArray(batchSize, "Arrow value buffer is too large"); + AllocateArrowArray(batchSize, reservedBytes, + "Arrow value buffer is too large"); arrowColumnProducer->ptrValueBuffer = arrowColumnProducer->int16Val.get(); break; case SQL_INTEGER: format = "i"; arrowColumnProducer->int32Val = - AllocateUniqueArray(batchSize, "Arrow value buffer is too large"); + AllocateArrowArray(batchSize, reservedBytes, + "Arrow value buffer is too large"); arrowColumnProducer->ptrValueBuffer = arrowColumnProducer->int32Val.get(); break; case SQL_BIGINT: format = "l"; arrowColumnProducer->int64Val = - AllocateUniqueArray(batchSize, "Arrow value buffer is too large"); + AllocateArrowArray(batchSize, reservedBytes, + "Arrow value buffer is too large"); arrowColumnProducer->ptrValueBuffer = arrowColumnProducer->int64Val.get(); break; case SQL_REAL: format = "f"; arrowColumnProducer->float32Val = - AllocateUniqueArray(batchSize, "Arrow value buffer is too large"); + AllocateArrowArray(batchSize, reservedBytes, + "Arrow value buffer is too large"); arrowColumnProducer->ptrValueBuffer = arrowColumnProducer->float32Val.get(); break; case SQL_FLOAT: case SQL_DOUBLE: format = "g"; arrowColumnProducer->float64Val = - AllocateUniqueArray(batchSize, "Arrow value buffer is too large"); + AllocateArrowArray(batchSize, reservedBytes, + "Arrow value buffer is too large"); arrowColumnProducer->ptrValueBuffer = arrowColumnProducer->float64Val.get(); break; case SQL_DECIMAL: @@ -5038,7 +5144,8 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, std::memcpy(arrowSchemaPrivateData[i]->format.get(), formatStr.c_str(), formatLen); format = arrowSchemaPrivateData[i]->format.get(); arrowColumnProducer->decimalVal = - AllocateUniqueArray(batchSize, "Arrow value buffer is too large"); + AllocateArrowArray(batchSize, reservedBytes, + "Arrow value buffer is too large"); arrowColumnProducer->ptrValueBuffer = arrowColumnProducer->decimalVal.get(); break; } @@ -5047,31 +5154,36 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, case SQL_DATETIME: format = "tsu:"; arrowColumnProducer->tsMicroVal = - AllocateUniqueArray(batchSize, "Arrow value buffer is too large"); + AllocateArrowArray(batchSize, reservedBytes, + "Arrow value buffer is too large"); arrowColumnProducer->ptrValueBuffer = arrowColumnProducer->tsMicroVal.get(); break; case SQL_SS_TIMESTAMPOFFSET: format = "tsu:+00:00"; arrowColumnProducer->tsMicroVal = - AllocateUniqueArray(batchSize, "Arrow value buffer is too large"); + AllocateArrowArray(batchSize, reservedBytes, + "Arrow value buffer is too large"); arrowColumnProducer->ptrValueBuffer = arrowColumnProducer->tsMicroVal.get(); break; case SQL_TYPE_DATE: format = "tdD"; arrowColumnProducer->dateVal = - AllocateUniqueArray(batchSize, "Arrow value buffer is too large"); + AllocateArrowArray(batchSize, reservedBytes, + "Arrow value buffer is too large"); arrowColumnProducer->ptrValueBuffer = arrowColumnProducer->dateVal.get(); break; case SQL_SS_TIME2: format = "ttn"; arrowColumnProducer->timeNanoVal = - AllocateUniqueArray(batchSize, "Arrow value buffer is too large"); + AllocateArrowArray(batchSize, reservedBytes, + "Arrow value buffer is too large"); arrowColumnProducer->ptrValueBuffer = arrowColumnProducer->timeNanoVal.get(); break; case SQL_BIT: format = "b"; arrowColumnProducer->bitVal = - AllocateUniqueArray(bitmapSize, "Arrow bitmap is too large"); + AllocateArrowArray(bitmapSize, reservedBytes, + "Arrow bitmap is too large"); std::memset(arrowColumnProducer->bitVal.get(), 0, bitmapSize); arrowColumnProducer->ptrValueBuffer = arrowColumnProducer->bitVal.get(); break; @@ -5094,16 +5206,22 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, } arrowColumnProducer->valid = - AllocateUniqueArray(bitmapSize, "Arrow bitmap is too large"); + AllocateArrowArray(bitmapSize, reservedBytes, "Arrow bitmap is too large"); // Initialize validity bitmap to all valid std::memset(arrowColumnProducer->valid.get(), 0xFF, bitmapSize); } // Initialize column buffers + ReserveNativeFetchBytes( + reservedBytes, + CheckedMultiplySize(static_cast(numCols), static_cast(fetchSize), + "Native fetch indicator buffer is too large"), + sizeof(SQLLEN)); ColumnBuffers buffers(numCols, fetchSize); if (!hasLobColumns && fetchSize > 0) { - ret = SQLBindColums(hStmt, buffers, columnNames, numCols, fetchSize, charCtype); + ret = SQLBindColums(hStmt, buffers, columnNames, numCols, fetchSize, reservedBytes, + charCtype); if (!SQL_SUCCEEDED(ret)) { LOG("Error when binding columns"); return ret; @@ -5153,7 +5271,7 @@ 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(), reservedBytes); if (!SQL_SUCCEEDED(ret)) { LOG("Error fetching BINARY LOB for column %d", idxCol + 1); return ret; @@ -5166,7 +5284,7 @@ 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(), reservedBytes); if (!SQL_SUCCEEDED(ret)) { LOG("Error fetching CHAR LOB data for column %d", idxCol + 1); return ret; @@ -5181,7 +5299,7 @@ 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(), reservedBytes); if (!SQL_SUCCEEDED(ret)) { LOG("Error fetching WCHAR LOB data for column %d", idxCol + 1); return ret; @@ -5405,9 +5523,10 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, uint64_t fetchBufferSize = columnSize /* bytes are not null terminated */; auto target_vec = &arrowColumnProducer->varData; auto start = arrowColumnProducer->varVal[idxRowArrow]; - while (target_vec->size() < start + dataLen) { - target_vec->resize(target_vec->size() * 2); - } + EnsureNativeFetchBufferSize( + *target_vec, + CheckedAddSize(start, dataLen, "Arrow value buffer is too large"), + reservedBytes); std::memcpy(&(*target_vec)[start], &buffers.charBuffers[idxCol][idxRowSql * fetchBufferSize], @@ -5426,9 +5545,10 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, #endif auto target_vec = &arrowColumnProducer->varData; auto start = arrowColumnProducer->varVal[idxRowArrow]; - while (target_vec->size() < start + dataLen) { - target_vec->resize(target_vec->size() * 2); - } + EnsureNativeFetchBufferSize( + *target_vec, + CheckedAddSize(start, dataLen, "Arrow value buffer is too large"), + reservedBytes); std::memcpy(&(*target_vec)[start], &buffers.charBuffers[idxCol][idxRowSql * fetchBufferSize], @@ -5452,11 +5572,13 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, static_assert(sizeof(SQLWCHAR) == sizeof(char16_t)); static_assert(alignof(SQLWCHAR) == alignof(char16_t)); const auto* utf16Source = reinterpret_cast(wcharSource); - size_t maxUtf8Size = dataLenW * 3; - - while (target_vec->size() < start + maxUtf8Size) { - target_vec->resize(target_vec->size() * 2); - } + size_t maxUtf8Size = + CheckedMultiplySize(dataLenW, 3, "Arrow value buffer is too large"); + EnsureNativeFetchBufferSize( + *target_vec, + CheckedAddSize(start, maxUtf8Size, + "Arrow value buffer is too large"), + reservedBytes); size_t bytesWritten = simdutf::convert_utf16le_to_utf8_with_replacement( utf16Source, dataLenW, @@ -5473,9 +5595,10 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, auto start = arrowColumnProducer->varVal[idxRowArrow]; // Ensure buffer has space for the GUID string + null terminator - while (target_vec->size() < start + 37) { - target_vec->resize(target_vec->size() * 2); - } + EnsureNativeFetchBufferSize( + *target_vec, + CheckedAddSize(start, 37, "Arrow value buffer is too large"), + reservedBytes); // Get the GUID from the buffer const SQLGUID& guidValue = buffers.guidBuffers[idxCol][idxRowSql]; @@ -5917,10 +6040,16 @@ SQLRETURN FetchAll_wrap(SqlHandlePtr StatementHandle, py::list& rows, } LOG("FetchAll_wrap: Fetching data in batch sizes of %d", fetchSize); + size_t reservedBytes = 0; + ReserveNativeFetchBytes( + reservedBytes, + CheckedMultiplySize(static_cast(numCols), static_cast(fetchSize), + "Native fetch indicator buffer is too large"), + sizeof(SQLLEN)); ColumnBuffers buffers(numCols, fetchSize); // Bind columns - ret = SQLBindColums(hStmt, buffers, columnNames, numCols, fetchSize, charCtype); + ret = SQLBindColums(hStmt, buffers, columnNames, numCols, fetchSize, reservedBytes, charCtype); if (!SQL_SUCCEEDED(ret)) { LOG("FetchAll_wrap: Error when binding columns - SQLRETURN=%d", ret); return ret; diff --git a/tests/test_004_cursor.py b/tests/test_004_cursor.py index 8dda60476..54b5d3504 100644 --- a/tests/test_004_cursor.py +++ b/tests/test_004_cursor.py @@ -1730,7 +1730,7 @@ def test_arraysize(cursor): assert cursor.arraysize == 5, "Arraysize mismatch after change" -@pytest.mark.parametrize("value", [0, -1, 2**31]) +@pytest.mark.parametrize("value", [0, -1, 1_000_001]) def test_arraysize_rejects_out_of_range_values(value): cursor = mssql_python.Cursor.__new__(mssql_python.Cursor) with pytest.raises(ValueError, match="arraysize"): @@ -1758,6 +1758,19 @@ def test_setinputsizes_rejects_boolean_sizes(size_info): cursor.setinputsizes([size_info]) +def test_executemany_rejects_text_for_binary_parameter(cursor): + cursor.setinputsizes([(mssql_python.SQL_VARBINARY, 10, 0)]) + with pytest.raises(RuntimeError, match="object type does not match"): + cursor.executemany("SELECT CAST(? AS VARBINARY(10))", [("text",)]) + + +def test_fetchmany_rejects_excessive_native_buffer(cursor): + cursor.execute("SELECT CAST('x' AS VARCHAR(8000))") + with pytest.raises(RuntimeError, match="256 MiB allocation limit"): + cursor.fetchmany(100_000) + assert cursor.fetchone()[0] == "x" + + def test_description(cursor): """Test description""" cursor.execute("SELECT * FROM #pytest_all_data_types WHERE id = 1") diff --git a/tests/test_010_pybind_functions.py b/tests/test_010_pybind_functions.py index 106b64ca3..afed788e4 100644 --- a/tests/test_010_pybind_functions.py +++ b/tests/test_010_pybind_functions.py @@ -37,6 +37,20 @@ ) +@pytest.mark.skipif(not DDBC_AVAILABLE, reason="ddbc_bindings not available") +@pytest.mark.parametrize("fetch_size", [-1, 1_000_001]) +def test_fetchmany_rejects_unsafe_size_before_handle_access(fetch_size): + with pytest.raises(RuntimeError, match="Fetch size"): + ddbc.DDBCSQLFetchMany(None, [], fetch_size) + + +@pytest.mark.skipif(not DDBC_AVAILABLE, reason="ddbc_bindings not available") +@pytest.mark.parametrize("batch_size", [-1, 1_000_001]) +def test_arrow_batch_rejects_unsafe_size_before_handle_access(batch_size): + with pytest.raises(RuntimeError, match="Arrow batch size"): + ddbc.DDBCSQLFetchArrowBatch(None, [], batch_size, 0) + + @pytest.mark.skipif(not DDBC_AVAILABLE, reason="ddbc_bindings not available") class TestPybindModuleInfo: """Test module information and architecture detection.""" From 3038841dffbdb4ff9a4f593052237d32fc5b54e1 Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Tue, 22 Sep 2026 11:12:18 +0530 Subject: [PATCH 04/39] FIX: Bound native parameter and Arrow source buffers --- mssql_python/cursor.py | 5 +- mssql_python/pybind/ddbc_bindings.cpp | 118 +++++++++++++++++++++++--- tests/test_004_cursor.py | 12 +++ 3 files changed, 120 insertions(+), 15 deletions(-) diff --git a/mssql_python/cursor.py b/mssql_python/cursor.py index 9ed94cd90..54dcbc174 100644 --- a/mssql_python/cursor.py +++ b/mssql_python/cursor.py @@ -53,6 +53,7 @@ MONEY_MAX: decimal.Decimal = decimal.Decimal("922337203685477.5807") # Bound each native fetch allocation; Arrow initially reserves 42 bytes per variable-width row. MAX_NATIVE_ROW_COUNT: int = 1_000_000 +MAX_NATIVE_PARAMETER_SIZE: int = 256 * 1024 * 1024 # SQL BIGINT is a signed 64-bit integer. Ints outside this range have no BIGINT # encoding and must be rejected at detect time on both paths (see _map_sql_type). BIGINT_MIN: int = -(2**63) @@ -1223,9 +1224,11 @@ def setinputsizes(self, sizes: List[Union[int, tuple]]) -> None: isinstance(column_size, bool) or not isinstance(column_size, int) or column_size < 0 + or column_size > MAX_NATIVE_PARAMETER_SIZE ): raise ValueError( - f"Invalid column size: {column_size}. Must be a non-negative integer." + f"Invalid column size: {column_size}. Must be a non-negative integer " + f"no greater than {MAX_NATIVE_PARAMETER_SIZE}." ) if ( diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index fd7f5ca41..aa5e276cf 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -348,6 +348,7 @@ size_t CheckedMultiplySize(size_t left, size_t right, const char* errorMessage) constexpr int MAX_NATIVE_ROW_COUNT = 1000000; constexpr size_t MAX_NATIVE_FETCH_BYTES = 256ULL * 1024 * 1024; +constexpr size_t MAX_NATIVE_PARAMETER_BYTES = 256ULL * 1024 * 1024; void ValidateNativeRowCount(int value, const char* name, bool allowZero) { const int minimum = allowZero ? 0 : 1; @@ -357,6 +358,79 @@ void ValidateNativeRowCount(int value, const char* name, bool allowZero) { } } +void ReserveNativeParameterBytes(size_t& reservedBytes, size_t count, size_t elementSize) { + const size_t allocationBytes = + CheckedMultiplySize(count, elementSize, "Parameter buffer size is too large"); + reservedBytes = + CheckedAddSize(reservedBytes, allocationBytes, "Parameter buffer size is too large"); + if (reservedBytes > MAX_NATIVE_PARAMETER_BYTES) { + ThrowStdException("Parameter buffers exceed the 256 MiB allocation limit"); + } +} + +size_t ParameterArrayElementSize(const ParamInfo& info) { + switch (info.paramCType) { + case SQL_C_LONG: + return sizeof(int); + case SQL_C_DOUBLE: + return sizeof(double); + case SQL_C_WCHAR: + return CheckedMultiplySize( + CheckedAddSize(info.columnSize, 1, "Wide-character parameter size is too large"), + sizeof(SQLWCHAR), "Wide-character parameter size is too large"); + case SQL_C_TINYINT: + case SQL_C_UTINYINT: + return sizeof(unsigned char); + case SQL_C_SHORT: + return sizeof(short); + case SQL_C_CHAR: + case SQL_C_BINARY: + return CheckedAddSize(info.columnSize, 1, "Character parameter size is too large"); + case SQL_C_BIT: + return sizeof(char); + case SQL_C_STINYINT: + case SQL_C_USHORT: + return sizeof(unsigned short); + case SQL_C_SBIGINT: + case SQL_C_SLONG: + case SQL_C_UBIGINT: + case SQL_C_ULONG: + return sizeof(int64_t); + case SQL_C_FLOAT: + return sizeof(float); + case SQL_C_TYPE_DATE: + return sizeof(SQL_DATE_STRUCT); + case SQL_C_TYPE_TIME: + return sizeof(SQL_TIME_STRUCT); + case SQL_C_TYPE_TIMESTAMP: + return sizeof(SQL_TIMESTAMP_STRUCT); + case SQL_C_SS_TIMESTAMPOFFSET: + return sizeof(DateTimeOffset); + case SQL_C_NUMERIC: + return sizeof(SQL_NUMERIC_STRUCT); + case SQL_C_GUID: + return sizeof(SQLGUID); + case SQL_C_DEFAULT: + return sizeof(char); + default: + ThrowStdException("Unsupported C type for parameter array allocation"); + } +} + +template +size_t CheckedArrowSourceOffset(const std::vector& buffer, size_t rowIndex, + size_t rowStride, size_t dataBytes) { + const size_t offset = + CheckedMultiplySize(rowIndex, rowStride, "Arrow source offset is too large"); + const size_t rowCapacity = CheckedMultiplySize( + rowStride, sizeof(ElementType), "Arrow source capacity is too large"); + if (offset > buffer.size() || rowStride > buffer.size() - offset || + dataBytes > rowCapacity) { + ThrowStdException("Driver data length exceeds the allocated fetch buffer"); + } + return offset; +} + template std::unique_ptr AllocateUniqueArray(size_t count, const char* errorMessage) { if (count > std::numeric_limits::max() / sizeof(ElementType)) { @@ -2282,6 +2356,12 @@ SQLRETURN BindParameterArray(SqlHandle& handle, SQLHANDLE hStmt, const py::list& try { // GH-627: resolve unknown NULL array param SQL types before binding any param. PreResolveUnknownNullTypes(handle, hStmt, paramInfos); + size_t reservedParameterBytes = 0; + for (const ParamInfo& info : paramInfos) { + ReserveNativeParameterBytes(reservedParameterBytes, paramSetSize, + ParameterArrayElementSize(info)); + ReserveNativeParameterBytes(reservedParameterBytes, paramSetSize, sizeof(SQLLEN)); + } for (int paramIndex = 0; paramIndex < columnwise_params.size(); ++paramIndex) { const py::list& columnValues = columnwise_params[paramIndex].cast(); ParamInfo& info = paramInfos[paramIndex]; @@ -5259,7 +5339,6 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, for (SQLUSMALLINT idxCol = 0; idxCol < numCols; idxCol++) { auto& arrowColumnProducer = arrowArrayPrivateData[idxCol]; auto dataType = dataTypes[idxCol]; - auto columnSize = columnSizes[idxCol]; if (hasLobColumns) { assert(idxRowSql == 0 && "GetData only works one row at a time"); @@ -5520,17 +5599,21 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, case SQL_BINARY: case SQL_VARBINARY: case SQL_LONGVARBINARY: { - uint64_t fetchBufferSize = columnSize /* bytes are not null terminated */; auto target_vec = &arrowColumnProducer->varData; auto start = arrowColumnProducer->varVal[idxRowArrow]; EnsureNativeFetchBufferSize( *target_vec, CheckedAddSize(start, dataLen, "Arrow value buffer is too large"), reservedBytes); + const size_t sourceStride = hasLobColumns + ? buffers.charBuffers[idxCol].size() + : buffers.charBuffers[idxCol].size() / + static_cast(fetchSize); + const size_t sourceOffset = CheckedArrowSourceOffset( + buffers.charBuffers[idxCol], idxRowSql, sourceStride, dataLen); std::memcpy(&(*target_vec)[start], - &buffers.charBuffers[idxCol][idxRowSql * fetchBufferSize], - dataLen); + &buffers.charBuffers[idxCol][sourceOffset], dataLen); arrowColumnProducer->varVal[idxRowArrow + 1] = start + dataLen; break; } @@ -5538,21 +5621,21 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, case SQL_VARCHAR: case SQL_LONGVARCHAR: { if (charCtype == SQL_C_CHAR) { -#if defined(__APPLE__) || defined(__linux__) - uint64_t fetchBufferSize = columnSize * 4 + 1 /*null-terminator*/; -#else - uint64_t fetchBufferSize = columnSize + 1 /*null-terminator*/; -#endif auto target_vec = &arrowColumnProducer->varData; auto start = arrowColumnProducer->varVal[idxRowArrow]; EnsureNativeFetchBufferSize( *target_vec, CheckedAddSize(start, dataLen, "Arrow value buffer is too large"), reservedBytes); + const size_t sourceStride = hasLobColumns + ? buffers.charBuffers[idxCol].size() + : buffers.charBuffers[idxCol].size() / + static_cast(fetchSize); + const size_t sourceOffset = CheckedArrowSourceOffset( + buffers.charBuffers[idxCol], idxRowSql, sourceStride, dataLen); std::memcpy(&(*target_vec)[start], - &buffers.charBuffers[idxCol][idxRowSql * fetchBufferSize], - dataLen); + &buffers.charBuffers[idxCol][sourceOffset], dataLen); arrowColumnProducer->varVal[idxRowArrow + 1] = start + dataLen; break; } @@ -5563,10 +5646,17 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, case SQL_WVARCHAR: case SQL_WLONGVARCHAR: { // We have previously fetched these as WCHARs, even for SQL_CHAR types. - assert(dataLen % sizeof(SQLWCHAR) == 0); + if (dataLen % sizeof(SQLWCHAR) != 0) { + ThrowStdException("Wide-character data has an invalid byte length"); + } auto dataLenW = dataLen / sizeof(SQLWCHAR); - auto wcharSource = - &buffers.wcharBuffers[idxCol][idxRowSql * (columnSize + 1)]; + const size_t sourceStride = hasLobColumns + ? buffers.wcharBuffers[idxCol].size() + : buffers.wcharBuffers[idxCol].size() / + static_cast(fetchSize); + const size_t sourceOffset = CheckedArrowSourceOffset( + buffers.wcharBuffers[idxCol], idxRowSql, sourceStride, dataLen); + auto wcharSource = &buffers.wcharBuffers[idxCol][sourceOffset]; auto start = arrowColumnProducer->varVal[idxRowArrow]; auto target_vec = &arrowColumnProducer->varData; static_assert(sizeof(SQLWCHAR) == sizeof(char16_t)); diff --git a/tests/test_004_cursor.py b/tests/test_004_cursor.py index 54b5d3504..4441308b9 100644 --- a/tests/test_004_cursor.py +++ b/tests/test_004_cursor.py @@ -1758,6 +1758,18 @@ def test_setinputsizes_rejects_boolean_sizes(size_info): cursor.setinputsizes([size_info]) +def test_setinputsizes_rejects_excessive_column_size(): + cursor = mssql_python.Cursor.__new__(mssql_python.Cursor) + with pytest.raises(ValueError, match="column size"): + cursor.setinputsizes([(mssql_python.SQL_VARCHAR, 1 << 40, 0)]) + + +def test_executemany_rejects_excessive_cumulative_parameter_buffers(cursor): + cursor.setinputsizes([(mssql_python.SQL_WVARCHAR, 50_000_000, 0)] * 3) + with pytest.raises(RuntimeError, match="Parameter buffers exceed the 256 MiB allocation limit"): + cursor.executemany("SELECT ?, ?, ?", [("", "", "")]) + + def test_executemany_rejects_text_for_binary_parameter(cursor): cursor.setinputsizes([(mssql_python.SQL_VARBINARY, 10, 0)]) with pytest.raises(RuntimeError, match="object type does not match"): From 710ff123aa06fc9556a6a021b18f15c07bc32772 Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Fri, 25 Sep 2026 15:18:52 +0530 Subject: [PATCH 05/39] FIX: Address native size validation review --- mssql_python/cursor.py | 15 +++- mssql_python/pybind/ddbc_bindings.cpp | 110 +++++++++++++++++--------- tests/test_004_cursor.py | 7 +- tests/test_004_cursor_arrow.py | 15 ++++ tests/test_024_bulkcopy_arrow.py | 4 + 5 files changed, 109 insertions(+), 42 deletions(-) diff --git a/mssql_python/cursor.py b/mssql_python/cursor.py index 54dcbc174..4c90f52a3 100644 --- a/mssql_python/cursor.py +++ b/mssql_python/cursor.py @@ -1224,7 +1224,14 @@ def setinputsizes(self, sizes: List[Union[int, tuple]]) -> None: isinstance(column_size, bool) or not isinstance(column_size, int) or column_size < 0 - or column_size > MAX_NATIVE_PARAMETER_SIZE + or ( + sql_type + not in ( + ddbc_sql_const.SQL_DECIMAL.value, + ddbc_sql_const.SQL_NUMERIC.value, + ) + and column_size > MAX_NATIVE_PARAMETER_SIZE + ) ): raise ValueError( f"Invalid column size: {column_size}. Must be a non-negative integer " @@ -3426,8 +3433,10 @@ def _bulkcopy_core_and_validate(table_name, batch_size, timeout): raise TypeError( f"batch_size must be a non-negative integer, got {type(batch_size).__name__}" ) - if batch_size < 0: - raise ValueError(f"batch_size must be non-negative, got {batch_size}") + if batch_size < 0 or batch_size > MAX_NATIVE_ROW_COUNT: + raise ValueError( + f"batch_size must be between 0 and {MAX_NATIVE_ROW_COUNT}, got {batch_size}" + ) if not isinstance(timeout, int) or isinstance(timeout, bool): raise TypeError(f"timeout must be a non-negative integer, got {type(timeout).__name__}") diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index aa5e276cf..53087c370 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -415,6 +415,7 @@ size_t ParameterArrayElementSize(const ParamInfo& info) { default: ThrowStdException("Unsupported C type for parameter array allocation"); } + return 0; } template @@ -4736,77 +4737,96 @@ SQLRETURN FetchBatchData(SQLHSTMT hStmt, ColumnBuffers& buffers, py::list& colum // Given a list of columns that are a part of single row in the result set, // calculates the max size of the row // TODO: Move to anonymous namespace, since it is not used outside this file -size_t calculateRowSize(py::list& columnNames, SQLUSMALLINT numCols) { - size_t rowSize = 0; +size_t calculateRowSize(py::list& columnNames, SQLUSMALLINT numCols, + int charCtype = SQL_C_WCHAR) { + size_t rowSize = CheckedMultiplySize(numCols, sizeof(SQLLEN), + "Native fetch indicator row size is too large"); for (SQLUSMALLINT col = 1; col <= numCols; col++) { auto columnMeta = columnNames[col - 1].cast(); SQLSMALLINT dataType = columnMeta["DataType"].cast(); SQLULEN columnSize = columnMeta["ColumnSize"].cast(); + size_t columnBytes = 0; switch (dataType) { case SQL_CHAR: case SQL_VARCHAR: - case SQL_LONGVARCHAR: - rowSize += columnSize; + case SQL_LONGVARCHAR: { + HandleZeroColumnSizeAtFetch(columnSize); + if (charCtype == SQL_C_WCHAR) { + columnBytes = CheckedMultiplySize( + CheckedAddSize(columnSize, 1, "Native fetch row size is too large"), + sizeof(SQLWCHAR), "Native fetch row size is too large"); + } else { +#if defined(__APPLE__) || defined(__linux__) + columnBytes = CheckedAddSize( + CheckedMultiplySize(columnSize, 4, + "Native fetch row size is too large"), + 1, "Native fetch row size is too large"); +#else + columnBytes = + CheckedAddSize(columnSize, 1, "Native fetch row size is too large"); +#endif + } break; + } case SQL_SS_XML: case SQL_WCHAR: case SQL_WVARCHAR: case SQL_WLONGVARCHAR: - rowSize += columnSize * sizeof(SQLWCHAR); + HandleZeroColumnSizeAtFetch(columnSize); + columnBytes = CheckedMultiplySize( + CheckedAddSize(columnSize, 1, "Native fetch row size is too large"), + sizeof(SQLWCHAR), "Native fetch row size is too large"); break; case SQL_INTEGER: - rowSize += sizeof(SQLINTEGER); + columnBytes = sizeof(SQLINTEGER); break; case SQL_SMALLINT: - rowSize += sizeof(SQLSMALLINT); + columnBytes = sizeof(SQLSMALLINT); break; case SQL_REAL: - rowSize += sizeof(SQLREAL); + columnBytes = sizeof(SQLREAL); break; case SQL_FLOAT: - rowSize += sizeof(SQLFLOAT); + columnBytes = sizeof(SQLFLOAT); break; case SQL_DOUBLE: - rowSize += sizeof(SQLDOUBLE); + columnBytes = sizeof(SQLDOUBLE); break; case SQL_DECIMAL: case SQL_NUMERIC: - rowSize += MAX_DIGITS_IN_NUMERIC; + columnBytes = MAX_DIGITS_IN_NUMERIC; break; case SQL_TIMESTAMP: case SQL_TYPE_TIMESTAMP: case SQL_DATETIME: - rowSize += sizeof(SQL_TIMESTAMP_STRUCT); + columnBytes = sizeof(SQL_TIMESTAMP_STRUCT); break; case SQL_BIGINT: - rowSize += sizeof(SQLBIGINT); + columnBytes = sizeof(SQLBIGINT); break; case SQL_TYPE_DATE: - rowSize += sizeof(SQL_DATE_STRUCT); + columnBytes = sizeof(SQL_DATE_STRUCT); break; case SQL_SS_TIME2: - rowSize += sizeof(SQL_SS_TIME2_STRUCT); + columnBytes = sizeof(SQL_SS_TIME2_STRUCT); break; case SQL_GUID: - rowSize += sizeof(SQLGUID); + columnBytes = sizeof(SQLGUID); break; case SQL_TINYINT: case SQL_BIT: - rowSize += sizeof(SQLCHAR); + columnBytes = sizeof(SQLCHAR); break; case SQL_SS_UDT: - rowSize += (static_cast(columnSize) == SQL_NO_TOTAL || columnSize == 0) - ? SQL_MAX_LOB_SIZE - : columnSize; - break; case SQL_BINARY: case SQL_VARBINARY: case SQL_LONGVARBINARY: - rowSize += columnSize; + HandleZeroColumnSizeAtFetch(columnSize); + columnBytes = columnSize; break; case SQL_SS_TIMESTAMPOFFSET: - rowSize += sizeof(DateTimeOffset); + columnBytes = sizeof(DateTimeOffset); break; default: std::string columnName = columnMeta["ColumnName"].cast(); @@ -4817,6 +4837,7 @@ size_t calculateRowSize(py::list& columnNames, SQLUSMALLINT numCols) { ThrowStdException(errorString.str()); break; } + rowSize = CheckedAddSize(rowSize, columnBytes, "Native fetch row size is too large"); } return rowSize; } @@ -5075,8 +5096,8 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, // Fetch narrow char data as SQL_C_CHAR if on Linux/macOS and configured by the user charCtype = EffectiveCharCtypeForFetch(charCtype, "utf-8"); - // An overly large fetch size doesn't seem to help performance - int fetchSize = 64; + // An overly large fetch size doesn't seem to help performance. + int fetchSize = arrowBatchSize > 0 ? std::min(arrowBatchSize, 64) : 1; SQLRETURN ret; SQLHSTMT hStmt = StatementHandle->get(); @@ -5291,15 +5312,28 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, std::memset(arrowColumnProducer->valid.get(), 0xFF, bitmapSize); } - // Initialize column buffers - ReserveNativeFetchBytes( - reservedBytes, - CheckedMultiplySize(static_cast(numCols), static_cast(fetchSize), - "Native fetch indicator buffer is too large"), - sizeof(SQLLEN)); - ColumnBuffers buffers(numCols, fetchSize); + if (arrowBatchSize > 0 && !hasLobColumns) { + const size_t rowSize = calculateRowSize(columnNames, numCols, charCtype); + const size_t remainingBytes = MAX_NATIVE_FETCH_BYTES - reservedBytes; + const size_t rowsWithinBudget = remainingBytes / rowSize; + if (rowsWithinBudget == 0) { + ThrowStdException("Native fetch buffers exceed the 256 MiB allocation limit"); + } + fetchSize = static_cast(std::min(fetchSize, rowsWithinBudget)); + } + + const int bufferFetchSize = arrowBatchSize > 0 ? fetchSize : 0; + if (bufferFetchSize > 0) { + ReserveNativeFetchBytes( + reservedBytes, + CheckedMultiplySize(static_cast(numCols), + static_cast(bufferFetchSize), + "Native fetch indicator buffer is too large"), + sizeof(SQLLEN)); + } + ColumnBuffers buffers(numCols, bufferFetchSize); - if (!hasLobColumns && fetchSize > 0) { + if (!hasLobColumns && bufferFetchSize > 0) { ret = SQLBindColums(hStmt, buffers, columnNames, numCols, fetchSize, reservedBytes, charCtype); if (!SQL_SUCCEEDED(ret)) { @@ -6090,14 +6124,12 @@ SQLRETURN FetchAll_wrap(SqlHandlePtr StatementHandle, py::list& rows, } // No LOBs detected - use binding path with batch fetching - // Define a memory limit (1 GB) - const size_t memoryLimit = 1ULL * 1024 * 1024 * 1024; - size_t totalRowSize = calculateRowSize(columnNames, numCols); + size_t totalRowSize = calculateRowSize(columnNames, numCols, charCtype); - // Calculate fetch size based on the total row size and memory limit + // Calculate fetch size from the same storage budget enforced during binding. size_t numRowsInMemLimit; if (totalRowSize > 0) { - numRowsInMemLimit = static_cast(memoryLimit / totalRowSize); + numRowsInMemLimit = MAX_NATIVE_FETCH_BYTES / totalRowSize; } else { // Handle case where totalRowSize is 0 to avoid division by zero. // This can happen for NVARCHAR(MAX) cols. SQLDescribeCol returns 0 @@ -6128,6 +6160,8 @@ SQLRETURN FetchAll_wrap(SqlHandlePtr StatementHandle, py::list& rows, } else { fetchSize = 1000; } + fetchSize = static_cast( + std::min(fetchSize, std::max(numRowsInMemLimit, 1))); LOG("FetchAll_wrap: Fetching data in batch sizes of %d", fetchSize); size_t reservedBytes = 0; diff --git a/tests/test_004_cursor.py b/tests/test_004_cursor.py index 4441308b9..c6154fe32 100644 --- a/tests/test_004_cursor.py +++ b/tests/test_004_cursor.py @@ -1767,7 +1767,7 @@ def test_setinputsizes_rejects_excessive_column_size(): def test_executemany_rejects_excessive_cumulative_parameter_buffers(cursor): cursor.setinputsizes([(mssql_python.SQL_WVARCHAR, 50_000_000, 0)] * 3) with pytest.raises(RuntimeError, match="Parameter buffers exceed the 256 MiB allocation limit"): - cursor.executemany("SELECT ?, ?, ?", [("", "", "")]) + cursor.executemany("SELECT ?, ?, ?", [(None, None, None)]) def test_executemany_rejects_text_for_binary_parameter(cursor): @@ -1783,6 +1783,11 @@ def test_fetchmany_rejects_excessive_native_buffer(cursor): assert cursor.fetchone()[0] == "x" +def test_fetchall_clamps_wide_result_batch_to_native_buffer_budget(cursor): + cursor.execute("SELECT " + ", ".join("CAST(N'x' AS NVARCHAR(4000))" for _ in range(34))) + assert cursor.fetchall() == [("x",) * 34] + + def test_description(cursor): """Test description""" cursor.execute("SELECT * FROM #pytest_all_data_types WHERE id = 1") diff --git a/tests/test_004_cursor_arrow.py b/tests/test_004_cursor_arrow.py index b815e524b..b6e4bfb80 100644 --- a/tests/test_004_cursor_arrow.py +++ b/tests/test_004_cursor_arrow.py @@ -225,6 +225,21 @@ def test_arrow_empty_fetch(cursor: mssql_python.Cursor): _test_arrow_test_data(cursor, [col_data], fetch_length=0) +def test_arrow_wide_schema_only_and_single_row_fit_native_buffer_budget( + cursor: mssql_python.Cursor, +): + columns = ", ".join(f"CAST(N'x' AS NVARCHAR(4000)) AS col_{index}" for index in range(600)) + cursor.execute(f"SELECT {columns}") + + schema_batch = cursor.arrow_batch(0) + assert schema_batch.num_rows == 0 + assert schema_batch.num_columns == 600 + + data_batch = cursor.arrow_batch(1) + assert data_batch.num_rows == 1 + assert data_batch.num_columns == 600 + + def test_arrow_table_batchsize_negative(cursor: mssql_python.Cursor): cursor.execute("select 1 a") with pytest.raises(ValueError, match="batch_size"): diff --git a/tests/test_024_bulkcopy_arrow.py b/tests/test_024_bulkcopy_arrow.py index ee9858262..c8718284b 100644 --- a/tests/test_024_bulkcopy_arrow.py +++ b/tests/test_024_bulkcopy_arrow.py @@ -202,6 +202,10 @@ def test_batch_size_negative(self): with pytest.raises(ValueError, match="batch_size"): _bare_cursor().bulkcopy_arrow("t", pa.table({"a": [1]}), batch_size=-1) + def test_batch_size_exceeds_native_row_limit(self): + with pytest.raises(ValueError, match="batch_size"): + _bare_cursor().bulkcopy_arrow("t", pa.table({"a": [1]}), batch_size=1_000_001) + def test_batch_size_bool_rejected(self): for flag in (False, True): with pytest.raises(TypeError, match="batch_size"): From a0966897fbd55c77e2726671e30cbd6cb8484bfc Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Wed, 7 Oct 2026 11:33:04 +0530 Subject: [PATCH 06/39] FIX: Close remaining native size validation gaps Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/pybind/ddbc_bindings.cpp | 49 +++++++++++++++++++++------ mssql_python/pybind/param_detect.hpp | 14 +++++--- tests/test_023_execute_path_parity.py | 12 +++++++ 3 files changed, 61 insertions(+), 14 deletions(-) diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index 4b580df3a..fc4aa8658 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -349,6 +349,23 @@ size_t CheckedMultiplySize(size_t left, size_t right, const char* errorMessage) return left * right; } +size_t CheckedFetchColumnSize(SQLULEN columnSize) { + const size_t result = static_cast(columnSize); + if (static_cast(result) != columnSize) { + ThrowStdException("Column size is too large"); + } + return result; +} + +SQLLEN CheckedFetchBufferLength(size_t elementCount, size_t elementSize) { + const size_t byteCount = + CheckedMultiplySize(elementCount, elementSize, "Column fetch stride is too large"); + if (byteCount > static_cast(std::numeric_limits::max())) { + ThrowStdException("Column fetch stride is too large"); + } + return static_cast(byteCount); +} + constexpr int MAX_NATIVE_ROW_COUNT = 1000000; constexpr size_t MAX_NATIVE_FETCH_BYTES = 256ULL * 1024 * 1024; constexpr size_t MAX_NATIVE_PARAMETER_BYTES = 256ULL * 1024 * 1024; @@ -4449,23 +4466,30 @@ SQLRETURN SQLBindColums(SQLHSTMT hStmt, ColumnBuffers& buffers, const Metadata& case SQL_VARCHAR: case SQL_LONGVARCHAR: { HandleZeroColumnSizeAtFetch(columnSize); + const size_t baseColumnSize = CheckedFetchColumnSize(columnSize); if (useWideChar) { // Bind VARCHAR columns as SQL_C_WCHAR so the ODBC driver // returns UTF-16 data, avoiding code-page decode issues. - uint64_t fetchBufferSize = columnSize + 1 /*null-terminator*/; + const size_t fetchBufferSize = CheckedAddSize( + baseColumnSize, 1, "Column fetch stride is too large"); ResizeNativeFetchBuffer(buffers.wcharBuffers[col - 1], CheckedMultiplySize(fetchSize, fetchBufferSize, "Native fetch buffer is too large"), reservedBytes); ret = SQLBindCol_ptr( hStmt, col, SQL_C_WCHAR, buffers.wcharBuffers[col - 1].data(), - fetchBufferSize * sizeof(SQLWCHAR), buffers.indicators[col - 1].data()); + CheckedFetchBufferLength(fetchBufferSize, sizeof(SQLWCHAR)), + buffers.indicators[col - 1].data()); } else { // Original narrow-char path #if defined(__APPLE__) || defined(__linux__) - uint64_t fetchBufferSize = columnSize * 4 + 1 /*null-terminator*/; + const size_t fetchBufferSize = CheckedAddSize( + CheckedMultiplySize(baseColumnSize, 4, + "Column fetch stride is too large"), + 1, "Column fetch stride is too large"); #else - uint64_t fetchBufferSize = columnSize + 1 /*null-terminator*/; + const size_t fetchBufferSize = CheckedAddSize( + baseColumnSize, 1, "Column fetch stride is too large"); #endif ResizeNativeFetchBuffer(buffers.charBuffers[col - 1], CheckedMultiplySize(fetchSize, fetchBufferSize, @@ -4473,7 +4497,8 @@ SQLRETURN SQLBindColums(SQLHSTMT hStmt, ColumnBuffers& buffers, const Metadata& reservedBytes); ret = SQLBindCol_ptr( hStmt, col, SQL_C_CHAR, buffers.charBuffers[col - 1].data(), - fetchBufferSize * sizeof(SQLCHAR), buffers.indicators[col - 1].data()); + CheckedFetchBufferLength(fetchBufferSize, sizeof(SQLCHAR)), + buffers.indicators[col - 1].data()); } break; } @@ -4483,13 +4508,14 @@ SQLRETURN SQLBindColums(SQLHSTMT hStmt, ColumnBuffers& buffers, const Metadata& // TODO: handle variable length data correctly. This logic wont // suffice HandleZeroColumnSizeAtFetch(columnSize); - uint64_t fetchBufferSize = columnSize + 1 /*null-terminator*/; + const size_t fetchBufferSize = CheckedAddSize( + CheckedFetchColumnSize(columnSize), 1, "Column fetch stride is too large"); ResizeNativeFetchBuffer(buffers.wcharBuffers[col - 1], CheckedMultiplySize(fetchSize, fetchBufferSize, "Native fetch buffer is too large"), reservedBytes); ret = SQLBindCol_ptr(hStmt, col, SQL_C_WCHAR, buffers.wcharBuffers[col - 1].data(), - fetchBufferSize * sizeof(SQLWCHAR), + CheckedFetchBufferLength(fetchBufferSize, sizeof(SQLWCHAR)), buffers.indicators[col - 1].data()); break; } @@ -4572,17 +4598,20 @@ SQLRETURN SQLBindColums(SQLHSTMT hStmt, ColumnBuffers& buffers, const Metadata& case SQL_SS_UDT: case SQL_BINARY: case SQL_VARBINARY: - case SQL_LONGVARBINARY: + case SQL_LONGVARBINARY: { // TODO: handle variable length data correctly. This logic wont // suffice HandleZeroColumnSizeAtFetch(columnSize); + const size_t fetchBufferSize = CheckedFetchColumnSize(columnSize); ResizeNativeFetchBuffer(buffers.charBuffers[col - 1], - CheckedMultiplySize(fetchSize, columnSize, + CheckedMultiplySize(fetchSize, fetchBufferSize, "Native fetch buffer is too large"), reservedBytes); ret = SQLBindCol_ptr(hStmt, col, SQL_C_BINARY, buffers.charBuffers[col - 1].data(), - columnSize, buffers.indicators[col - 1].data()); + CheckedFetchBufferLength(fetchBufferSize, 1), + buffers.indicators[col - 1].data()); break; + } case SQL_SS_TIMESTAMPOFFSET: ResizeNativeFetchBuffer(buffers.datetimeoffsetBuffers[col - 1], fetchSize, reservedBytes); diff --git a/mssql_python/pybind/param_detect.hpp b/mssql_python/pybind/param_detect.hpp index 172a98b51..fe5cdeaa2 100644 --- a/mssql_python/pybind/param_detect.hpp +++ b/mssql_python/pybind/param_detect.hpp @@ -257,10 +257,16 @@ inline void ApplyInputSizeOverride(PyObject* params, PyObject* inputSize, Py_ssi } } - info.isDAE = - (PyUnicode_Check(obj) && PyLongGreaterThan(columnSize, MAX_INLINE_CHAR)) || - ((PyBytes_Check(obj) || PyByteArray_Check(obj)) && - PyLongGreaterThan(columnSize, MAX_INLINE_BINARY)); + const bool textNeedsDAE = + PyUnicode_Check(obj) && + (PyLongGreaterThan(columnSize, MAX_INLINE_CHAR) || + PyUnicode_GET_LENGTH(obj) > MAX_INLINE_CHAR); + const bool binaryNeedsDAE = + (PyBytes_Check(obj) || PyByteArray_Check(obj)) && + (PyLongGreaterThan(columnSize, MAX_INLINE_BINARY) || + (PyBytes_Check(obj) ? PyBytes_GET_SIZE(obj) : PyByteArray_GET_SIZE(obj)) > + MAX_INLINE_BINARY); + info.isDAE = textNeedsDAE || binaryNeedsDAE; if (PyTime_Check(obj) && info.paramCType == PARAM_C_TYPE_TEXT) { NormalizeTimeParam(params, index, info.columnSize); diff --git a/tests/test_023_execute_path_parity.py b/tests/test_023_execute_path_parity.py index d9bada323..9453cb85f 100644 --- a/tests/test_023_execute_path_parity.py +++ b/tests/test_023_execute_path_parity.py @@ -296,6 +296,18 @@ def test_setinputsizes_text_binding_normalizes_time(cursor): ) +@pytest.mark.parametrize( + ("value", "sql_type"), + [ + ("x" * 4001, ddbc_sql_const.SQL_VARCHAR.value), + (b"x" * 8001, ddbc_sql_const.SQL_VARBINARY.value), + ], +) +def test_setinputsizes_uses_actual_value_size_for_dae(cursor, value, sql_type): + """A too-small declared size cannot bypass streaming for a large value.""" + assert _override_roundtrip(cursor, value, sql_type, 1) == value + + @pytest.mark.parametrize("sql_type", [None, ddbc_sql_const.SQL_VARCHAR.value]) def test_time_isoformat_must_return_string(cursor, sql_type): """Native time normalization rejects a broken subclass contract on either path.""" From 231ae90a543c4ef2637990ed1d27568d942537d3 Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Wed, 7 Oct 2026 11:43:24 +0530 Subject: [PATCH 07/39] FIX: Measure override payloads in native units Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/pybind/ddbc_bindings.cpp | 3 +- mssql_python/pybind/param_detect.hpp | 52 +++++++++++++++++++-------- tests/test_023_execute_path_parity.py | 1 + 3 files changed, 41 insertions(+), 15 deletions(-) diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index fc4aa8658..d19b4013a 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -2290,7 +2290,8 @@ SQLRETURN SQLExecute_wrap(const SqlHandlePtr statementHandle, // Run DetectParamTypes BEFORE SQLPrepare so that type-detection errors // (unsupported type, NaN Decimal, precision overflow) don't leave the // cursor in a half-prepared state. - std::vector paramInfos = DetectParamTypes(params.ptr(), input_sizes.ptr()); + std::vector paramInfos = + DetectParamTypes(params.ptr(), input_sizes.ptr(), charEncoding); RETCODE rc; bool already_prepared = is_stmt_prepared[0].cast(); diff --git a/mssql_python/pybind/param_detect.hpp b/mssql_python/pybind/param_detect.hpp index fe5cdeaa2..9d28c63a5 100644 --- a/mssql_python/pybind/param_detect.hpp +++ b/mssql_python/pybind/param_detect.hpp @@ -190,6 +190,29 @@ inline bool PyLongGreaterThan(PyObject* value, long long threshold) { return overflow > 0 || (overflow == 0 && result > threshold); } +inline Py_ssize_t UnicodeUtf16Length(PyObject* value) { + const Py_ssize_t length = PyUnicode_GET_LENGTH(value); + if (PyUnicode_KIND(value) <= PyUnicode_2BYTE_KIND) { + return length; + } + + Py_ssize_t utf16Length = 0; + const Py_UCS4* data = PyUnicode_4BYTE_DATA(value); + for (Py_ssize_t index = 0; index < length; ++index) { + utf16Length += data[index] > 0xFFFF ? 2 : 1; + } + return utf16Length; +} + +inline Py_ssize_t EncodedUnicodeLength(PyObject* value, const std::string& encoding) { + py::object encoded = + steal(PyUnicode_AsEncodedString(value, encoding.c_str(), "strict")); + if (!encoded) { + throw py::error_already_set(); + } + return PyBytes_GET_SIZE(encoded.ptr()); +} + inline PyObject* FormatDecimalParam(PyObject* params, Py_ssize_t index, PyObject* value) { py::object formatted = steal(PyObject_CallMethod(value, "__format__", "s", "f")); if (!formatted) throw py::error_already_set(); @@ -216,7 +239,7 @@ inline void NormalizeTimeParam(PyObject* params, Py_ssize_t index, SQLULEN& colu } inline void ApplyInputSizeOverride(PyObject* params, PyObject* inputSize, Py_ssize_t index, - ParamInfo& info) { + ParamInfo& info, const std::string& charEncoding) { py::tuple values = borrow(inputSize); info.paramSQLType = values[0].cast(); info.paramCType = values[1].cast(); @@ -257,16 +280,25 @@ inline void ApplyInputSizeOverride(PyObject* params, PyObject* inputSize, Py_ssi } } + Py_ssize_t actualTextLength = 0; + if (PyUnicode_Check(obj)) { + actualTextLength = info.paramCType == SQL_C_CHAR + ? EncodedUnicodeLength(obj, charEncoding) + : UnicodeUtf16Length(obj); + } const bool textNeedsDAE = PyUnicode_Check(obj) && (PyLongGreaterThan(columnSize, MAX_INLINE_CHAR) || - PyUnicode_GET_LENGTH(obj) > MAX_INLINE_CHAR); + actualTextLength > MAX_INLINE_CHAR); const bool binaryNeedsDAE = (PyBytes_Check(obj) || PyByteArray_Check(obj)) && (PyLongGreaterThan(columnSize, MAX_INLINE_BINARY) || (PyBytes_Check(obj) ? PyBytes_GET_SIZE(obj) : PyByteArray_GET_SIZE(obj)) > MAX_INLINE_BINARY); info.isDAE = textNeedsDAE || binaryNeedsDAE; + if (info.isDAE && !PyLongGreaterThan(columnSize, MAX_INLINE_BINARY)) { + info.columnSize = 0; + } if (PyTime_Check(obj) && info.paramCType == PARAM_C_TYPE_TEXT) { NormalizeTimeParam(params, index, info.columnSize); @@ -302,7 +334,8 @@ inline void ApplyInputSizeOverride(PyObject* params, PyObject* inputSize, Py_ssi // // Takes raw PyObject* lists. Caller guarantees params is a fresh copy (cursor.py // does list(actual_params)), so in-place mutation via PyList_SetItem is safe. -inline std::vector DetectParamTypes(PyObject* params, PyObject* inputSizes) { +inline std::vector DetectParamTypes(PyObject* params, PyObject* inputSizes, + const std::string& charEncoding = "utf-8") { PyTypeCache::initialize(); const Py_ssize_t n = PyList_GET_SIZE(params); @@ -318,7 +351,7 @@ inline std::vector DetectParamTypes(PyObject* params, PyObject* input info.isDAE = false; if (i < inputSizeCount) { - ApplyInputSizeOverride(params, PyList_GET_ITEM(inputSizes, i), i, info); + ApplyInputSizeOverride(params, PyList_GET_ITEM(inputSizes, i), i, info, charEncoding); continue; } @@ -404,16 +437,7 @@ inline std::vector DetectParamTypes(PyObject* params, PyObject* input unsigned int kind = PyUnicode_KIND(obj); const void* udata = PyUnicode_DATA(obj); - Py_ssize_t utf16_len; - if (kind <= PyUnicode_2BYTE_KIND) { - utf16_len = length; - } else { - utf16_len = 0; - const Py_UCS4* data = PyUnicode_4BYTE_DATA(obj); - for (Py_ssize_t j = 0; j < length; ++j) { - utf16_len += (data[j] > 0xFFFF) ? 2 : 1; - } - } + const Py_ssize_t utf16_len = UnicodeUtf16Length(obj); // Detect whether the string needs wide-char (NVARCHAR) or narrow (VARCHAR) binding. // PyUnicode_IS_COMPACT_ASCII is a struct field check (O(1)), not a content scan. diff --git a/tests/test_023_execute_path_parity.py b/tests/test_023_execute_path_parity.py index 9453cb85f..cc39db777 100644 --- a/tests/test_023_execute_path_parity.py +++ b/tests/test_023_execute_path_parity.py @@ -300,6 +300,7 @@ def test_setinputsizes_text_binding_normalizes_time(cursor): ("value", "sql_type"), [ ("x" * 4001, ddbc_sql_const.SQL_VARCHAR.value), + ("\U0001f600" * 3000, ddbc_sql_const.SQL_WVARCHAR.value), (b"x" * 8001, ddbc_sql_const.SQL_VARBINARY.value), ], ) From 444fb7f17c5eff596598592d166727abc0428ea2 Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Wed, 7 Oct 2026 11:53:14 +0530 Subject: [PATCH 08/39] FIX: Enforce peak native fetch memory limit Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/pybind/ddbc_bindings.cpp | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index 48963d516..a634144a7 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -474,8 +474,19 @@ void ReserveNativeFetchBytes(size_t& reservedBytes, size_t count, size_t element template void ResizeNativeFetchBuffer(std::vector& buffer, size_t count, size_t& reservedBytes) { + if (buffer.size() >= count) { + return; + } + + const size_t oldBytes = CheckedMultiplySize( + buffer.size(), sizeof(ElementType), "Native fetch buffer size is too large"); ReserveNativeFetchBytes(reservedBytes, count, sizeof(ElementType)); - buffer.resize(count); + { + std::vector resizedBuffer(count); + std::copy(buffer.begin(), buffer.end(), resizedBuffer.begin()); + buffer.swap(resizedBuffer); + } + reservedBytes -= oldBytes; } template @@ -488,8 +499,7 @@ void EnsureNativeFetchBufferSize(std::vector& buffer, size_t requir while (newSize < requiredSize) { newSize = CheckedMultiplySize(newSize, 2, "Native fetch buffer size is too large"); } - ReserveNativeFetchBytes(reservedBytes, newSize - buffer.size(), sizeof(ElementType)); - buffer.resize(newSize); + ResizeNativeFetchBuffer(buffer, newSize, reservedBytes); } template @@ -5442,8 +5452,7 @@ SQLRETURN GetDataVar(SQLHSTMT hStmt, SQLUSMALLINT colNumber, SQLSMALLINT cType, start = dataVec.size() - sizeNullTerminator; // Resize buffer for next iteration - ReserveNativeFetchBytes(reservedBytes, end - dataVec.size(), sizeof(T)); - dataVec.resize(end); + ResizeNativeFetchBuffer(dataVec, end, reservedBytes); } else { // Unexpected return code return ret; From d5b44313d98889cc062fca73f155c083edb45bc0 Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Wed, 7 Oct 2026 11:59:37 +0530 Subject: [PATCH 09/39] FIX: Bound streamed LOB native memory Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/pybind/ddbc_bindings.cpp | 9 ++++++++- mssql_python/pybind/param_detect.hpp | 2 +- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index a634144a7..2b995328b 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -3471,6 +3471,8 @@ py::object FetchLobColumnData(SQLHSTMT hStmt, SQLUSMALLINT colIndex, SQLSMALLINT py::handle messages) { PERF_TIMER("FetchLobColumnData"); std::vector buffer; + size_t reservedBytes = 0; + ReserveNativeFetchBytes(reservedBytes, DAE_CHUNK_SIZE, sizeof(char)); SQLRETURN ret = SQL_SUCCESS_WITH_INFO; int loopCount = 0; @@ -3540,7 +3542,11 @@ py::object FetchLobColumnData(SQLHSTMT hStmt, SQLUSMALLINT colIndex, SQLSMALLINT } } if (bytesRead > 0) { - buffer.insert(buffer.end(), chunk.begin(), chunk.begin() + bytesRead); + const size_t previousSize = buffer.size(); + const size_t requiredSize = + CheckedAddSize(previousSize, bytesRead, "LOB fetch buffer is too large"); + ResizeNativeFetchBuffer(buffer, requiredSize, reservedBytes); + std::memcpy(buffer.data() + previousSize, chunk.data(), bytesRead); LOG("FetchLobColumnData: Appended %zu bytes at loop %d", bytesRead, loopCount); } if (ret == SQL_SUCCESS) { @@ -3561,6 +3567,7 @@ py::object FetchLobColumnData(SQLHSTMT hStmt, SQLUSMALLINT colIndex, SQLSMALLINT ThrowStdException("Wide-character LOB data has an invalid byte length"); } size_t wcharCount = buffer.size() / sizeof(SQLWCHAR); + ReserveNativeFetchBytes(reservedBytes, wcharCount, sizeof(SQLWCHAR)); std::vector alignedBuf(wcharCount); std::memcpy(alignedBuf.data(), buffer.data(), buffer.size()); return py::cast(dupeSqlWCharAsUtf16Le(alignedBuf.data(), wcharCount)); diff --git a/mssql_python/pybind/param_detect.hpp b/mssql_python/pybind/param_detect.hpp index 9d28c63a5..ae0cf7bf8 100644 --- a/mssql_python/pybind/param_detect.hpp +++ b/mssql_python/pybind/param_detect.hpp @@ -296,7 +296,7 @@ inline void ApplyInputSizeOverride(PyObject* params, PyObject* inputSize, Py_ssi (PyBytes_Check(obj) ? PyBytes_GET_SIZE(obj) : PyByteArray_GET_SIZE(obj)) > MAX_INLINE_BINARY); info.isDAE = textNeedsDAE || binaryNeedsDAE; - if (info.isDAE && !PyLongGreaterThan(columnSize, MAX_INLINE_BINARY)) { + if (info.isDAE) { info.columnSize = 0; } From 3b60b8cd37389caf0b5ee6635cbd9a2bb795c489 Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Wed, 7 Oct 2026 12:08:14 +0530 Subject: [PATCH 10/39] FIX: Validate executemany payload sizes Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/cursor.py | 25 ++++++--- mssql_python/pybind/ddbc_bindings.cpp | 77 ++++++++++++++++++++------- tests/test_023_execute_path_parity.py | 20 +++++++ 3 files changed, 97 insertions(+), 25 deletions(-) diff --git a/mssql_python/cursor.py b/mssql_python/cursor.py index 7ebc0c35b..39c2f6267 100644 --- a/mssql_python/cursor.py +++ b/mssql_python/cursor.py @@ -2580,6 +2580,7 @@ def executemany( # pylint: disable=too-many-locals,too-many-branches,too-many-s param_info = ddbc_bindings.ParamInfo parameters_type = [] any_dae = False + encoding_settings = self._get_encoding_settings() # Check if we have explicit input sizes set if self._inputsizes: @@ -2615,10 +2616,22 @@ def executemany( # pylint: disable=too-many-locals,too-many-branches,too-many-s ): c_type = ddbc_sql_const.SQL_C_CHAR.value - # Check if this should be a DAE (data at execution) parameter based on column size + # Check if this should be a DAE (data at execution) parameter. if sample_value is not None: - if isinstance(sample_value, str) and column_size > MAX_INLINE_CHAR: - is_dae = True + if isinstance(sample_value, str): + if c_type == ddbc_sql_const.SQL_C_CHAR.value: + text_size = max( + len(value.encode(encoding_settings["encoding"])) + for value in column + if isinstance(value, str) + ) + else: + text_size = max( + sum(2 if ord(char) > 0xFFFF else 1 for char in value) + for value in column + if isinstance(value, str) + ) + is_dae = column_size > MAX_INLINE_CHAR or text_size > MAX_INLINE_CHAR elif isinstance(sample_value, (bytes, bytearray)) and column_size > 8000: is_dae = True @@ -2652,6 +2665,9 @@ def executemany( # pylint: disable=too-many-locals,too-many-branches,too-many-s # Always ensure at least a minimum size of 1 for empty strings column_size = max(max_binary_size, 1) + if is_dae: + column_size = 0 + paraminfo = param_info() paraminfo.paramCType = c_type paraminfo.paramSQLType = sql_type @@ -2816,9 +2832,6 @@ def executemany( # pylint: disable=too-many-locals,too-many-branches,too-many-s processed_parameters ) - # Get encoding settings - encoding_settings = self._get_encoding_settings() - # Debug logging: emit batch metadata only. Never log parameter values or # row representations here -- rows may contain PII (SSNs, emails, # balances) that would leak into log files and APM/log shippers even diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index 2b995328b..d2a1f9d02 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -2569,16 +2569,41 @@ SQLRETURN BindParameterArray(SqlHandle& handle, SQLHANDLE hStmt, const py::list& std::memset(wcharArray + i * (info.columnSize + 1), 0, (info.columnSize + 1) * sizeof(SQLWCHAR)); } else { - std::u16string wstr = columnValues[i].cast(); - // u16string is already UTF-16, so the - // original check is sufficient - if (wstr.length() > info.columnSize) { + if (!py::isinstance(columnValues[i])) { + ThrowStdException( + MakeParamMismatchErrorStr(info.paramCType, paramIndex)); + } + if (static_cast( + PyUnicode_GET_LENGTH(columnValues[i].ptr())) > + info.columnSize) { ThrowStdException("Input string exceeds allowed column size " "at parameter index " + std::to_string(paramIndex)); } - std::memcpy(wcharArray + i * (info.columnSize + 1), wstr.c_str(), - (wstr.length() + 1) * sizeof(SQLWCHAR)); + py::object encoded = + columnValues[i].attr("encode")("utf-16-le", "strict"); + char* encodedData = nullptr; + Py_ssize_t encodedSize = 0; + if (PyBytes_AsStringAndSize(encoded.ptr(), &encodedData, + &encodedSize) != 0) { + throw py::error_already_set(); + } + if (encodedSize < 0 || + static_cast(encodedSize) % sizeof(SQLWCHAR) != 0) { + ThrowStdException( + "Wide-character parameter has an invalid byte length"); + } + const size_t wcharLength = + static_cast(encodedSize) / sizeof(SQLWCHAR); + if (wcharLength > info.columnSize) { + ThrowStdException("Input string exceeds allowed column size " + "at parameter index " + + std::to_string(paramIndex)); + } + SQLWCHAR* destination = wcharArray + i * elementWidth; + std::copy_n(reinterpret_cast(encodedData), + wcharLength, destination); + destination[wcharLength] = 0; strLenOrIndArray[i] = SQL_NTS; } } @@ -2679,19 +2704,24 @@ SQLRETURN BindParameterArray(SqlHandle& handle, SQLHANDLE hStmt, const py::list& ThrowStdException(MakeParamMismatchErrorStr(info.paramCType, paramIndex)); } - std::string encodedStr; - + py::object encoded; + char* encodedData = nullptr; + Py_ssize_t encodedSize = 0; if (py::isinstance(columnValues[i])) { // Use Python's codec system to encode the string with specified // encoding try { - py::object encoded = + encoded = columnValues[i].attr("encode")(charEncoding, "strict"); - encodedStr = encoded.cast(); + if (PyBytes_AsStringAndSize(encoded.ptr(), &encodedData, + &encodedSize) != 0) { + throw py::error_already_set(); + } LOG("BindParameterArray: param[%d] row[%zu] SQL_C_CHAR - " "Encoded with '%s', " "size=%zu bytes", - paramIndex, i, charEncoding.c_str(), encodedStr.size()); + paramIndex, i, charEncoding.c_str(), + static_cast(encodedSize)); } catch (const py::error_already_set& e) { LOG_ERROR("BindParameterArray: param[%d] row[%zu] SQL_C_CHAR - " "Failed to encode " @@ -2702,22 +2732,31 @@ SQLRETURN BindParameterArray(SqlHandle& handle, SQLHANDLE hStmt, const py::list& std::to_string(paramIndex) + " row " + std::to_string(i) + " with encoding '" + charEncoding + "': " + e.what()); } + } else if (py::isinstance(columnValues[i])) { + encoded = py::reinterpret_borrow(columnValues[i]); + if (PyBytes_AsStringAndSize(encoded.ptr(), &encodedData, + &encodedSize) != 0) { + throw py::error_already_set(); + } } else { - // bytes/bytearray - use as-is (already encoded) - encodedStr = columnValues[i].cast(); + encodedData = PyByteArray_AsString(columnValues[i].ptr()); + encodedSize = PyByteArray_Size(columnValues[i].ptr()); + if (encodedData == nullptr || encodedSize < 0) { + throw py::error_already_set(); + } } - if (encodedStr.size() > info.columnSize) { + const size_t dataSize = static_cast(encodedSize); + if (dataSize > info.columnSize) { LOG("BindParameterArray: String/binary too " "long - param_index=%d, row=%zu, size=%zu, " "max=%zu", - paramIndex, i, encodedStr.size(), info.columnSize); + paramIndex, i, dataSize, info.columnSize); ThrowStdException("Input exceeds column size at index " + std::to_string(i)); } - std::memcpy(charArray + i * (info.columnSize + 1), encodedStr.c_str(), - encodedStr.size()); - strLenOrIndArray[i] = static_cast(encodedStr.size()); + std::copy_n(encodedData, dataSize, charArray + i * elementWidth); + strLenOrIndArray[i] = static_cast(dataSize); } } LOG("BindParameterArray: SQL_C_CHAR/BINARY bound - " @@ -3546,7 +3585,7 @@ py::object FetchLobColumnData(SQLHSTMT hStmt, SQLUSMALLINT colIndex, SQLSMALLINT const size_t requiredSize = CheckedAddSize(previousSize, bytesRead, "LOB fetch buffer is too large"); ResizeNativeFetchBuffer(buffer, requiredSize, reservedBytes); - std::memcpy(buffer.data() + previousSize, chunk.data(), bytesRead); + std::copy_n(chunk.data(), bytesRead, buffer.data() + previousSize); LOG("FetchLobColumnData: Appended %zu bytes at loop %d", bytesRead, loopCount); } if (ret == SQL_SUCCESS) { diff --git a/tests/test_023_execute_path_parity.py b/tests/test_023_execute_path_parity.py index cc39db777..595ae9adc 100644 --- a/tests/test_023_execute_path_parity.py +++ b/tests/test_023_execute_path_parity.py @@ -309,6 +309,26 @@ def test_setinputsizes_uses_actual_value_size_for_dae(cursor, value, sql_type): assert _override_roundtrip(cursor, value, sql_type, 1) == value +@pytest.mark.parametrize( + ("value", "sql_type"), + [ + pytest.param("x" * 4001, ddbc_sql_const.SQL_VARCHAR.value, id="varchar-encoded-bytes"), + pytest.param( + "\U0001f600" * 2001, + ddbc_sql_const.SQL_WVARCHAR.value, + id="nvarchar-utf16-units", + ), + ], +) +def test_executemany_setinputsizes_uses_actual_text_size_for_dae(cursor, value, sql_type): + """A too-small declared size cannot bypass streaming for an array-bound text value.""" + cursor.execute("CREATE TABLE #dae_batch (value NVARCHAR(MAX))") + cursor.setinputsizes([(sql_type, 1, 0)]) + cursor.executemany("INSERT INTO #dae_batch (value) VALUES (?)", [(value,), ("short",)]) + cursor.execute("SELECT value FROM #dae_batch ORDER BY LEN(value) DESC") + assert cursor.fetchall() == [(value,), ("short",)] + + @pytest.mark.parametrize("sql_type", [None, ddbc_sql_const.SQL_VARCHAR.value]) def test_time_isoformat_must_return_string(cursor, sql_type): """Native time normalization rejects a broken subclass contract on either path.""" From 9f5c11af1563feda8f6724fa0f375936073cb421 Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Wed, 7 Oct 2026 12:18:06 +0530 Subject: [PATCH 11/39] FIX: Copy encoded wide text as bytes Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/pybind/ddbc_bindings.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index d2a1f9d02..422e6fc8b 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -2601,8 +2601,8 @@ SQLRETURN BindParameterArray(SqlHandle& handle, SQLHANDLE hStmt, const py::list& std::to_string(paramIndex)); } SQLWCHAR* destination = wcharArray + i * elementWidth; - std::copy_n(reinterpret_cast(encodedData), - wcharLength, destination); + std::copy_n(encodedData, static_cast(encodedSize), + reinterpret_cast(destination)); destination[wcharLength] = 0; strLenOrIndArray[i] = SQL_NTS; } From 5b26c36bf155a64c23e7287d6559697a5b64522f Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Wed, 7 Oct 2026 12:31:31 +0530 Subject: [PATCH 12/39] FIX: Stream DAE payloads in bounded chunks Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/cursor.py | 29 +++- mssql_python/pybind/ddbc_bindings.cpp | 196 +++++++++++++------------- mssql_python/pybind/param_detect.hpp | 26 +++- 3 files changed, 144 insertions(+), 107 deletions(-) diff --git a/mssql_python/cursor.py b/mssql_python/cursor.py index 39c2f6267..f80fdbeb4 100644 --- a/mssql_python/cursor.py +++ b/mssql_python/cursor.py @@ -14,6 +14,7 @@ import decimal import logging import uuid +import codecs import datetime import warnings from typing import List, Mapping, Union, Any, Optional, Tuple, Sequence, TYPE_CHECKING, Iterable @@ -54,6 +55,22 @@ # Bound each native fetch allocation; Arrow initially reserves 42 bytes per variable-width row. MAX_NATIVE_ROW_COUNT: int = 1_000_000 MAX_NATIVE_PARAMETER_SIZE: int = 256 * 1024 * 1024 + + +def _encoded_length_exceeds(value: str, encoding: str, limit: int) -> bool: + encoder = codecs.getincrementalencoder(encoding)(errors="strict") + total = 0 + chunk_size = 4096 + for offset in range(0, len(value), chunk_size): + end = min(offset + chunk_size, len(value)) + total += len(encoder.encode(value[offset:end], final=end == len(value))) + if total > limit: + return True + if not value: + total += len(encoder.encode("", final=True)) + return total > limit + + # SQL BIGINT is a signed 64-bit integer. Ints outside this range have no BIGINT # encoding and must be rejected at detect time on both paths (see _map_sql_type). BIGINT_MIN: int = -(2**63) @@ -2620,8 +2637,12 @@ def executemany( # pylint: disable=too-many-locals,too-many-branches,too-many-s if sample_value is not None: if isinstance(sample_value, str): if c_type == ddbc_sql_const.SQL_C_CHAR.value: - text_size = max( - len(value.encode(encoding_settings["encoding"])) + is_dae = column_size > MAX_INLINE_CHAR or any( + _encoded_length_exceeds( + value, + encoding_settings["encoding"], + MAX_INLINE_CHAR, + ) for value in column if isinstance(value, str) ) @@ -2631,7 +2652,9 @@ def executemany( # pylint: disable=too-many-locals,too-many-branches,too-many-s for value in column if isinstance(value, str) ) - is_dae = column_size > MAX_INLINE_CHAR or text_size > MAX_INLINE_CHAR + is_dae = ( + column_size > MAX_INLINE_CHAR or text_size > MAX_INLINE_CHAR + ) elif isinstance(sample_value, (bytes, bytearray)) and column_size > 8000: is_dae = True diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index 422e6fc8b..66fea9790 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -535,6 +535,78 @@ static SQLRETURN stream_dae_chunks(const void* data, size_t total_bytes, PutData return SQL_SUCCESS; } +template +static SQLRETURN stream_unicode_dae_chunks(PyObject* value, const std::string& encoding, + PutDataFn put_data_fn) { + py::object encoderFactory = + py::module_::import("codecs").attr("getincrementalencoder")(encoding); + py::object encoder = encoderFactory("strict"); + const Py_ssize_t length = PyUnicode_GET_LENGTH(value); + constexpr Py_ssize_t chunkSize = 4096; + for (Py_ssize_t offset = 0; offset < length; offset += chunkSize) { + const Py_ssize_t end = std::min(offset + chunkSize, length); + py::object chunk = steal(PyUnicode_Substring(value, offset, end)); + if (!chunk) throw py::error_already_set(); + py::object encoded = encoder.attr("encode")(chunk, end == length); + char* data = nullptr; + Py_ssize_t size = 0; + if (PyBytes_AsStringAndSize(encoded.ptr(), &data, &size) != 0) { + throw py::error_already_set(); + } + if (size > 0) { + SQLRETURN rc = put_data_fn(data, static_cast(size)); + if (!SQL_SUCCEEDED(rc)) return rc; + } + } + if (length == 0) { + py::object encoded = encoder.attr("encode")(py::str(), true); + char* data = nullptr; + Py_ssize_t size = 0; + if (PyBytes_AsStringAndSize(encoded.ptr(), &data, &size) != 0) { + throw py::error_already_set(); + } + if (size > 0) return put_data_fn(data, static_cast(size)); + } + return SQL_SUCCESS; +} + +static SQLRETURN StreamDAEParameter(SQLHSTMT hStmt, const ParamInfo& info, + const std::string& charEncoding) { + PyObject* value = info.dataPtr.ptr(); + if (!value || value == Py_None) { + py::gil_scoped_release release; + return SQLPutData_ptr(hStmt, nullptr, 0); + } + + auto putImmutableData = [&](SQLPOINTER data, SQLLEN length) { + py::gil_scoped_release release; + return SQLPutData_ptr(hStmt, data, length); + }; + if (PyUnicode_Check(value)) { + if (info.paramCType == SQL_C_WCHAR) { + return stream_unicode_dae_chunks(value, "utf-16-le", putImmutableData); + } + if (info.paramCType == SQL_C_CHAR) { + return stream_unicode_dae_chunks(value, charEncoding, putImmutableData); + } + ThrowStdException("DAE only supports text C types for str values"); + } + if (PyBytes_Check(value)) { + return stream_dae_chunks(PyBytes_AS_STRING(value), + static_cast(PyBytes_GET_SIZE(value)), + putImmutableData); + } + if (PyByteArray_Check(value)) { + auto putMutableData = [&](SQLPOINTER data, SQLLEN length) { + return SQLPutData_ptr(hStmt, data, length); + }; + return stream_dae_chunks(PyByteArray_AS_STRING(value), + static_cast(PyByteArray_GET_SIZE(value)), + putMutableData); + } + ThrowStdException("DAE only supports str, bytes, or bytearray values"); +} + // GH-610: Resolve SQL type for a NULL parameter using per-handle cache. // On cache miss, calls SQLDescribeParam and stores the result. static DescribedParamInfo ResolveNullParamType(SqlHandle& handle, SQLHANDLE hStmt, int paramIndex) { @@ -2377,10 +2449,6 @@ SQLRETURN SQLExecute_wrap(const SqlHandlePtr statementHandle, // GIL is released around each ODBC call to match slow-path concurrency. if (rc == SQL_NEED_DATA) { SQLPOINTER paramToken = nullptr; - auto putData = [&](SQLPOINTER data, SQLLEN len) { - py::gil_scoped_release release; - return SQLPutData_ptr(hStmt, data, len); - }; while (true) { { py::gil_scoped_release release; @@ -2399,58 +2467,8 @@ SQLRETURN SQLExecute_wrap(const SqlHandlePtr statementHandle, if (matchedInfo < first || matchedInfo >= last) { ThrowStdException("SQLExecute: unrecognized paramToken from SQLParamData"); } - PyObject* pyObj = matchedInfo->dataPtr.ptr(); - if (!pyObj || pyObj == Py_None) { - py::gil_scoped_release release; - SQLPutData_ptr(hStmt, nullptr, 0); - continue; - } - - if (PyUnicode_Check(pyObj)) { - if (matchedInfo->paramCType == SQL_C_WCHAR) { - std::u16string u16 = - borrow(pyObj).cast(); - rc = stream_dae_chunks( - reinterpretU16stringAsSqlWChar(u16), - u16.size() * sizeof(SQLWCHAR), - putData); - if (!SQL_SUCCEEDED(rc)) return rc; - } else if (matchedInfo->paramCType == SQL_C_CHAR) { - std::string encodedStr; - py::object encoded = borrow(pyObj) - .attr("encode")(charEncoding, "strict"); - encodedStr = encoded.cast(); - rc = stream_dae_chunks(encodedStr.data(), encodedStr.size(), putData); - if (!SQL_SUCCEEDED(rc)) return rc; - } else { - ThrowStdException("SQLExecute: unsupported C type for str in DAE"); - } - } else if (PyBytes_Check(pyObj) || PyByteArray_Check(pyObj)) { - // matchedInfo->dataPtr holds a strong ref to pyObj for the whole loop. - const char* dataPtr = nullptr; - size_t totalBytes = 0; - std::string bytesStorage; // only used for the bytearray copy below - - if (PyBytes_Check(pyObj)) { - // bytes is immutable and kept alive by the strong ref above, so stream - // straight from its internal buffer with no copy. This is the large-blob - // DAE path, so skipping a full payload copy is the whole point. - dataPtr = PyBytes_AS_STRING(pyObj); - totalBytes = static_cast(PyBytes_GET_SIZE(pyObj)); - } else { - // bytearray is mutable and the GIL is released mid-stream, so copy to a - // stable buffer before streaming. - bytesStorage.assign(PyByteArray_AS_STRING(pyObj), - static_cast(PyByteArray_GET_SIZE(pyObj))); - dataPtr = bytesStorage.data(); - totalBytes = bytesStorage.size(); - } - - rc = stream_dae_chunks(dataPtr, totalBytes, putData); - if (!SQL_SUCCEEDED(rc)) return rc; - } else { - ThrowStdException("SQLExecute: DAE only supported for str or bytes"); - } + rc = StreamDAEParameter(hStmt, *matchedInfo, charEncoding); + if (!SQL_SUCCEEDED(rc)) return rc; } if (!SQL_SUCCEEDED(rc) && rc != SQL_NO_DATA) return rc; } @@ -3245,11 +3263,23 @@ SQLRETURN SQLExecuteMany_wrap(const SqlHandlePtr statementHandle, const std::u16 return rc; } else { LOG("SQLExecuteMany: Using DAE (data-at-execution) - row_count=%zu", - columnwise_params.size()); - size_t rowCount = columnwise_params.size(); + paramSetSize); + if (columnwise_params.size() != paramInfos.size()) { + ThrowStdException("Parameter count does not match parameter metadata count"); + } + rc = SQLSetStmtAttr_ptr(hStmt, SQL_ATTR_PARAMSET_SIZE, (SQLPOINTER)1, 0); + if (!SQL_SUCCEEDED(rc)) return rc; + size_t rowCount = paramSetSize; for (size_t rowIndex = 0; rowIndex < rowCount; ++rowIndex) { LOG("SQLExecuteMany: Processing DAE row %zu of %zu", rowIndex + 1, rowCount); - py::list rowParams = columnwise_params[rowIndex]; + py::list rowParams; + for (py::handle column : columnwise_params) { + py::sequence values = py::reinterpret_borrow(column); + if (static_cast(py::len(values)) != paramSetSize) { + ThrowStdException("Parameter column length does not match parameter set size"); + } + rowParams.append(values[rowIndex]); + } std::vector> paramBuffers; rc = BindParameters(*statementHandle, hStmt, rowParams, paramInfos, @@ -3284,47 +3314,15 @@ SQLRETURN SQLExecuteMany_wrap(const SqlHandlePtr statementHandle, const std::u16 return rc; } - py::object* py_obj_ptr = reinterpret_cast(token); - if (!py_obj_ptr) { - LOG("SQLExecuteMany: NULL token pointer in DAE - chunk=%zu", dae_chunk_count); - return SQL_ERROR; - } - - if (py::isinstance(*py_obj_ptr)) { - std::string data = py_obj_ptr->cast(); - SQLLEN data_len = static_cast(data.size()); - LOG("SQLExecuteMany: Sending string DAE data - chunk=%zu, " - "length=%lld", - dae_chunk_count, static_cast(data_len)); - rc = [&] { - py::gil_scoped_release release; - return SQLPutData_ptr(hStmt, (SQLPOINTER)data.c_str(), data_len); - }(); - if (!SQL_SUCCEEDED(rc) && rc != SQL_NEED_DATA) { - LOG("SQLExecuteMany: SQLPutData(string) failed - " - "chunk=%zu, rc=%d", - dae_chunk_count, rc); - } - } else if (py::isinstance(*py_obj_ptr) || - py::isinstance(*py_obj_ptr)) { - std::string data = py_obj_ptr->cast(); - SQLLEN data_len = static_cast(data.size()); - LOG("SQLExecuteMany: Sending bytes/bytearray DAE data - " - "chunk=%zu, length=%lld", - dae_chunk_count, static_cast(data_len)); - rc = [&] { - py::gil_scoped_release release; - return SQLPutData_ptr(hStmt, (SQLPOINTER)data.c_str(), data_len); - }(); - if (!SQL_SUCCEEDED(rc) && rc != SQL_NEED_DATA) { - LOG("SQLExecuteMany: SQLPutData(bytes) failed - " - "chunk=%zu, rc=%d", - dae_chunk_count, rc); - } - } else { - LOG("SQLExecuteMany: Unsupported DAE data type - chunk=%zu", dae_chunk_count); - return SQL_ERROR; + const ParamInfo* matchedInfo = reinterpret_cast(token); + const ParamInfo* first = paramInfos.data(); + const ParamInfo* last = first + paramInfos.size(); + if (matchedInfo < first || matchedInfo >= last) { + ThrowStdException( + "SQLExecuteMany: unrecognized paramToken from SQLParamData"); } + rc = StreamDAEParameter(hStmt, *matchedInfo, charEncoding); + if (!SQL_SUCCEEDED(rc)) return rc; dae_chunk_count++; } LOG("SQLExecuteMany: DAE completed for row %zu - total_chunks=%zu, " diff --git a/mssql_python/pybind/param_detect.hpp b/mssql_python/pybind/param_detect.hpp index ae0cf7bf8..f8a3e7d4f 100644 --- a/mssql_python/pybind/param_detect.hpp +++ b/mssql_python/pybind/param_detect.hpp @@ -205,12 +205,28 @@ inline Py_ssize_t UnicodeUtf16Length(PyObject* value) { } inline Py_ssize_t EncodedUnicodeLength(PyObject* value, const std::string& encoding) { - py::object encoded = - steal(PyUnicode_AsEncodedString(value, encoding.c_str(), "strict")); - if (!encoded) { - throw py::error_already_set(); + py::object encoderFactory = + py::module_::import("codecs").attr("getincrementalencoder")(encoding); + py::object encoder = encoderFactory("strict"); + const Py_ssize_t length = PyUnicode_GET_LENGTH(value); + constexpr Py_ssize_t chunkSize = 4096; + Py_ssize_t total = 0; + for (Py_ssize_t offset = 0; offset < length; offset += chunkSize) { + const Py_ssize_t end = std::min(offset + chunkSize, length); + py::object chunk = steal(PyUnicode_Substring(value, offset, end)); + if (!chunk) throw py::error_already_set(); + py::object encoded = encoder.attr("encode")(chunk, end == length); + const Py_ssize_t encodedSize = PyBytes_GET_SIZE(encoded.ptr()); + if (encodedSize > MAX_INLINE_CHAR - total) { + return MAX_INLINE_CHAR + 1; + } + total += encodedSize; + } + if (length == 0) { + py::object encoded = encoder.attr("encode")(py::str(), true); + total = PyBytes_GET_SIZE(encoded.ptr()); } - return PyBytes_GET_SIZE(encoded.ptr()); + return total; } inline PyObject* FormatDecimalParam(PyObject* params, Py_ssize_t index, PyObject* value) { From 3a57124f6337af1ab338d1de67690ce37e295381 Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Wed, 7 Oct 2026 12:42:46 +0530 Subject: [PATCH 13/39] FIX: Complete bounded DAE streaming Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/pybind/ddbc_bindings.cpp | 23 +++++++++++++++++------ 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index 66fea9790..15808e442 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -525,6 +525,7 @@ std::string DescribeChar(unsigned char ch) { template // The callable hides whether the caller wraps SQLPutData with GIL management; chunk sizing stays shared. static SQLRETURN stream_dae_chunks(const void* data, size_t total_bytes, PutDataFn put_data_fn) { + if (total_bytes == 0) return put_data_fn(nullptr, 0); const char* bytes = static_cast(data); for (size_t offset = 0; offset < total_bytes; offset += DAE_CHUNK_SIZE) { size_t len = std::min(static_cast(DAE_CHUNK_SIZE), total_bytes - offset); @@ -566,6 +567,7 @@ static SQLRETURN stream_unicode_dae_chunks(PyObject* value, const std::string& e throw py::error_already_set(); } if (size > 0) return put_data_fn(data, static_cast(size)); + return put_data_fn(nullptr, 0); } return SQL_SUCCESS; } @@ -597,12 +599,20 @@ static SQLRETURN StreamDAEParameter(SQLHSTMT hStmt, const ParamInfo& info, putImmutableData); } if (PyByteArray_Check(value)) { - auto putMutableData = [&](SQLPOINTER data, SQLLEN length) { - return SQLPutData_ptr(hStmt, data, length); - }; - return stream_dae_chunks(PyByteArray_AS_STRING(value), - static_cast(PyByteArray_GET_SIZE(value)), - putMutableData); + const size_t totalBytes = static_cast(PyByteArray_GET_SIZE(value)); + if (totalBytes == 0) return putImmutableData(nullptr, 0); + std::vector chunk(std::min(static_cast(DAE_CHUNK_SIZE), totalBytes)); + for (size_t offset = 0; offset < totalBytes; offset += chunk.size()) { + const size_t currentSize = static_cast(PyByteArray_GET_SIZE(value)); + const size_t length = std::min(chunk.size(), totalBytes - offset); + if (currentSize < offset + length) { + ThrowStdException("bytearray changed size during DAE streaming"); + } + std::copy_n(PyByteArray_AS_STRING(value) + offset, length, chunk.data()); + SQLRETURN rc = putImmutableData(chunk.data(), static_cast(length)); + if (!SQL_SUCCEEDED(rc)) return rc; + } + return SQL_SUCCESS; } ThrowStdException("DAE only supports str, bytes, or bytearray values"); } @@ -3313,6 +3323,7 @@ SQLRETURN SQLExecuteMany_wrap(const SqlHandlePtr statementHandle, const std::u16 dae_chunk_count, rc); return rc; } + if (rc != SQL_NEED_DATA) break; const ParamInfo* matchedInfo = reinterpret_cast(token); const ParamInfo* first = paramInfos.data(); From c64fbddbc96558f335a795122eae27a00f78c30e Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Wed, 7 Oct 2026 12:53:04 +0530 Subject: [PATCH 14/39] FIX: Align DAE and array metadata Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/cursor.py | 51 +++++++++++++++------------ mssql_python/pybind/ddbc_bindings.cpp | 7 ++-- tests/test_023_execute_path_parity.py | 10 ++++++ 3 files changed, 43 insertions(+), 25 deletions(-) diff --git a/mssql_python/cursor.py b/mssql_python/cursor.py index f80fdbeb4..b6cd81cf2 100644 --- a/mssql_python/cursor.py +++ b/mssql_python/cursor.py @@ -2633,30 +2633,35 @@ def executemany( # pylint: disable=too-many-locals,too-many-branches,too-many-s ): c_type = ddbc_sql_const.SQL_C_CHAR.value - # Check if this should be a DAE (data at execution) parameter. - if sample_value is not None: - if isinstance(sample_value, str): - if c_type == ddbc_sql_const.SQL_C_CHAR.value: - is_dae = column_size > MAX_INLINE_CHAR or any( - _encoded_length_exceeds( - value, - encoding_settings["encoding"], - MAX_INLINE_CHAR, - ) - for value in column - if isinstance(value, str) - ) - else: - text_size = max( - sum(2 if ord(char) > 0xFFFF else 1 for char in value) - for value in column - if isinstance(value, str) - ) - is_dae = ( - column_size > MAX_INLINE_CHAR or text_size > MAX_INLINE_CHAR + # Check every compatible value before selecting array binding. Mixed + # str/bytes columns cannot rely on the sample value's type. + if c_type in ( + ddbc_sql_const.SQL_C_CHAR.value, + ddbc_sql_const.SQL_C_WCHAR.value, + ): + text_values = [value for value in column if isinstance(value, str)] + if c_type == ddbc_sql_const.SQL_C_CHAR.value: + text_is_large = any( + _encoded_length_exceeds( + value, + encoding_settings["encoding"], + MAX_INLINE_CHAR, ) - elif isinstance(sample_value, (bytes, bytearray)) and column_size > 8000: - is_dae = True + for value in text_values + ) + binary_is_large = any( + len(value) > MAX_INLINE_CHAR + for value in column + if isinstance(value, (bytes, bytearray)) + ) + else: + text_is_large = any( + sum(2 if ord(char) > 0xFFFF else 1 for char in value) + > MAX_INLINE_CHAR + for value in text_values + ) + binary_is_large = False + is_dae = column_size > MAX_INLINE_CHAR or text_is_large or binary_is_large # Sanitize precision/scale for numeric types if sql_type in ( diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index 15808e442..6dd960106 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -391,7 +391,7 @@ void ReserveNativeParameterBytes(size_t& reservedBytes, size_t count, size_t ele size_t ParameterArrayElementSize(const ParamInfo& info) { switch (info.paramCType) { case SQL_C_LONG: - return sizeof(int); + return sizeof(int64_t); case SQL_C_DOUBLE: return sizeof(double); case SQL_C_WCHAR: @@ -773,6 +773,9 @@ SQLRETURN BindParameters(SqlHandle& handle, SQLHANDLE hStmt, const py::list& par for (int paramIndex = 0; paramIndex < params.size(); paramIndex++) { const auto& param = params[paramIndex]; ParamInfo& paramInfo = paramInfos[paramIndex]; + if (paramInfo.isDAE) { + paramInfo.dataPtr = py::reinterpret_borrow(param); + } LOG("BindParameters: Processing param[%d] - C_Type=%d, SQL_Type=%d, " "ColumnSize=%lu, DecimalDigits=%d, InputOutputType=%d", paramIndex, paramInfo.paramCType, paramInfo.paramSQLType, @@ -4725,7 +4728,7 @@ SQLRETURN SQLBindColums(SQLHSTMT hStmt, ColumnBuffers& buffers, const Metadata& reservedBytes); ret = SQLBindCol_ptr(hStmt, col, SQL_C_SS_TIMESTAMPOFFSET, buffers.datetimeoffsetBuffers[col - 1].data(), - sizeof(DateTimeOffset) * fetchSize, + sizeof(DateTimeOffset), buffers.indicators[col - 1].data()); break; default: diff --git a/tests/test_023_execute_path_parity.py b/tests/test_023_execute_path_parity.py index 595ae9adc..3a3c38935 100644 --- a/tests/test_023_execute_path_parity.py +++ b/tests/test_023_execute_path_parity.py @@ -329,6 +329,16 @@ def test_executemany_setinputsizes_uses_actual_text_size_for_dae(cursor, value, assert cursor.fetchall() == [(value,), ("short",)] +def test_executemany_setinputsizes_scans_mixed_text_values_for_dae(cursor): + """A bytes sample cannot hide a later oversized string from DAE selection.""" + value = "x" * 4001 + cursor.execute("CREATE TABLE #mixed_dae_batch (value VARCHAR(MAX))") + cursor.setinputsizes([(ddbc_sql_const.SQL_VARCHAR.value, 1, 0)]) + cursor.executemany("INSERT INTO #mixed_dae_batch (value) VALUES (?)", [(b"x",), (value,)]) + cursor.execute("SELECT value FROM #mixed_dae_batch ORDER BY LEN(value) DESC") + assert cursor.fetchall() == [(value,), ("x",)] + + @pytest.mark.parametrize("sql_type", [None, ddbc_sql_const.SQL_VARCHAR.value]) def test_time_isoformat_must_return_string(cursor, sql_type): """Native time normalization rejects a broken subclass contract on either path.""" From b8a8873f37839f1f4591e28cea07289f24b3fc4c Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Wed, 7 Oct 2026 13:14:41 +0530 Subject: [PATCH 15/39] FIX: Satisfy strict DAE return analysis Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/pybind/ddbc_bindings.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index 6dd960106..f20549905 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -615,6 +615,7 @@ static SQLRETURN StreamDAEParameter(SQLHSTMT hStmt, const ParamInfo& info, return SQL_SUCCESS; } ThrowStdException("DAE only supports str, bytes, or bytearray values"); + return SQL_ERROR; } // GH-610: Resolve SQL type for a NULL parameter using per-handle cache. From 9a55372eb6b8f5611713b4727b1e4f669b7b09a8 Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Wed, 7 Oct 2026 13:36:12 +0530 Subject: [PATCH 16/39] FIX: Close remaining native size gaps Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/pybind/ddbc_bindings.cpp | 33 +++++++++++++-------------- mssql_python/pybind/param_detect.hpp | 5 ++++ tests/test_023_execute_path_parity.py | 7 ++++++ 3 files changed, 28 insertions(+), 17 deletions(-) diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index f20549905..20987ed1c 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -391,7 +391,7 @@ void ReserveNativeParameterBytes(size_t& reservedBytes, size_t count, size_t ele size_t ParameterArrayElementSize(const ParamInfo& info) { switch (info.paramCType) { case SQL_C_LONG: - return sizeof(int64_t); + return sizeof(int); case SQL_C_DOUBLE: return sizeof(double); case SQL_C_WCHAR: @@ -883,8 +883,7 @@ SQLRETURN BindParameters(SqlHandle& handle, SQLHANDLE hStmt, const py::list& par break; } case SQL_C_WCHAR: { - if (!py::isinstance(param) && !py::isinstance(param) && - !py::isinstance(param)) { + if (!py::isinstance(param)) { ThrowStdException(MakeParamMismatchErrorStr(paramInfo.paramCType, paramIndex)); } if (paramInfo.isDAE) { @@ -5568,7 +5567,6 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, bool hasLobColumns = false; std::vector dataTypes(numCols); - std::vector columnSizes(numCols); std::vector columnNullable(numCols); std::vector columnVarLen(numCols, false); std::vector nullCounts(numCols, 0); @@ -5587,7 +5585,6 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, SQLSMALLINT nullable = colMeta["Nullable"].cast(); dataTypes[i] = dataType; - columnSizes[i] = columnSize; columnNullable[i] = (nullable != SQL_NO_NULLS); if ((dataType == SQL_WVARCHAR || dataType == SQL_WLONGVARCHAR || dataType == SQL_VARCHAR || @@ -6104,19 +6101,21 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, case SQL_LONGVARBINARY: { auto target_vec = &arrowColumnProducer->varData; auto start = arrowColumnProducer->varVal[idxRowArrow]; - EnsureNativeFetchBufferSize( - *target_vec, - CheckedAddSize(start, dataLen, "Arrow value buffer is too large"), - reservedBytes); - const size_t sourceStride = hasLobColumns - ? buffers.charBuffers[idxCol].size() - : buffers.charBuffers[idxCol].size() / - static_cast(fetchSize); - const size_t sourceOffset = CheckedArrowSourceOffset( - buffers.charBuffers[idxCol], idxRowSql, sourceStride, dataLen); + if (dataLen > 0) { + EnsureNativeFetchBufferSize( + *target_vec, + CheckedAddSize(start, dataLen, "Arrow value buffer is too large"), + reservedBytes); + const size_t sourceStride = hasLobColumns + ? buffers.charBuffers[idxCol].size() + : buffers.charBuffers[idxCol].size() / + static_cast(fetchSize); + const size_t sourceOffset = CheckedArrowSourceOffset( + buffers.charBuffers[idxCol], idxRowSql, sourceStride, dataLen); - std::memcpy(&(*target_vec)[start], - &buffers.charBuffers[idxCol][sourceOffset], dataLen); + std::memcpy(&(*target_vec)[start], + &buffers.charBuffers[idxCol][sourceOffset], dataLen); + } arrowColumnProducer->varVal[idxRowArrow + 1] = start + dataLen; break; } diff --git a/mssql_python/pybind/param_detect.hpp b/mssql_python/pybind/param_detect.hpp index f8a3e7d4f..32bff929f 100644 --- a/mssql_python/pybind/param_detect.hpp +++ b/mssql_python/pybind/param_detect.hpp @@ -296,6 +296,11 @@ inline void ApplyInputSizeOverride(PyObject* params, PyObject* inputSize, Py_ssi } } + if (info.paramCType == SQL_C_WCHAR && + (PyBytes_Check(obj) || PyByteArray_Check(obj))) { + throw py::type_error("bytes values cannot be bound as SQL_C_WCHAR"); + } + Py_ssize_t actualTextLength = 0; if (PyUnicode_Check(obj)) { actualTextLength = info.paramCType == SQL_C_CHAR diff --git a/tests/test_023_execute_path_parity.py b/tests/test_023_execute_path_parity.py index 3a3c38935..b7550131d 100644 --- a/tests/test_023_execute_path_parity.py +++ b/tests/test_023_execute_path_parity.py @@ -339,6 +339,13 @@ def test_executemany_setinputsizes_scans_mixed_text_values_for_dae(cursor): assert cursor.fetchall() == [(value,), ("x",)] +def test_setinputsizes_rejects_bytes_for_wide_character_binding(cursor): + """Wide-character overrides require text rather than raw encoded bytes.""" + cursor.setinputsizes([(ddbc_sql_const.SQL_WVARCHAR.value, 1, 0)]) + with pytest.raises(TypeError, match="SQL_C_WCHAR"): + cursor.execute("SELECT ?", [b"x" * 8001]) + + @pytest.mark.parametrize("sql_type", [None, ddbc_sql_const.SQL_VARCHAR.value]) def test_time_isoformat_must_return_string(cursor, sql_type): """Native time normalization rejects a broken subclass contract on either path.""" From 7bb3a1afc4c799a33842c185cfd933bb8eb8f6fa Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Wed, 7 Oct 2026 13:58:54 +0530 Subject: [PATCH 17/39] FIX: Restore driver-compatible DAE binding Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/cursor.py | 68 +++++++++++++++------------ mssql_python/pybind/ddbc_bindings.cpp | 1 + mssql_python/pybind/param_detect.hpp | 28 ++++++----- 3 files changed, 54 insertions(+), 43 deletions(-) diff --git a/mssql_python/cursor.py b/mssql_python/cursor.py index b6cd81cf2..934087a4e 100644 --- a/mssql_python/cursor.py +++ b/mssql_python/cursor.py @@ -1347,9 +1347,9 @@ def _reset_inputsizes(self) -> None: def _get_sql_to_c_type_map(cls): if cls._SQL_TO_C_TYPE is None: cls._SQL_TO_C_TYPE = { - ddbc_sql_const.SQL_CHAR.value: ddbc_sql_const.SQL_C_CHAR.value, - ddbc_sql_const.SQL_VARCHAR.value: ddbc_sql_const.SQL_C_CHAR.value, - ddbc_sql_const.SQL_LONGVARCHAR.value: ddbc_sql_const.SQL_C_CHAR.value, + ddbc_sql_const.SQL_CHAR.value: ddbc_sql_const.SQL_CHAR.value, + ddbc_sql_const.SQL_VARCHAR.value: ddbc_sql_const.SQL_CHAR.value, + ddbc_sql_const.SQL_LONGVARCHAR.value: ddbc_sql_const.SQL_CHAR.value, ddbc_sql_const.SQL_WCHAR.value: ddbc_sql_const.SQL_C_WCHAR.value, ddbc_sql_const.SQL_WVARCHAR.value: ddbc_sql_const.SQL_C_WCHAR.value, ddbc_sql_const.SQL_WLONGVARCHAR.value: ddbc_sql_const.SQL_C_WCHAR.value, @@ -2633,35 +2633,39 @@ def executemany( # pylint: disable=too-many-locals,too-many-branches,too-many-s ): c_type = ddbc_sql_const.SQL_C_CHAR.value - # Check every compatible value before selecting array binding. Mixed - # str/bytes columns cannot rely on the sample value's type. - if c_type in ( - ddbc_sql_const.SQL_C_CHAR.value, - ddbc_sql_const.SQL_C_WCHAR.value, + # Check every compatible value before selecting array binding. Use + # the SQL type because the public SQL_C_CHAR enum aliases SQL_C_WCHAR. + if sql_type in ( + ddbc_sql_const.SQL_CHAR.value, + ddbc_sql_const.SQL_VARCHAR.value, + ddbc_sql_const.SQL_LONGVARCHAR.value, ): text_values = [value for value in column if isinstance(value, str)] - if c_type == ddbc_sql_const.SQL_C_CHAR.value: - text_is_large = any( - _encoded_length_exceeds( - value, - encoding_settings["encoding"], - MAX_INLINE_CHAR, - ) - for value in text_values - ) - binary_is_large = any( - len(value) > MAX_INLINE_CHAR - for value in column - if isinstance(value, (bytes, bytearray)) + text_is_large = any( + _encoded_length_exceeds( + value, + encoding_settings["encoding"], + MAX_INLINE_CHAR, ) - else: - text_is_large = any( - sum(2 if ord(char) > 0xFFFF else 1 for char in value) - > MAX_INLINE_CHAR - for value in text_values - ) - binary_is_large = False - is_dae = column_size > MAX_INLINE_CHAR or text_is_large or binary_is_large + for value in text_values + ) + binary_is_large = any( + len(value) > MAX_INLINE_CHAR + for value in column + if isinstance(value, (bytes, bytearray)) + ) + is_dae = text_is_large or binary_is_large + elif sql_type in ( + ddbc_sql_const.SQL_WCHAR.value, + ddbc_sql_const.SQL_WVARCHAR.value, + ddbc_sql_const.SQL_WLONGVARCHAR.value, + ddbc_sql_const.SQL_SS_XML.value, + ): + is_dae = any( + sum(2 if ord(char) > 0xFFFF else 1 for char in value) > MAX_INLINE_CHAR + for value in column + if isinstance(value, str) + ) # Sanitize precision/scale for numeric types if sql_type in ( @@ -2676,6 +2680,7 @@ def executemany( # pylint: disable=too-many-locals,too-many-branches,too-many-s ddbc_sql_const.SQL_BINARY.value, ddbc_sql_const.SQL_VARBINARY.value, ddbc_sql_const.SQL_LONGVARBINARY.value, + ddbc_sql_const.SQL_SS_UDT.value, ): # Find the maximum size needed for any row's binary data max_binary_size = 0 @@ -2685,8 +2690,9 @@ def executemany( # pylint: disable=too-many-locals,too-many-branches,too-many-s max_binary_size = max(max_binary_size, len(value)) # For SQL Server VARBINARY(MAX), we need to use large object binding - if column_size > 8000 or max_binary_size > 8000: - sql_type = ddbc_sql_const.SQL_LONGVARBINARY.value + if max_binary_size > 8000: + if sql_type != ddbc_sql_const.SQL_SS_UDT.value: + sql_type = ddbc_sql_const.SQL_LONGVARBINARY.value is_dae = True # Update column_size to actual maximum size if it's larger diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index 20987ed1c..1a90c0ea9 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -4767,6 +4767,7 @@ SQLRETURN FetchBatchData(SQLHSTMT hStmt, ColumnBuffers& buffers, const Metadata& LOG("FetchBatchData: Fetching data in batches"); SQLRETURN ret; { + numRowsFetched = 0; // Release the GIL during the blocking ODBC fetch py::gil_scoped_release release; PERF_TIMER("FetchBatchData::SQLFetchScroll_call"); diff --git a/mssql_python/pybind/param_detect.hpp b/mssql_python/pybind/param_detect.hpp index 32bff929f..46adcb08c 100644 --- a/mssql_python/pybind/param_detect.hpp +++ b/mssql_python/pybind/param_detect.hpp @@ -308,17 +308,23 @@ inline void ApplyInputSizeOverride(PyObject* params, PyObject* inputSize, Py_ssi : UnicodeUtf16Length(obj); } const bool textNeedsDAE = - PyUnicode_Check(obj) && - (PyLongGreaterThan(columnSize, MAX_INLINE_CHAR) || - actualTextLength > MAX_INLINE_CHAR); + PyUnicode_Check(obj) && actualTextLength > MAX_INLINE_CHAR; const bool binaryNeedsDAE = (PyBytes_Check(obj) || PyByteArray_Check(obj)) && - (PyLongGreaterThan(columnSize, MAX_INLINE_BINARY) || - (PyBytes_Check(obj) ? PyBytes_GET_SIZE(obj) : PyByteArray_GET_SIZE(obj)) > - MAX_INLINE_BINARY); + (PyBytes_Check(obj) ? PyBytes_GET_SIZE(obj) : PyByteArray_GET_SIZE(obj)) > + MAX_INLINE_BINARY; info.isDAE = textNeedsDAE || binaryNeedsDAE; if (info.isDAE) { info.columnSize = 0; + if (textNeedsDAE) { + if (info.paramSQLType == SQL_CHAR || info.paramSQLType == SQL_VARCHAR) { + info.paramSQLType = SQL_LONGVARCHAR; + } else if (info.paramSQLType == SQL_WCHAR || info.paramSQLType == SQL_WVARCHAR) { + info.paramSQLType = SQL_WLONGVARCHAR; + } + } else if (info.paramSQLType == SQL_BINARY || info.paramSQLType == SQL_VARBINARY) { + info.paramSQLType = SQL_LONGVARBINARY; + } } if (PyTime_Check(obj) && info.paramCType == PARAM_C_TYPE_TEXT) { @@ -490,16 +496,13 @@ inline std::vector DetectParamTypes(PyObject* params, PyObject* input // Strings > 4000 UTF-16 code units exceed SQL Server's inline NVARCHAR(MAX) // threshold. Switch to data-at-execution (DAE) streaming: ODBC driver pulls // data in chunks via SQLPutData, avoiding a single massive buffer allocation. - // DAE path: match slow-path types exactly. - // Non-unicode (ASCII) → SQL_VARCHAR + PARAM_C_TYPE_TEXT, which is - // SQL_C_WCHAR and matches the slow path's SQL_C_CHAR (numerically - // -8 == SQL_C_WCHAR — a long-standing alias in the Python layer). - // Unicode → SQL_WVARCHAR + SQL_C_WCHAR (wide-char streaming) + // Long SQL types prevent a zero DAE ColumnSize from being interpreted + // as a fixed-width VARCHAR/NVARCHAR precision by the driver. info.isDAE = true; info.columnSize = 0; info.utf16Len = utf16_len; info.dataPtr = borrow(obj); - info.paramSQLType = is_unicode ? SQL_WVARCHAR : SQL_VARCHAR; + info.paramSQLType = is_unicode ? SQL_WLONGVARCHAR : SQL_LONGVARCHAR; info.paramCType = is_unicode ? SQL_C_WCHAR : PARAM_C_TYPE_TEXT; } else { info.columnSize = is_unicode ? utf16_len : length; @@ -520,6 +523,7 @@ inline std::vector DetectParamTypes(PyObject* params, PyObject* input info.isDAE = true; info.columnSize = 0; info.dataPtr = borrow(obj); + info.paramSQLType = SQL_LONGVARBINARY; } else { info.columnSize = std::max(length, 1); } From 37ec527db5b003d7bfd0ae1ec4983d27035c1db1 Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Wed, 7 Oct 2026 14:10:37 +0530 Subject: [PATCH 18/39] FIX: Bound wide LOB conversion memory Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/cursor.py | 9 +++++---- mssql_python/pybind/ddbc_bindings.cpp | 13 ++++++++----- 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/mssql_python/cursor.py b/mssql_python/cursor.py index 934087a4e..85621d5b8 100644 --- a/mssql_python/cursor.py +++ b/mssql_python/cursor.py @@ -48,6 +48,7 @@ MAX_INLINE_CHAR: int = ( 4000 # NVARCHAR/VARCHAR inline limit; this triggers NVARCHAR(MAX)/VARCHAR(MAX) + DAE ) +MAX_INLINE_BINARY: int = 8000 SMALLMONEY_MIN: decimal.Decimal = decimal.Decimal("-214748.3648") SMALLMONEY_MAX: decimal.Decimal = decimal.Decimal("214748.3647") MONEY_MIN: decimal.Decimal = decimal.Decimal("-922337203685477.5808") @@ -980,7 +981,7 @@ def _map_sql_type( # pylint: disable=too-many-arguments,too-many-positional-arg if isinstance(param, (bytes, bytearray)): length = len(param) - if length > 8000: # Use VARBINARY(MAX) for large blobs + if length > MAX_INLINE_BINARY: # Use VARBINARY(MAX) for large blobs return ( ddbc_sql_const.SQL_VARBINARY.value, ddbc_sql_const.SQL_C_BINARY.value, @@ -2650,7 +2651,7 @@ def executemany( # pylint: disable=too-many-locals,too-many-branches,too-many-s for value in text_values ) binary_is_large = any( - len(value) > MAX_INLINE_CHAR + len(value) > MAX_INLINE_BINARY for value in column if isinstance(value, (bytes, bytearray)) ) @@ -2690,7 +2691,7 @@ def executemany( # pylint: disable=too-many-locals,too-many-branches,too-many-s max_binary_size = max(max_binary_size, len(value)) # For SQL Server VARBINARY(MAX), we need to use large object binding - if max_binary_size > 8000: + if max_binary_size > MAX_INLINE_BINARY: if sql_type != ddbc_sql_const.SQL_SS_UDT.value: sql_type = ddbc_sql_const.SQL_LONGVARBINARY.value is_dae = True @@ -2780,7 +2781,7 @@ def executemany( # pylint: disable=too-many-locals,too-many-branches,too-many-s max_binary_size = max(max_binary_size, len(value)) # For SQL Server VARBINARY(MAX), we need to use large object binding - if max_binary_size > 8000: + if max_binary_size > MAX_INLINE_BINARY: paraminfo.paramSQLType = ddbc_sql_const.SQL_LONGVARBINARY.value paraminfo.isDAE = True diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index 1a90c0ea9..8a48c7d09 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -3617,11 +3617,14 @@ py::object FetchLobColumnData(SQLHSTMT hStmt, SQLUSMALLINT colIndex, SQLSMALLINT if (buffer.size() % sizeof(SQLWCHAR) != 0) { ThrowStdException("Wide-character LOB data has an invalid byte length"); } - size_t wcharCount = buffer.size() / sizeof(SQLWCHAR); - ReserveNativeFetchBytes(reservedBytes, wcharCount, sizeof(SQLWCHAR)); - std::vector alignedBuf(wcharCount); - std::memcpy(alignedBuf.data(), buffer.data(), buffer.size()); - return py::cast(dupeSqlWCharAsUtf16Le(alignedBuf.data(), wcharCount)); + int byteOrder = -1; + PyObject* decoded = + PyUnicode_DecodeUTF16(buffer.data(), static_cast(buffer.size()), + "strict", &byteOrder); + if (decoded == nullptr) { + throw py::error_already_set(); + } + return py::reinterpret_steal(decoded); } if (isBinary) { LOG("FetchLobColumnData: Returning binary data - %zu bytes for column " From 7b5a3b5ca6fa4854a771786d300a0bab18e61a3a Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Wed, 7 Oct 2026 14:25:14 +0530 Subject: [PATCH 19/39] FIX: Align streamed parameter metadata Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/cursor.py | 55 ++++++++++++++-------------- mssql_python/pybind/param_detect.hpp | 27 +++++++++----- 2 files changed, 45 insertions(+), 37 deletions(-) diff --git a/mssql_python/cursor.py b/mssql_python/cursor.py index 85621d5b8..1bdf51ce2 100644 --- a/mssql_python/cursor.py +++ b/mssql_python/cursor.py @@ -1348,9 +1348,9 @@ def _reset_inputsizes(self) -> None: def _get_sql_to_c_type_map(cls): if cls._SQL_TO_C_TYPE is None: cls._SQL_TO_C_TYPE = { - ddbc_sql_const.SQL_CHAR.value: ddbc_sql_const.SQL_CHAR.value, - ddbc_sql_const.SQL_VARCHAR.value: ddbc_sql_const.SQL_CHAR.value, - ddbc_sql_const.SQL_LONGVARCHAR.value: ddbc_sql_const.SQL_CHAR.value, + ddbc_sql_const.SQL_CHAR.value: ddbc_sql_const.SQL_C_CHAR.value, + ddbc_sql_const.SQL_VARCHAR.value: ddbc_sql_const.SQL_C_CHAR.value, + ddbc_sql_const.SQL_LONGVARCHAR.value: ddbc_sql_const.SQL_C_CHAR.value, ddbc_sql_const.SQL_WCHAR.value: ddbc_sql_const.SQL_C_WCHAR.value, ddbc_sql_const.SQL_WVARCHAR.value: ddbc_sql_const.SQL_C_WCHAR.value, ddbc_sql_const.SQL_WLONGVARCHAR.value: ddbc_sql_const.SQL_C_WCHAR.value, @@ -2613,6 +2613,7 @@ def executemany( # pylint: disable=too-many-locals,too-many-branches,too-many-s # Prepare parameter type information with perf_phase("py::executemany::param_type_detection"): for col_index in range(param_count): + requires_row_fallback = False column = ( [row[col_index] for row in seq_of_parameters] if hasattr(seq_of_parameters, "__getitem__") @@ -2634,39 +2635,37 @@ def executemany( # pylint: disable=too-many-locals,too-many-branches,too-many-s ): c_type = ddbc_sql_const.SQL_C_CHAR.value - # Check every compatible value before selecting array binding. Use - # the SQL type because the public SQL_C_CHAR enum aliases SQL_C_WCHAR. - if sql_type in ( + # Text values stay on the driver's wide binding path. Bytes supplied + # for narrow SQL text types use real SQL_C_CHAR in scalar execution. + narrow_text_type = sql_type in ( ddbc_sql_const.SQL_CHAR.value, ddbc_sql_const.SQL_VARCHAR.value, ddbc_sql_const.SQL_LONGVARCHAR.value, - ): - text_values = [value for value in column if isinstance(value, str)] - text_is_large = any( - _encoded_length_exceeds( - value, - encoding_settings["encoding"], - MAX_INLINE_CHAR, - ) - for value in text_values - ) - binary_is_large = any( - len(value) > MAX_INLINE_BINARY - for value in column - if isinstance(value, (bytes, bytearray)) - ) - is_dae = text_is_large or binary_is_large - elif sql_type in ( + ) + wide_text_type = sql_type in ( ddbc_sql_const.SQL_WCHAR.value, ddbc_sql_const.SQL_WVARCHAR.value, ddbc_sql_const.SQL_WLONGVARCHAR.value, ddbc_sql_const.SQL_SS_XML.value, - ): - is_dae = any( + ) + if narrow_text_type or wide_text_type: + text_values = [value for value in column if isinstance(value, str)] + binary_values = [ + value for value in column if isinstance(value, (bytes, bytearray)) + ] + text_is_large = any( sum(2 if ord(char) > 0xFFFF else 1 for char in value) > MAX_INLINE_CHAR - for value in column - if isinstance(value, str) + for value in text_values + ) + binary_is_large = narrow_text_type and any( + len(value) > MAX_INLINE_BINARY for value in binary_values ) + requires_row_fallback = bool( + narrow_text_type and text_values and binary_values + ) + if narrow_text_type and binary_values and not text_values: + c_type = ddbc_sql_const.SQL_CHAR.value + is_dae = text_is_large or binary_is_large # Sanitize precision/scale for numeric types if sql_type in ( @@ -2790,7 +2789,7 @@ def executemany( # pylint: disable=too-many-locals,too-many-branches,too-many-s paraminfo.columnSize = max(max_binary_size, 1) parameters_type.append(paraminfo) - if paraminfo.isDAE: + if paraminfo.isDAE or requires_row_fallback: any_dae = True if any_dae: diff --git a/mssql_python/pybind/param_detect.hpp b/mssql_python/pybind/param_detect.hpp index 46adcb08c..96b53246f 100644 --- a/mssql_python/pybind/param_detect.hpp +++ b/mssql_python/pybind/param_detect.hpp @@ -296,6 +296,12 @@ inline void ApplyInputSizeOverride(PyObject* params, PyObject* inputSize, Py_ssi } } + if ((PyBytes_Check(obj) || PyByteArray_Check(obj)) && + (info.paramSQLType == SQL_CHAR || info.paramSQLType == SQL_VARCHAR || + info.paramSQLType == SQL_LONGVARCHAR)) { + info.paramCType = SQL_C_CHAR; + } + if (info.paramCType == SQL_C_WCHAR && (PyBytes_Check(obj) || PyByteArray_Check(obj))) { throw py::type_error("bytes values cannot be bound as SQL_C_WCHAR"); @@ -315,13 +321,16 @@ inline void ApplyInputSizeOverride(PyObject* params, PyObject* inputSize, Py_ssi MAX_INLINE_BINARY; info.isDAE = textNeedsDAE || binaryNeedsDAE; if (info.isDAE) { - info.columnSize = 0; - if (textNeedsDAE) { - if (info.paramSQLType == SQL_CHAR || info.paramSQLType == SQL_VARCHAR) { - info.paramSQLType = SQL_LONGVARCHAR; - } else if (info.paramSQLType == SQL_WCHAR || info.paramSQLType == SQL_WVARCHAR) { - info.paramSQLType = SQL_WLONGVARCHAR; - } + info.columnSize = + textNeedsDAE + ? static_cast(actualTextLength) + : static_cast(PyBytes_Check(obj) ? PyBytes_GET_SIZE(obj) + : PyByteArray_GET_SIZE(obj)); + if (info.paramSQLType == SQL_CHAR || info.paramSQLType == SQL_VARCHAR) { + info.paramSQLType = SQL_LONGVARCHAR; + } else if (textNeedsDAE && + (info.paramSQLType == SQL_WCHAR || info.paramSQLType == SQL_WVARCHAR)) { + info.paramSQLType = SQL_WLONGVARCHAR; } else if (info.paramSQLType == SQL_BINARY || info.paramSQLType == SQL_VARBINARY) { info.paramSQLType = SQL_LONGVARBINARY; } @@ -499,7 +508,7 @@ inline std::vector DetectParamTypes(PyObject* params, PyObject* input // Long SQL types prevent a zero DAE ColumnSize from being interpreted // as a fixed-width VARCHAR/NVARCHAR precision by the driver. info.isDAE = true; - info.columnSize = 0; + info.columnSize = utf16_len; info.utf16Len = utf16_len; info.dataPtr = borrow(obj); info.paramSQLType = is_unicode ? SQL_WLONGVARCHAR : SQL_LONGVARCHAR; @@ -521,7 +530,7 @@ inline std::vector DetectParamTypes(PyObject* params, PyObject* input info.decimalDigits = 0; if (length > MAX_INLINE_BINARY) { info.isDAE = true; - info.columnSize = 0; + info.columnSize = static_cast(length); info.dataPtr = borrow(obj); info.paramSQLType = SQL_LONGVARBINARY; } else { From 190760f29090e6b0b1bfa0189388ddb34093dd7e Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Wed, 7 Oct 2026 14:38:21 +0530 Subject: [PATCH 20/39] FIX: Preserve MAX parameter type semantics Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/pybind/param_detect.hpp | 15 +++------------ 1 file changed, 3 insertions(+), 12 deletions(-) diff --git a/mssql_python/pybind/param_detect.hpp b/mssql_python/pybind/param_detect.hpp index 96b53246f..288c16ddf 100644 --- a/mssql_python/pybind/param_detect.hpp +++ b/mssql_python/pybind/param_detect.hpp @@ -326,14 +326,6 @@ inline void ApplyInputSizeOverride(PyObject* params, PyObject* inputSize, Py_ssi ? static_cast(actualTextLength) : static_cast(PyBytes_Check(obj) ? PyBytes_GET_SIZE(obj) : PyByteArray_GET_SIZE(obj)); - if (info.paramSQLType == SQL_CHAR || info.paramSQLType == SQL_VARCHAR) { - info.paramSQLType = SQL_LONGVARCHAR; - } else if (textNeedsDAE && - (info.paramSQLType == SQL_WCHAR || info.paramSQLType == SQL_WVARCHAR)) { - info.paramSQLType = SQL_WLONGVARCHAR; - } else if (info.paramSQLType == SQL_BINARY || info.paramSQLType == SQL_VARBINARY) { - info.paramSQLType = SQL_LONGVARBINARY; - } } if (PyTime_Check(obj) && info.paramCType == PARAM_C_TYPE_TEXT) { @@ -505,13 +497,13 @@ inline std::vector DetectParamTypes(PyObject* params, PyObject* input // Strings > 4000 UTF-16 code units exceed SQL Server's inline NVARCHAR(MAX) // threshold. Switch to data-at-execution (DAE) streaming: ODBC driver pulls // data in chunks via SQLPutData, avoiding a single massive buffer allocation. - // Long SQL types prevent a zero DAE ColumnSize from being interpreted - // as a fixed-width VARCHAR/NVARCHAR precision by the driver. + // Advertise the validated payload size rather than zero or the caller's + // declared size so the driver selects the corresponding MAX representation. info.isDAE = true; info.columnSize = utf16_len; info.utf16Len = utf16_len; info.dataPtr = borrow(obj); - info.paramSQLType = is_unicode ? SQL_WLONGVARCHAR : SQL_LONGVARCHAR; + info.paramSQLType = is_unicode ? SQL_WVARCHAR : SQL_VARCHAR; info.paramCType = is_unicode ? SQL_C_WCHAR : PARAM_C_TYPE_TEXT; } else { info.columnSize = is_unicode ? utf16_len : length; @@ -532,7 +524,6 @@ inline std::vector DetectParamTypes(PyObject* params, PyObject* input info.isDAE = true; info.columnSize = static_cast(length); info.dataPtr = borrow(obj); - info.paramSQLType = SQL_LONGVARBINARY; } else { info.columnSize = std::max(length, 1); } From 5d0a2e323ff3d5cc69e036bae25be18f3ed032d9 Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Wed, 7 Oct 2026 14:42:03 +0530 Subject: [PATCH 21/39] FIX: Size text by active C binding Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/cursor.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/mssql_python/cursor.py b/mssql_python/cursor.py index 1bdf51ce2..bafc19de6 100644 --- a/mssql_python/cursor.py +++ b/mssql_python/cursor.py @@ -2653,10 +2653,21 @@ def executemany( # pylint: disable=too-many-locals,too-many-branches,too-many-s binary_values = [ value for value in column if isinstance(value, (bytes, bytearray)) ] - text_is_large = any( - sum(2 if ord(char) > 0xFFFF else 1 for char in value) > MAX_INLINE_CHAR - for value in text_values - ) + if c_type == ddbc_sql_const.SQL_CHAR.value: + text_is_large = any( + _encoded_length_exceeds( + value, + encoding_settings["encoding"], + MAX_INLINE_CHAR, + ) + for value in text_values + ) + else: + text_is_large = any( + sum(2 if ord(char) > 0xFFFF else 1 for char in value) + > MAX_INLINE_CHAR + for value in text_values + ) binary_is_large = narrow_text_type and any( len(value) > MAX_INLINE_BINARY for value in binary_values ) From e98d558391d3e7513acef9c97608e18c318a795c Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Wed, 7 Oct 2026 14:54:31 +0530 Subject: [PATCH 22/39] FIX: Bind NULL rows independently in DAE batches Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/pybind/ddbc_bindings.cpp | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index 8a48c7d09..5f73fa572 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -3294,8 +3294,16 @@ SQLRETURN SQLExecuteMany_wrap(const SqlHandlePtr statementHandle, const std::u16 rowParams.append(values[rowIndex]); } + std::vector rowParamInfos = paramInfos; + for (size_t paramIndex = 0; paramIndex < rowParamInfos.size(); ++paramIndex) { + if (rowParams[paramIndex].is_none()) { + rowParamInfos[paramIndex].paramCType = SQL_C_DEFAULT; + rowParamInfos[paramIndex].isDAE = false; + rowParamInfos[paramIndex].dataPtr = py::none(); + } + } std::vector> paramBuffers; - rc = BindParameters(*statementHandle, hStmt, rowParams, paramInfos, + rc = BindParameters(*statementHandle, hStmt, rowParams, rowParamInfos, paramBuffers, charEncoding); if (!SQL_SUCCEEDED(rc)) { LOG("SQLExecuteMany: BindParameters failed for row %zu - rc=%d", rowIndex, rc); From 99e837cae8db6857cb28e1de68749777bf5cacdc Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Wed, 7 Oct 2026 15:12:11 +0530 Subject: [PATCH 23/39] FIX: Validate row-local DAE tokens Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/pybind/ddbc_bindings.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index 5f73fa572..aaaf7d0a1 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -3337,8 +3337,8 @@ SQLRETURN SQLExecuteMany_wrap(const SqlHandlePtr statementHandle, const std::u16 if (rc != SQL_NEED_DATA) break; const ParamInfo* matchedInfo = reinterpret_cast(token); - const ParamInfo* first = paramInfos.data(); - const ParamInfo* last = first + paramInfos.size(); + const ParamInfo* first = rowParamInfos.data(); + const ParamInfo* last = first + rowParamInfos.size(); if (matchedInfo < first || matchedInfo >= last) { ThrowStdException( "SQLExecuteMany: unrecognized paramToken from SQLParamData"); From a72841c8d08a8bfa2283ec45604d61f8c35d66bd Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Wed, 7 Oct 2026 15:25:44 +0530 Subject: [PATCH 24/39] FIX: Use MAX sentinel for oversized DAE Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/pybind/param_detect.hpp | 35 +++++++++++++++++++++++----- 1 file changed, 29 insertions(+), 6 deletions(-) diff --git a/mssql_python/pybind/param_detect.hpp b/mssql_python/pybind/param_detect.hpp index 288c16ddf..5a3e98d4f 100644 --- a/mssql_python/pybind/param_detect.hpp +++ b/mssql_python/pybind/param_detect.hpp @@ -140,6 +140,26 @@ inline constexpr int MAX_INLINE_CHAR = 4000; // Binary data longer than this uses DAE streaming (SQL Server max for non-MAX types) inline constexpr int MAX_INLINE_BINARY = 8000; +inline SQLULEN DAEColumnSize(SQLSMALLINT sqlType, SQLULEN actualSize) { + switch (sqlType) { + case SQL_CHAR: + case SQL_VARCHAR: + return actualSize > MAX_INLINE_BINARY ? 0 : actualSize; + case SQL_WCHAR: + case SQL_WVARCHAR: + return actualSize > MAX_INLINE_CHAR ? 0 : actualSize; + case SQL_BINARY: + case SQL_VARBINARY: + return actualSize > MAX_INLINE_BINARY ? 0 : actualSize; + case SQL_LONGVARCHAR: + case SQL_WLONGVARCHAR: + case SQL_LONGVARBINARY: + return 0; + default: + return actualSize; + } +} + // SQL Server maximum numeric precision inline constexpr int MAX_NUMERIC_PRECISION = 38; @@ -321,11 +341,12 @@ inline void ApplyInputSizeOverride(PyObject* params, PyObject* inputSize, Py_ssi MAX_INLINE_BINARY; info.isDAE = textNeedsDAE || binaryNeedsDAE; if (info.isDAE) { - info.columnSize = + const SQLULEN actualSize = textNeedsDAE ? static_cast(actualTextLength) : static_cast(PyBytes_Check(obj) ? PyBytes_GET_SIZE(obj) : PyByteArray_GET_SIZE(obj)); + info.columnSize = DAEColumnSize(info.paramSQLType, actualSize); } if (PyTime_Check(obj) && info.paramCType == PARAM_C_TYPE_TEXT) { @@ -497,13 +518,14 @@ inline std::vector DetectParamTypes(PyObject* params, PyObject* input // Strings > 4000 UTF-16 code units exceed SQL Server's inline NVARCHAR(MAX) // threshold. Switch to data-at-execution (DAE) streaming: ODBC driver pulls // data in chunks via SQLPutData, avoiding a single massive buffer allocation. - // Advertise the validated payload size rather than zero or the caller's - // declared size so the driver selects the corresponding MAX representation. + // Use the validated payload size when it is legal fixed-width metadata; + // larger values use the driver's MAX-length sentinel. info.isDAE = true; - info.columnSize = utf16_len; + const SQLSMALLINT sqlType = is_unicode ? SQL_WVARCHAR : SQL_VARCHAR; + info.columnSize = DAEColumnSize(sqlType, utf16_len); info.utf16Len = utf16_len; info.dataPtr = borrow(obj); - info.paramSQLType = is_unicode ? SQL_WVARCHAR : SQL_VARCHAR; + info.paramSQLType = sqlType; info.paramCType = is_unicode ? SQL_C_WCHAR : PARAM_C_TYPE_TEXT; } else { info.columnSize = is_unicode ? utf16_len : length; @@ -522,7 +544,8 @@ inline std::vector DetectParamTypes(PyObject* params, PyObject* input info.decimalDigits = 0; if (length > MAX_INLINE_BINARY) { info.isDAE = true; - info.columnSize = static_cast(length); + info.columnSize = + DAEColumnSize(SQL_VARBINARY, static_cast(length)); info.dataPtr = borrow(obj); } else { info.columnSize = std::max(length, 1); From 5919670f18d23b5511be0412d7a906ac65f9d6ec Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Wed, 7 Oct 2026 15:38:46 +0530 Subject: [PATCH 25/39] FIX: Size DAE metadata through MAX boundary Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/pybind/param_detect.hpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mssql_python/pybind/param_detect.hpp b/mssql_python/pybind/param_detect.hpp index 5a3e98d4f..fa2717038 100644 --- a/mssql_python/pybind/param_detect.hpp +++ b/mssql_python/pybind/param_detect.hpp @@ -154,7 +154,7 @@ inline SQLULEN DAEColumnSize(SQLSMALLINT sqlType, SQLULEN actualSize) { case SQL_LONGVARCHAR: case SQL_WLONGVARCHAR: case SQL_LONGVARBINARY: - return 0; + return actualSize; default: return actualSize; } @@ -237,8 +237,8 @@ inline Py_ssize_t EncodedUnicodeLength(PyObject* value, const std::string& encod if (!chunk) throw py::error_already_set(); py::object encoded = encoder.attr("encode")(chunk, end == length); const Py_ssize_t encodedSize = PyBytes_GET_SIZE(encoded.ptr()); - if (encodedSize > MAX_INLINE_CHAR - total) { - return MAX_INLINE_CHAR + 1; + if (encodedSize > MAX_INLINE_BINARY - total) { + return MAX_INLINE_BINARY + 1; } total += encodedSize; } From db18317db25563051a3780a09af0de6440385706 Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Wed, 7 Oct 2026 15:53:03 +0530 Subject: [PATCH 26/39] FIX: align DAE metadata with wide bindings Use the Unicode-safe MAX boundary for VARCHAR data-at-execution metadata and compare the wide fetch regression by row values rather than generated Row identity. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/pybind/param_detect.hpp | 2 +- tests/test_004_cursor.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/mssql_python/pybind/param_detect.hpp b/mssql_python/pybind/param_detect.hpp index fa2717038..d586daac3 100644 --- a/mssql_python/pybind/param_detect.hpp +++ b/mssql_python/pybind/param_detect.hpp @@ -144,7 +144,7 @@ inline SQLULEN DAEColumnSize(SQLSMALLINT sqlType, SQLULEN actualSize) { switch (sqlType) { case SQL_CHAR: case SQL_VARCHAR: - return actualSize > MAX_INLINE_BINARY ? 0 : actualSize; + return actualSize > MAX_INLINE_CHAR ? 0 : actualSize; case SQL_WCHAR: case SQL_WVARCHAR: return actualSize > MAX_INLINE_CHAR ? 0 : actualSize; diff --git a/tests/test_004_cursor.py b/tests/test_004_cursor.py index 518dd15b4..02dec2e9f 100644 --- a/tests/test_004_cursor.py +++ b/tests/test_004_cursor.py @@ -1748,7 +1748,7 @@ def test_fetchmany_rejects_excessive_native_buffer(cursor): def test_fetchall_clamps_wide_result_batch_to_native_buffer_budget(cursor): cursor.execute("SELECT " + ", ".join("CAST(N'x' AS NVARCHAR(4000))" for _ in range(34))) - assert cursor.fetchall() == [("x",) * 34] + assert [tuple(row) for row in cursor.fetchall()] == [("x",) * 34] def test_description(cursor): From 400bd0f2191ac6ff3d2da45855a296e2f8ef16bc Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Wed, 7 Oct 2026 16:11:21 +0530 Subject: [PATCH 27/39] FIX: preserve actual parameter metadata size Prevent a too-small setinputsizes declaration from truncating later non-DAE text or binary rows after executemany falls back to row execution. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/pybind/param_detect.hpp | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/mssql_python/pybind/param_detect.hpp b/mssql_python/pybind/param_detect.hpp index d586daac3..b7ec6a189 100644 --- a/mssql_python/pybind/param_detect.hpp +++ b/mssql_python/pybind/param_detect.hpp @@ -334,19 +334,23 @@ inline void ApplyInputSizeOverride(PyObject* params, PyObject* inputSize, Py_ssi : UnicodeUtf16Length(obj); } const bool textNeedsDAE = - PyUnicode_Check(obj) && actualTextLength > MAX_INLINE_CHAR; + !isNumeric && PyUnicode_Check(obj) && actualTextLength > MAX_INLINE_CHAR; const bool binaryNeedsDAE = (PyBytes_Check(obj) || PyByteArray_Check(obj)) && (PyBytes_Check(obj) ? PyBytes_GET_SIZE(obj) : PyByteArray_GET_SIZE(obj)) > MAX_INLINE_BINARY; info.isDAE = textNeedsDAE || binaryNeedsDAE; - if (info.isDAE) { - const SQLULEN actualSize = - textNeedsDAE - ? static_cast(actualTextLength) - : static_cast(PyBytes_Check(obj) ? PyBytes_GET_SIZE(obj) - : PyByteArray_GET_SIZE(obj)); - info.columnSize = DAEColumnSize(info.paramSQLType, actualSize); + if (!isNumeric && PyUnicode_Check(obj)) { + const SQLULEN actualSize = static_cast(actualTextLength); + info.columnSize = + info.isDAE ? DAEColumnSize(info.paramSQLType, actualSize) + : std::max(info.columnSize, actualSize); + } else if (PyBytes_Check(obj) || PyByteArray_Check(obj)) { + const SQLULEN actualSize = static_cast( + PyBytes_Check(obj) ? PyBytes_GET_SIZE(obj) : PyByteArray_GET_SIZE(obj)); + info.columnSize = + info.isDAE ? DAEColumnSize(info.paramSQLType, actualSize) + : std::max(info.columnSize, actualSize); } if (PyTime_Check(obj) && info.paramCType == PARAM_C_TYPE_TEXT) { From 26c4763d712585870f3b7fb9b1aa16c00607dec6 Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Wed, 7 Oct 2026 16:31:29 +0530 Subject: [PATCH 28/39] FIX: address native validation review findings Reject boolean row counts before conversion, guard empty Arrow character copies, validate bytearray inputs before C API access, and make the DAE regressions compare row values reliably. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/pybind/ddbc_bindings.cpp | 57 ++++++++++++++++++--------- tests/test_010_pybind_functions.py | 4 +- tests/test_023_execute_path_parity.py | 15 ++++--- 3 files changed, 50 insertions(+), 26 deletions(-) diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index aaaf7d0a1..c9b2c0dce 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -378,6 +378,15 @@ void ValidateNativeRowCount(int value, const char* name, bool allowZero) { } } +int ValidateNativeRowCountArgument(py::handle value, const char* name, bool allowZero) { + if (PyBool_Check(value.ptr())) { + ThrowStdException(std::string(name) + " must be an integer, not bool"); + } + const int converted = value.cast(); + ValidateNativeRowCount(converted, name, allowZero); + return converted; +} + void ReserveNativeParameterBytes(size_t& reservedBytes, size_t count, size_t elementSize) { const size_t allocationBytes = CheckedMultiplySize(count, elementSize, "Parameter buffer size is too large"); @@ -2769,12 +2778,15 @@ SQLRETURN BindParameterArray(SqlHandle& handle, SQLHANDLE hStmt, const py::list& &encodedSize) != 0) { throw py::error_already_set(); } - } else { + } else if (py::isinstance(columnValues[i])) { encodedData = PyByteArray_AsString(columnValues[i].ptr()); encodedSize = PyByteArray_Size(columnValues[i].ptr()); if (encodedData == nullptr || encodedSize < 0) { throw py::error_already_set(); } + } else { + ThrowStdException( + MakeParamMismatchErrorStr(info.paramCType, paramIndex)); } const size_t dataSize = static_cast(encodedSize); @@ -5312,12 +5324,13 @@ struct FetchStateGuard { // the result set and populates the provided Python list with the row data. If // there are no more rows to fetch, it returns SQL_NO_DATA. If an error occurs // during fetching, it throws a runtime error. -SQLRETURN FetchMany_wrap(SqlHandlePtr StatementHandle, py::list& rows, int fetchSize, +SQLRETURN FetchMany_wrap(SqlHandlePtr StatementHandle, py::list& rows, py::handle fetchSizeArg, const std::string& charEncoding = "utf-16le", const std::string& wcharEncoding = "utf-16le", int charCtype = SQL_C_WCHAR, py::handle messages = {}) { PERF_TIMER("FetchMany_wrap"); - ValidateNativeRowCount(fetchSize, "Fetch size", false); + const int fetchSize = + ValidateNativeRowCountArgument(fetchSizeArg, "Fetch size", false); // 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. charCtype = EffectiveCharCtypeForFetch(charCtype, charEncoding); @@ -5543,10 +5556,12 @@ 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, py::handle messages = {}) { +SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, + py::handle arrowBatchSizeArg, int charCtype, + py::handle messages = {}) { PERF_TIMER("FetchArrowBatch_wrap"); - ValidateNativeRowCount(arrowBatchSize, "Arrow batch size", true); + const int arrowBatchSize = + ValidateNativeRowCountArgument(arrowBatchSizeArg, "Arrow batch size", true); const size_t batchSize = static_cast(arrowBatchSize); const size_t offsetCount = CheckedAddSize(batchSize, 1, "Arrow batch size is too large"); const size_t initialVarDataSize = @@ -6137,19 +6152,23 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, if (charCtype == SQL_C_CHAR) { auto target_vec = &arrowColumnProducer->varData; auto start = arrowColumnProducer->varVal[idxRowArrow]; - EnsureNativeFetchBufferSize( - *target_vec, - CheckedAddSize(start, dataLen, "Arrow value buffer is too large"), - reservedBytes); - const size_t sourceStride = hasLobColumns - ? buffers.charBuffers[idxCol].size() - : buffers.charBuffers[idxCol].size() / - static_cast(fetchSize); - const size_t sourceOffset = CheckedArrowSourceOffset( - buffers.charBuffers[idxCol], idxRowSql, sourceStride, dataLen); - - std::memcpy(&(*target_vec)[start], - &buffers.charBuffers[idxCol][sourceOffset], dataLen); + if (dataLen > 0) { + EnsureNativeFetchBufferSize( + *target_vec, + CheckedAddSize(start, dataLen, + "Arrow value buffer is too large"), + reservedBytes); + const size_t sourceStride = + hasLobColumns + ? buffers.charBuffers[idxCol].size() + : buffers.charBuffers[idxCol].size() / + static_cast(fetchSize); + const size_t sourceOffset = CheckedArrowSourceOffset( + buffers.charBuffers[idxCol], idxRowSql, sourceStride, dataLen); + + std::memcpy(&(*target_vec)[start], + &buffers.charBuffers[idxCol][sourceOffset], dataLen); + } arrowColumnProducer->varVal[idxRowArrow + 1] = start + dataLen; break; } diff --git a/tests/test_010_pybind_functions.py b/tests/test_010_pybind_functions.py index afed788e4..ec85f14ed 100644 --- a/tests/test_010_pybind_functions.py +++ b/tests/test_010_pybind_functions.py @@ -38,14 +38,14 @@ @pytest.mark.skipif(not DDBC_AVAILABLE, reason="ddbc_bindings not available") -@pytest.mark.parametrize("fetch_size", [-1, 1_000_001]) +@pytest.mark.parametrize("fetch_size", [-1, 1_000_001, True]) def test_fetchmany_rejects_unsafe_size_before_handle_access(fetch_size): with pytest.raises(RuntimeError, match="Fetch size"): ddbc.DDBCSQLFetchMany(None, [], fetch_size) @pytest.mark.skipif(not DDBC_AVAILABLE, reason="ddbc_bindings not available") -@pytest.mark.parametrize("batch_size", [-1, 1_000_001]) +@pytest.mark.parametrize("batch_size", [-1, 1_000_001, True]) def test_arrow_batch_rejects_unsafe_size_before_handle_access(batch_size): with pytest.raises(RuntimeError, match="Arrow batch size"): ddbc.DDBCSQLFetchArrowBatch(None, [], batch_size, 0) diff --git a/tests/test_023_execute_path_parity.py b/tests/test_023_execute_path_parity.py index b7550131d..daeb42514 100644 --- a/tests/test_023_execute_path_parity.py +++ b/tests/test_023_execute_path_parity.py @@ -322,11 +322,16 @@ def test_setinputsizes_uses_actual_value_size_for_dae(cursor, value, sql_type): ) def test_executemany_setinputsizes_uses_actual_text_size_for_dae(cursor, value, sql_type): """A too-small declared size cannot bypass streaming for an array-bound text value.""" - cursor.execute("CREATE TABLE #dae_batch (value NVARCHAR(MAX))") + table_name = ( + "#dae_batch_varchar" + if sql_type == ddbc_sql_const.SQL_VARCHAR.value + else "#dae_batch_nvarchar" + ) + cursor.execute(f"CREATE TABLE {table_name} (value NVARCHAR(MAX))") cursor.setinputsizes([(sql_type, 1, 0)]) - cursor.executemany("INSERT INTO #dae_batch (value) VALUES (?)", [(value,), ("short",)]) - cursor.execute("SELECT value FROM #dae_batch ORDER BY LEN(value) DESC") - assert cursor.fetchall() == [(value,), ("short",)] + cursor.executemany(f"INSERT INTO {table_name} (value) VALUES (?)", [(value,), ("short",)]) + cursor.execute(f"SELECT value FROM {table_name} ORDER BY LEN(value) DESC") + assert [tuple(row) for row in cursor.fetchall()] == [(value,), ("short",)] def test_executemany_setinputsizes_scans_mixed_text_values_for_dae(cursor): @@ -336,7 +341,7 @@ def test_executemany_setinputsizes_scans_mixed_text_values_for_dae(cursor): cursor.setinputsizes([(ddbc_sql_const.SQL_VARCHAR.value, 1, 0)]) cursor.executemany("INSERT INTO #mixed_dae_batch (value) VALUES (?)", [(b"x",), (value,)]) cursor.execute("SELECT value FROM #mixed_dae_batch ORDER BY LEN(value) DESC") - assert cursor.fetchall() == [(value,), ("x",)] + assert [tuple(row) for row in cursor.fetchall()] == [(value,), ("x",)] def test_setinputsizes_rejects_bytes_for_wide_character_binding(cursor): From f53323085fcf86a70bc68804b17e8b031a5818b5 Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Wed, 7 Oct 2026 16:46:04 +0530 Subject: [PATCH 29/39] FIX: reset DAE bindings after each row Release ODBC parameter bindings before row-local buffers and SQLParamData tokens leave scope, and surface reset failures without disturbing execution diagnostics on error paths. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/pybind/ddbc_bindings.cpp | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index c9b2c0dce..4190b27e2 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -3367,6 +3367,16 @@ SQLRETURN SQLExecuteMany_wrap(const SqlHandlePtr statementHandle, const std::u16 LOG("SQLExecuteMany: DAE row %zu failed - rc=%d", rowIndex, rc); return rc; } + + const SQLRETURN rowRc = rc; + rc = SQLFreeStmt_ptr(hStmt, SQL_RESET_PARAMS); + if (!SQL_SUCCEEDED(rc)) { + LOG("SQLExecuteMany: SQL_RESET_PARAMS failed for row %zu - " + "rc=%d", + rowIndex, rc); + return rc; + } + rc = rowRc; } LOG("SQLExecuteMany: All DAE rows processed successfully - " "total_rows=%zu", From 48a2f23af05e1f2205c0a38f54e0c95758e43bc7 Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Thu, 8 Oct 2026 11:51:57 +0530 Subject: [PATCH 30/39] FIX: restore geometric LOB buffer growth Track LOB payload length separately from geometric allocation size, validate native executemany counts and input-size metadata before handle mutation, and add direct boundary and growth-count regressions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/pybind/ddbc_bindings.cpp | 56 +++++++++++++---- mssql_python/pybind/param_detect.hpp | 88 +++++++++++++++++++++++++++ tests/test_010_pybind_functions.py | 36 +++++++++++ 3 files changed, 167 insertions(+), 13 deletions(-) diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index 4190b27e2..21ab00f33 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -511,6 +511,29 @@ void EnsureNativeFetchBufferSize(std::vector& buffer, size_t requir ResizeNativeFetchBuffer(buffer, newSize, reservedBytes); } +size_t ExtendNativeFetchBuffer(std::vector& buffer, size_t currentSize, size_t appendSize, + size_t& reservedBytes) { + const size_t requiredSize = + CheckedAddSize(currentSize, appendSize, "Native fetch buffer size is too large"); + EnsureNativeFetchBufferSize(buffer, requiredSize, reservedBytes); + return requiredSize; +} + +py::tuple TestLobBufferGrowth(size_t chunkCount, size_t chunkSize) { + std::vector buffer; + size_t dataSize = 0; + size_t reservedBytes = 0; + size_t growthCount = 0; + for (size_t chunk = 0; chunk < chunkCount; ++chunk) { + const size_t previousAllocation = buffer.size(); + dataSize = ExtendNativeFetchBuffer(buffer, dataSize, chunkSize, reservedBytes); + if (buffer.size() != previousAllocation) { + ++growthCount; + } + } + return py::make_tuple(growthCount, dataSize, buffer.size(), reservedBytes); +} + template std::unique_ptr AllocateArrowArray(size_t count, size_t& reservedBytes, const char* errorMessage) { @@ -2396,6 +2419,7 @@ SQLRETURN SQLExecute_wrap(const SqlHandlePtr statementHandle, bool use_prepare, const py::dict& encoding_settings) { PERF_TIMER("SQLExecute_wrap"); + ValidateInputSizes(input_sizes.ptr()); if (!statementHandle || !statementHandle->get()) { return SQL_INVALID_HANDLE; } @@ -3218,9 +3242,14 @@ SQLRETURN BindParameterArray(SqlHandle& handle, SQLHANDLE hStmt, const py::list& SQLRETURN SQLExecuteMany_wrap(const SqlHandlePtr statementHandle, const std::u16string& query, const py::list& columnwise_params, - std::vector& paramInfos, size_t paramSetSize, + std::vector& paramInfos, py::handle paramSetSizeArg, const py::dict& encodingSettings) { PERF_TIMER("SQLExecuteMany_wrap"); + const size_t paramSetSize = static_cast( + ValidateNativeRowCountArgument(paramSetSizeArg, "Parameter set size", false)); + if (!statementHandle || !statementHandle->get()) { + return SQL_INVALID_HANDLE; + } statementHandle->resultMetadata.clear(); LOG("SQLExecuteMany: Starting batch execution - param_count=%zu, " "param_set_size=%zu", @@ -3552,6 +3581,7 @@ py::object FetchLobColumnData(SQLHSTMT hStmt, SQLUSMALLINT colIndex, SQLSMALLINT py::handle messages) { PERF_TIMER("FetchLobColumnData"); std::vector buffer; + size_t dataSize = 0; size_t reservedBytes = 0; ReserveNativeFetchBytes(reservedBytes, DAE_CHUNK_SIZE, sizeof(char)); SQLRETURN ret = SQL_SUCCESS_WITH_INFO; @@ -3623,10 +3653,8 @@ py::object FetchLobColumnData(SQLHSTMT hStmt, SQLUSMALLINT colIndex, SQLSMALLINT } } if (bytesRead > 0) { - const size_t previousSize = buffer.size(); - const size_t requiredSize = - CheckedAddSize(previousSize, bytesRead, "LOB fetch buffer is too large"); - ResizeNativeFetchBuffer(buffer, requiredSize, reservedBytes); + const size_t previousSize = dataSize; + dataSize = ExtendNativeFetchBuffer(buffer, dataSize, bytesRead, reservedBytes); std::copy_n(chunk.data(), bytesRead, buffer.data() + previousSize); LOG("FetchLobColumnData: Appended %zu bytes at loop %d", bytesRead, loopCount); } @@ -3635,21 +3663,21 @@ py::object FetchLobColumnData(SQLHSTMT hStmt, SQLUSMALLINT colIndex, SQLSMALLINT break; } } - LOG("FetchLobColumnData: Total bytes collected=%zu for column %d", buffer.size(), colIndex); + LOG("FetchLobColumnData: Total bytes collected=%zu for column %d", dataSize, colIndex); - if (buffer.empty()) { + if (dataSize == 0) { if (isBinary) { return py::bytes(""); } return py::str(""); } if (isWideChar) { - if (buffer.size() % sizeof(SQLWCHAR) != 0) { + if (dataSize % sizeof(SQLWCHAR) != 0) { ThrowStdException("Wide-character LOB data has an invalid byte length"); } int byteOrder = -1; PyObject* decoded = - PyUnicode_DecodeUTF16(buffer.data(), static_cast(buffer.size()), + PyUnicode_DecodeUTF16(buffer.data(), static_cast(dataSize), "strict", &byteOrder); if (decoded == nullptr) { throw py::error_already_set(); @@ -3659,18 +3687,18 @@ py::object FetchLobColumnData(SQLHSTMT hStmt, SQLUSMALLINT colIndex, SQLSMALLINT if (isBinary) { LOG("FetchLobColumnData: Returning binary data - %zu bytes for column " "%d", - buffer.size(), colIndex); - return py::bytes(buffer.data(), buffer.size()); + dataSize, colIndex); + return py::bytes(buffer.data(), dataSize); } // For SQL_C_CHAR data, decode using the appropriate encoding. const std::string effectiveCharEncoding = GetEffectiveCharDecoding(charEncoding); - py::bytes raw_bytes(buffer.data(), buffer.size()); + py::bytes raw_bytes(buffer.data(), dataSize); try { py::object decoded = raw_bytes.attr("decode")(effectiveCharEncoding, "strict"); LOG("FetchLobColumnData: Decoded narrow string with '%s' - %zu bytes -> %zu chars for " "column %d", - effectiveCharEncoding.c_str(), buffer.size(), py::len(decoded), colIndex); + effectiveCharEncoding.c_str(), dataSize, py::len(decoded), colIndex); return decoded; } catch (const py::error_already_set& e) { LOG_ERROR("FetchLobColumnData: Failed to decode with '%s' for column %d: %s", @@ -7000,6 +7028,8 @@ PYBIND11_MODULE(ddbc_bindings, m) { "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("_test_lob_buffer_growth", &TestLobBufferGrowth, py::arg("chunk_count"), + py::arg("chunk_size")); 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/param_detect.hpp b/mssql_python/pybind/param_detect.hpp index b7ec6a189..6443d3d23 100644 --- a/mssql_python/pybind/param_detect.hpp +++ b/mssql_python/pybind/param_detect.hpp @@ -274,6 +274,89 @@ inline void NormalizeTimeParam(PyObject* params, Py_ssize_t index, SQLULEN& colu } } +inline long long ValidatedInputSizeInteger(PyObject* value, const char* fieldName) { + if (!PyLong_Check(value) || PyBool_Check(value)) { + throw py::type_error(std::string(fieldName) + " must be an integer"); + } + int overflow = 0; + const long long result = PyLong_AsLongLongAndOverflow(value, &overflow); + if ((result == -1 && PyErr_Occurred()) || overflow != 0) { + PyErr_Clear(); + throw py::value_error(std::string(fieldName) + " is out of range"); + } + return result; +} + +inline void ValidateInputSizes(PyObject* inputSizes) { + if (inputSizes == Py_None) { + return; + } + if (!PyList_Check(inputSizes)) { + throw py::type_error("inputSizes must be None or a list"); + } + + const Py_ssize_t count = PyList_GET_SIZE(inputSizes); + for (Py_ssize_t index = 0; index < count; ++index) { + PyObject* entry = PyList_GET_ITEM(inputSizes, index); + if (!PyTuple_Check(entry) || PyTuple_GET_SIZE(entry) != 4) { + throw py::type_error("each inputSizes entry must be a four-item tuple"); + } + + const long long sqlType = + ValidatedInputSizeInteger(PyTuple_GET_ITEM(entry, 0), "SQL type"); + const long long cType = + ValidatedInputSizeInteger(PyTuple_GET_ITEM(entry, 1), "C type"); + PyObject* columnSize = PyTuple_GET_ITEM(entry, 2); + PyObject* decimalDigits = PyTuple_GET_ITEM(entry, 3); + if (!PyLong_Check(columnSize) || PyBool_Check(columnSize) || + !PyLong_Check(decimalDigits) || PyBool_Check(decimalDigits)) { + throw py::type_error("column size and decimal digits must be integers"); + } + + if (sqlType < std::numeric_limits::min() || + sqlType > std::numeric_limits::max() || + cType < std::numeric_limits::min() || + cType > std::numeric_limits::max()) { + throw py::value_error("SQL and C types must fit in SQLSMALLINT"); + } + const bool isNumeric = sqlType == SQL_DECIMAL || sqlType == SQL_NUMERIC; + py::int_ zero(0); + const int negativeColumnSize = + PyObject_RichCompareBool(columnSize, zero.ptr(), Py_LT); + const int negativeDecimalDigits = + PyObject_RichCompareBool(decimalDigits, zero.ptr(), Py_LT); + if (negativeColumnSize == -1 || negativeDecimalDigits == -1) { + throw py::error_already_set(); + } + if (negativeColumnSize == 1) { + throw py::value_error("column size must be non-negative"); + } + if (negativeDecimalDigits == 1) { + throw py::value_error("decimal digits must be non-negative"); + } + if (!isNumeric) { + const unsigned long long requestedSize = PyLong_AsUnsignedLongLong(columnSize); + if (requestedSize == static_cast(-1) && PyErr_Occurred()) { + PyErr_Clear(); + throw py::value_error("column size is out of range"); + } + if (requestedSize > std::numeric_limits::max()) { + throw py::value_error("column size is out of range"); + } + const unsigned long long requestedDigits = + PyLong_AsUnsignedLongLong(decimalDigits); + if (requestedDigits == static_cast(-1) && PyErr_Occurred()) { + PyErr_Clear(); + throw py::value_error("decimal digits are out of range"); + } + if (requestedDigits > + static_cast(std::numeric_limits::max())) { + throw py::value_error("decimal digits are out of range"); + } + } + } +} + inline void ApplyInputSizeOverride(PyObject* params, PyObject* inputSize, Py_ssize_t index, ParamInfo& info, const std::string& charEncoding) { py::tuple values = borrow(inputSize); @@ -391,6 +474,11 @@ inline std::vector DetectParamTypes(PyObject* params, PyObject* input const std::string& charEncoding = "utf-8") { PyTypeCache::initialize(); + if (!PyList_Check(params)) { + throw py::type_error("params must be a list"); + } + ValidateInputSizes(inputSizes); + const Py_ssize_t n = PyList_GET_SIZE(params); const Py_ssize_t inputSizeCount = inputSizes == Py_None ? 0 : PyList_GET_SIZE(inputSizes); std::vector infos(n); diff --git a/tests/test_010_pybind_functions.py b/tests/test_010_pybind_functions.py index ec85f14ed..ce4ec61db 100644 --- a/tests/test_010_pybind_functions.py +++ b/tests/test_010_pybind_functions.py @@ -51,6 +51,42 @@ def test_arrow_batch_rejects_unsafe_size_before_handle_access(batch_size): ddbc.DDBCSQLFetchArrowBatch(None, [], batch_size, 0) +@pytest.mark.skipif(not DDBC_AVAILABLE, reason="ddbc_bindings not available") +@pytest.mark.parametrize("param_set_size", [-1, 0, 1_000_001, True]) +def test_executemany_rejects_unsafe_size_before_handle_access(param_set_size): + with pytest.raises(RuntimeError, match="Parameter set size"): + ddbc.SQLExecuteMany(None, "", [], [], param_set_size, {}) + + +@pytest.mark.skipif(not DDBC_AVAILABLE, reason="ddbc_bindings not available") +@pytest.mark.parametrize( + ("input_sizes", "message"), + [ + ((12, -8, 1, 0), "inputSizes must be None or a list"), + ([None], "four-item tuple"), + ([(12, -8, 1)], "four-item tuple"), + ([(12, -8, True, 0)], "must be integers"), + ([(12, -8, -1, 0)], "column size must be non-negative"), + ([(12, -8, 2**128, 0)], "column size is out of range"), + ([(12, -8, 1, 2**128)], "decimal digits are out of range"), + ], +) +def test_execute_rejects_malformed_input_sizes_before_handle_access(input_sizes, message): + with pytest.raises((TypeError, ValueError), match=message): + ddbc.DDBCSQLExecute(None, "", [], input_sizes, [False], True, {}) + + +@pytest.mark.skipif(not DDBC_AVAILABLE, reason="ddbc_bindings not available") +def test_lob_buffer_growth_is_geometric(): + growth_count, data_size, allocation_size, reserved_bytes = ddbc._test_lob_buffer_growth( + 64, 64 * 1024 + ) + assert data_size == 4 * 1024 * 1024 + assert allocation_size == data_size + assert reserved_bytes == allocation_size + assert growth_count == 7 + + @pytest.mark.skipif(not DDBC_AVAILABLE, reason="ddbc_bindings not available") class TestPybindModuleInfo: """Test module information and architecture detection.""" From d0912f643faf4b0e789cb7b1a4c8e5b38929387b Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Thu, 8 Oct 2026 12:35:44 +0530 Subject: [PATCH 31/39] FIX: bound LOB growth test helper Reject boolean, zero, oversized chunk counts and invalid chunk sizes before the internal production test binding can enter its simulation loop. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/pybind/ddbc_bindings.cpp | 11 ++++++++++- tests/test_010_pybind_functions.py | 15 +++++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index 21ab00f33..b89ca350c 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -519,7 +519,16 @@ size_t ExtendNativeFetchBuffer(std::vector& buffer, size_t currentSize, si return requiredSize; } -py::tuple TestLobBufferGrowth(size_t chunkCount, size_t chunkSize) { +py::tuple TestLobBufferGrowth(py::handle chunkCountArg, py::handle chunkSizeArg) { + const size_t chunkCount = static_cast( + ValidateNativeRowCountArgument(chunkCountArg, "LOB growth chunk count", false)); + if (!PyLong_Check(chunkSizeArg.ptr()) || PyBool_Check(chunkSizeArg.ptr())) { + ThrowStdException("LOB growth chunk size must be a positive integer"); + } + const size_t chunkSize = chunkSizeArg.cast(); + if (chunkSize == 0 || chunkSize > MAX_NATIVE_FETCH_BYTES) { + ThrowStdException("LOB growth chunk size must be between 1 and 256 MiB"); + } std::vector buffer; size_t dataSize = 0; size_t reservedBytes = 0; diff --git a/tests/test_010_pybind_functions.py b/tests/test_010_pybind_functions.py index ce4ec61db..41e175af4 100644 --- a/tests/test_010_pybind_functions.py +++ b/tests/test_010_pybind_functions.py @@ -87,6 +87,21 @@ def test_lob_buffer_growth_is_geometric(): assert growth_count == 7 +@pytest.mark.skipif(not DDBC_AVAILABLE, reason="ddbc_bindings not available") +@pytest.mark.parametrize( + ("chunk_count", "chunk_size", "message"), + [ + (1_000_001, 1, "chunk count"), + (True, 1, "chunk count"), + (1, 0, "chunk size"), + (1, True, "chunk size"), + ], +) +def test_lob_buffer_growth_helper_rejects_unsafe_sizes(chunk_count, chunk_size, message): + with pytest.raises((RuntimeError, TypeError), match=message): + ddbc._test_lob_buffer_growth(chunk_count, chunk_size) + + @pytest.mark.skipif(not DDBC_AVAILABLE, reason="ddbc_bindings not available") class TestPybindModuleInfo: """Test module information and architecture detection.""" From 3bd80c42ea46b9f1b10022f7d95bf4d744f4b7e2 Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Mon, 12 Oct 2026 12:25:59 +0530 Subject: [PATCH 32/39] FIX: address remaining parameter sizing reviews Normalize time values before DAE sizing and scan automatic executemany text columns by UTF-16 code units. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/cursor.py | 15 +++++++++++++ mssql_python/pybind/param_detect.hpp | 10 ++++----- tests/test_023_execute_path_parity.py | 31 +++++++++++++++++++++++++++ 3 files changed, 51 insertions(+), 5 deletions(-) diff --git a/mssql_python/cursor.py b/mssql_python/cursor.py index bafc19de6..c814194a4 100644 --- a/mssql_python/cursor.py +++ b/mssql_python/cursor.py @@ -2777,6 +2777,21 @@ def executemany( # pylint: disable=too-many-locals,too-many-branches,too-many-s ): paraminfo.columnSize = max_decimal_len + if paraminfo.paramSQLType in ( + ddbc_sql_const.SQL_CHAR.value, + ddbc_sql_const.SQL_VARCHAR.value, + ddbc_sql_const.SQL_LONGVARCHAR.value, + ddbc_sql_const.SQL_WCHAR.value, + ddbc_sql_const.SQL_WVARCHAR.value, + ddbc_sql_const.SQL_WLONGVARCHAR.value, + ) and any( + isinstance(value, str) + and sum(2 if ord(char) > 0xFFFF else 1 for char in value) > MAX_INLINE_CHAR + for value in column + ): + paraminfo.isDAE = True + paraminfo.columnSize = 0 + # Special handling for binary data in auto-detected types if paraminfo.paramSQLType in ( ddbc_sql_const.SQL_BINARY.value, diff --git a/mssql_python/pybind/param_detect.hpp b/mssql_python/pybind/param_detect.hpp index 6443d3d23..247b45114 100644 --- a/mssql_python/pybind/param_detect.hpp +++ b/mssql_python/pybind/param_detect.hpp @@ -410,6 +410,11 @@ inline void ApplyInputSizeOverride(PyObject* params, PyObject* inputSize, Py_ssi throw py::type_error("bytes values cannot be bound as SQL_C_WCHAR"); } + if (PyTime_Check(obj) && info.paramCType == PARAM_C_TYPE_TEXT) { + NormalizeTimeParam(params, index, info.columnSize); + obj = PyList_GET_ITEM(params, index); + } + Py_ssize_t actualTextLength = 0; if (PyUnicode_Check(obj)) { actualTextLength = info.paramCType == SQL_C_CHAR @@ -436,11 +441,6 @@ inline void ApplyInputSizeOverride(PyObject* params, PyObject* inputSize, Py_ssi : std::max(info.columnSize, actualSize); } - if (PyTime_Check(obj) && info.paramCType == PARAM_C_TYPE_TEXT) { - NormalizeTimeParam(params, index, info.columnSize); - obj = PyList_GET_ITEM(params, index); - } - if (info.isDAE) { info.dataPtr = borrow(obj); } diff --git a/tests/test_023_execute_path_parity.py b/tests/test_023_execute_path_parity.py index daeb42514..088ce5372 100644 --- a/tests/test_023_execute_path_parity.py +++ b/tests/test_023_execute_path_parity.py @@ -296,6 +296,24 @@ def test_setinputsizes_text_binding_normalizes_time(cursor): ) +def test_setinputsizes_measures_normalized_time_for_dae(cursor): + """A custom time string is measured after normalization and streamed when oversized.""" + + class LargeTime(datetime.time): + def isoformat(self, *args, **kwargs): + return "x" * 4001 + + assert ( + _override_roundtrip( + cursor, + LargeTime(1, 2, 3), + ddbc_sql_const.SQL_VARCHAR.value, + 1, + ) + == "x" * 4001 + ) + + @pytest.mark.parametrize( ("value", "sql_type"), [ @@ -344,6 +362,19 @@ def test_executemany_setinputsizes_scans_mixed_text_values_for_dae(cursor): assert [tuple(row) for row in cursor.fetchall()] == [(value,), ("x",)] +def test_executemany_auto_detection_scans_utf16_lengths_for_dae(cursor): + """A shorter astral string cannot hide behind a longer BMP sample.""" + bmp_value = "x" * 4000 + astral_value = "\U0001f600" * 3000 + cursor.execute("CREATE TABLE #auto_utf16_dae (value NVARCHAR(MAX))") + cursor.executemany( + "INSERT INTO #auto_utf16_dae (value) VALUES (?)", + [(bmp_value,), (astral_value,)], + ) + cursor.execute("SELECT value FROM #auto_utf16_dae ORDER BY DATALENGTH(value) DESC") + assert [tuple(row) for row in cursor.fetchall()] == [(astral_value,), (bmp_value,)] + + def test_setinputsizes_rejects_bytes_for_wide_character_binding(cursor): """Wide-character overrides require text rather than raw encoded bytes.""" cursor.setinputsizes([(ddbc_sql_const.SQL_WVARCHAR.value, 1, 0)]) From 97100e0e3bc0856fbc4484c1e8ba2e033c4b891c Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Mon, 12 Oct 2026 12:31:30 +0530 Subject: [PATCH 33/39] FIX: initialize executemany DAE token Avoid formatting an indeterminate SQLParamData token when the driver reports completion without writing it. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/pybind/ddbc_bindings.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index b89ca350c..acc0c4c1e 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -3369,7 +3369,7 @@ SQLRETURN SQLExecuteMany_wrap(const SqlHandlePtr statementHandle, const std::u16 LOG("SQLExecuteMany: SQLExecute for row %zu - initial_rc=%d", rowIndex, rc); size_t dae_chunk_count = 0; while (rc == SQL_NEED_DATA) { - SQLPOINTER token; + SQLPOINTER token = nullptr; { // Release the GIL around the blocking SQLParamData call. py::gil_scoped_release release; From f63e905d7eb3cdaeff2eea73b81633fb0cd836b1 Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Mon, 12 Oct 2026 12:39:08 +0530 Subject: [PATCH 34/39] FIX: complete DAE rows and bound numeric text Continue SQLParamData after each streamed executemany parameter and reject oversized numeric override text before native copying. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/pybind/ddbc_bindings.cpp | 5 +++-- mssql_python/pybind/param_detect.hpp | 3 +++ tests/test_023_execute_path_parity.py | 9 +++++++++ 3 files changed, 15 insertions(+), 2 deletions(-) diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index acc0c4c1e..301f80dd9 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -3393,8 +3393,9 @@ SQLRETURN SQLExecuteMany_wrap(const SqlHandlePtr statementHandle, const std::u16 ThrowStdException( "SQLExecuteMany: unrecognized paramToken from SQLParamData"); } - rc = StreamDAEParameter(hStmt, *matchedInfo, charEncoding); - if (!SQL_SUCCEEDED(rc)) return rc; + const SQLRETURN streamRc = + StreamDAEParameter(hStmt, *matchedInfo, charEncoding); + if (!SQL_SUCCEEDED(streamRc)) return streamRc; dae_chunk_count++; } LOG("SQLExecuteMany: DAE completed for row %zu - total_chunks=%zu, " diff --git a/mssql_python/pybind/param_detect.hpp b/mssql_python/pybind/param_detect.hpp index 247b45114..7f71125ec 100644 --- a/mssql_python/pybind/param_detect.hpp +++ b/mssql_python/pybind/param_detect.hpp @@ -421,6 +421,9 @@ inline void ApplyInputSizeOverride(PyObject* params, PyObject* inputSize, Py_ssi ? EncodedUnicodeLength(obj, charEncoding) : UnicodeUtf16Length(obj); } + if (isNumeric && PyUnicode_Check(obj) && actualTextLength > MAX_INLINE_CHAR) { + throw py::value_error("numeric parameter text is too large"); + } const bool textNeedsDAE = !isNumeric && PyUnicode_Check(obj) && actualTextLength > MAX_INLINE_CHAR; const bool binaryNeedsDAE = diff --git a/tests/test_023_execute_path_parity.py b/tests/test_023_execute_path_parity.py index 088ce5372..d0e37b281 100644 --- a/tests/test_023_execute_path_parity.py +++ b/tests/test_023_execute_path_parity.py @@ -416,6 +416,15 @@ def __format__(self, format_spec): cursor.setinputsizes(None) +def test_setinputsizes_rejects_oversized_numeric_text(cursor): + cursor.setinputsizes([(ddbc_sql_const.SQL_DECIMAL.value, 18, 0)]) + try: + with pytest.raises(ValueError, match="numeric parameter text is too large"): + cursor.execute("SELECT ?", [decimal.Decimal("1e100000")]) + finally: + cursor.setinputsizes(None) + + def test_setinputsizes_binary_dae(cursor): """Declared binary sizes over 8000 stream through the native DAE path.""" value = b"\xab" * 10000 From 73ab6f89514df6e4358e13a24f8c91744815bcc9 Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Mon, 12 Oct 2026 12:48:46 +0530 Subject: [PATCH 35/39] FIX: use bounded Unicode encoding paths Encode string subclasses from their immutable Unicode contents so overridden encode methods cannot bypass native sizing. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/pybind/ddbc_bindings.cpp | 13 ++++++++++--- tests/test_023_execute_path_parity.py | 14 ++++++++++++++ 2 files changed, 24 insertions(+), 3 deletions(-) diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index 301f80dd9..51fe81414 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -613,6 +613,12 @@ static SQLRETURN stream_unicode_dae_chunks(PyObject* value, const std::string& e return SQL_SUCCESS; } +static py::object EncodeUnicodeValue(PyObject* value, const std::string& encoding) { + PyObject* encoded = PyUnicode_AsEncodedString(value, encoding.c_str(), "strict"); + if (!encoded) throw py::error_already_set(); + return steal(encoded); +} + static SQLRETURN StreamDAEParameter(SQLHSTMT hStmt, const ParamInfo& info, const std::string& charEncoding) { PyObject* value = info.dataPtr.ptr(); @@ -850,7 +856,8 @@ SQLRETURN BindParameters(SqlHandle& handle, SQLHANDLE hStmt, const py::list& par if (py::isinstance(param)) { // Encode Unicode string using the specified encoding try { - py::object encoded = param.attr("encode")(charEncoding, "strict"); + py::object encoded = + EncodeUnicodeValue(param.ptr(), charEncoding); encodedStr = encoded.cast(); LOG("BindParameters: param[%d] SQL_C_CHAR - Encoded with '%s', " "size=%zu bytes", @@ -2654,7 +2661,7 @@ SQLRETURN BindParameterArray(SqlHandle& handle, SQLHANDLE hStmt, const py::list& std::to_string(paramIndex)); } py::object encoded = - columnValues[i].attr("encode")("utf-16-le", "strict"); + EncodeUnicodeValue(columnValues[i].ptr(), "utf-16-le"); char* encodedData = nullptr; Py_ssize_t encodedSize = 0; if (PyBytes_AsStringAndSize(encoded.ptr(), &encodedData, @@ -2785,7 +2792,7 @@ SQLRETURN BindParameterArray(SqlHandle& handle, SQLHANDLE hStmt, const py::list& // encoding try { encoded = - columnValues[i].attr("encode")(charEncoding, "strict"); + EncodeUnicodeValue(columnValues[i].ptr(), charEncoding); if (PyBytes_AsStringAndSize(encoded.ptr(), &encodedData, &encodedSize) != 0) { throw py::error_already_set(); diff --git a/tests/test_023_execute_path_parity.py b/tests/test_023_execute_path_parity.py index d0e37b281..516e9eeb4 100644 --- a/tests/test_023_execute_path_parity.py +++ b/tests/test_023_execute_path_parity.py @@ -375,6 +375,20 @@ def test_executemany_auto_detection_scans_utf16_lengths_for_dae(cursor): assert [tuple(row) for row in cursor.fetchall()] == [(astral_value,), (bmp_value,)] +def test_executemany_ignores_str_subclass_encode_override(cursor): + class HostileStr(str): + def encode(self, *args, **kwargs): + return b"x" * 100_000 + + cursor.execute("CREATE TABLE #str_subclass_encoding (value NVARCHAR(20))") + cursor.executemany( + "INSERT INTO #str_subclass_encoding (value) VALUES (?)", + [(HostileStr("safe"),), (HostileStr("text"),)], + ) + cursor.execute("SELECT value FROM #str_subclass_encoding ORDER BY value") + assert [tuple(row) for row in cursor.fetchall()] == [("safe",), ("text",)] + + def test_setinputsizes_rejects_bytes_for_wide_character_binding(cursor): """Wide-character overrides require text rather than raw encoded bytes.""" cursor.setinputsizes([(ddbc_sql_const.SQL_WVARCHAR.value, 1, 0)]) From 8fdde26aa030e2c08b721c1baf33ea49f06e8374 Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Mon, 12 Oct 2026 15:03:31 +0530 Subject: [PATCH 36/39] FIX: size base Unicode contents in batches Bypass string subclass protocol overrides during bounded narrow and UTF-16 executemany sizing. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql_python/cursor.py | 21 +++++++++++++-------- tests/test_023_execute_path_parity.py | 18 ++++++++++++++---- 2 files changed, 27 insertions(+), 12 deletions(-) diff --git a/mssql_python/cursor.py b/mssql_python/cursor.py index c814194a4..ef0065de5 100644 --- a/mssql_python/cursor.py +++ b/mssql_python/cursor.py @@ -62,16 +62,23 @@ def _encoded_length_exceeds(value: str, encoding: str, limit: int) -> bool: encoder = codecs.getincrementalencoder(encoding)(errors="strict") total = 0 chunk_size = 4096 - for offset in range(0, len(value), chunk_size): - end = min(offset + chunk_size, len(value)) - total += len(encoder.encode(value[offset:end], final=end == len(value))) + value_length = str.__len__(value) + for offset in range(0, value_length, chunk_size): + end = min(offset + chunk_size, value_length) + chunk = str.__getitem__(value, slice(offset, end)) + total += len(encoder.encode(chunk, final=end == value_length)) if total > limit: return True - if not value: + if value_length == 0: total += len(encoder.encode("", final=True)) return total > limit +def _utf16_length_exceeds(value: str, limit: int) -> bool: + bounded_value = str.__getitem__(value, slice(0, limit + 1)) + return len(str.encode(bounded_value, "utf-16-le")) // 2 > limit + + # SQL BIGINT is a signed 64-bit integer. Ints outside this range have no BIGINT # encoding and must be rejected at detect time on both paths (see _map_sql_type). BIGINT_MIN: int = -(2**63) @@ -2664,8 +2671,7 @@ def executemany( # pylint: disable=too-many-locals,too-many-branches,too-many-s ) else: text_is_large = any( - sum(2 if ord(char) > 0xFFFF else 1 for char in value) - > MAX_INLINE_CHAR + _utf16_length_exceeds(value, MAX_INLINE_CHAR) for value in text_values ) binary_is_large = narrow_text_type and any( @@ -2785,8 +2791,7 @@ def executemany( # pylint: disable=too-many-locals,too-many-branches,too-many-s ddbc_sql_const.SQL_WVARCHAR.value, ddbc_sql_const.SQL_WLONGVARCHAR.value, ) and any( - isinstance(value, str) - and sum(2 if ord(char) > 0xFFFF else 1 for char in value) > MAX_INLINE_CHAR + isinstance(value, str) and _utf16_length_exceeds(value, MAX_INLINE_CHAR) for value in column ): paraminfo.isDAE = True diff --git a/tests/test_023_execute_path_parity.py b/tests/test_023_execute_path_parity.py index 516e9eeb4..4833de9d3 100644 --- a/tests/test_023_execute_path_parity.py +++ b/tests/test_023_execute_path_parity.py @@ -377,16 +377,26 @@ def test_executemany_auto_detection_scans_utf16_lengths_for_dae(cursor): def test_executemany_ignores_str_subclass_encode_override(cursor): class HostileStr(str): + def __len__(self): + return 1 + + def __iter__(self): + return iter("x") + + def __getitem__(self, key): + return "x" + def encode(self, *args, **kwargs): return b"x" * 100_000 - cursor.execute("CREATE TABLE #str_subclass_encoding (value NVARCHAR(20))") + large_value = HostileStr("\U0001f600" * 3000) + cursor.execute("CREATE TABLE #str_subclass_encoding (value NVARCHAR(MAX))") cursor.executemany( "INSERT INTO #str_subclass_encoding (value) VALUES (?)", - [(HostileStr("safe"),), (HostileStr("text"),)], + [(HostileStr("safe"),), (large_value,)], ) - cursor.execute("SELECT value FROM #str_subclass_encoding ORDER BY value") - assert [tuple(row) for row in cursor.fetchall()] == [("safe",), ("text",)] + cursor.execute("SELECT value FROM #str_subclass_encoding ORDER BY DATALENGTH(value)") + assert [tuple(row) for row in cursor.fetchall()] == [("safe",), (str(large_value),)] def test_setinputsizes_rejects_bytes_for_wide_character_binding(cursor): From 748ceb2da655a1def9211787cd2fac4ff83d7487 Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Mon, 19 Oct 2026 13:44:48 +0530 Subject: [PATCH 37/39] FIX: Address native size validation review feedback Reject oversized formatted numeric batch parameters before native binding. Cover the rejection boundary and the existing wide automatic text binding, and document row caps, allocation budgets, and compatibility changes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- CHANGELOG.md | 15 +++++ README.md | 46 ++++++++++++++ mssql_python/cursor.py | 46 ++++++++++++-- tests/test_023_execute_path_parity.py | 89 +++++++++++++++++++++++++++ 4 files changed, 190 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2cac01d4d..f9f2eb4d2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -57,6 +57,21 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), does not change the default provider or ship any Rust driver binaries. ### Changed +- Cursor fetch/Arrow batch sizes and bulk-copy batch/commit sizes are now capped + at 1,000,000 rows; `arraysize` must be positive. Negative row counts raise + `ValueError` and boolean row counts raise `TypeError`. Zero-size fetches remain + empty, zero-size Arrow requests are schema-only, and bulk-copy zero retains + native default batching. Native array-bound `executemany()` calls are also + capped at 1,000,000 parameter sets. +- Native fetch buffers and parameter-array buffers each have a cumulative + 256 MiB allocation budget. Array limits depend on batch size and column widths, + including slots for NULLs; reduce batch sizes or oversized `setinputsizes()` + hints when a budget is exceeded. Non-numeric hints are capped at 268,435,456 + size units, while numeric precision/scale normalization is preserved. + Negative/boolean size hints are rejected. These are native buffer limits, not + process-memory or total-result limits; see the README's cursor-size guidance. +- Oversized formatted DECIMAL/NUMERIC text is rejected with `ValueError` before + native binding on both scalar and array-bound execution paths. - Fetches reuse owned native metadata for stable columns within a result set; `fetchmany()` avoids the Python metadata-dictionary roundtrip. Re-execution, result transitions and statement/connection cleanup invalidate this metadata. diff --git a/README.md b/README.md index 40fc409e7..ab1387325 100644 --- a/README.md +++ b/README.md @@ -170,6 +170,52 @@ example, requesting ASCII with `SQL_CHAR` does not cause non-ASCII parameters to raise encoding errors. Use `setencoding()` with no arguments to restore the supported defaults. `setdecoding()` independently controls how result data is read. +### Cursor size and native buffer limits + +Row counts must be integers, not booleans. The following limits are inclusive: + +| API setting | Accepted row count | Meaning of zero | +| --- | --- | --- | +| `Cursor.arraysize` | 1 to 1,000,000 | Not accepted | +| `Cursor.fetchmany(size)` | 0 to 1,000,000; `None` uses `arraysize` | Returns an empty list without fetching | +| `Cursor.arrow_batch(batch_size)`, `arrow(batch_size)`, `arrow_reader(batch_size)` | 0 to 1,000,000 | Schema only; no data rows | +| `Cursor.bulkcopy(batch_size)`, `bulkcopy_arrow(batch_size)` | 0 to 1,000,000 | Uses the native core's default batch/commit behavior | + +Native array-bound `executemany()` calls also accept at most 1,000,000 parameter +sets per call; an empty input remains a no-op. These are per-batch limits, not +limits on the total number of rows in a result or table. Bulk-copy `batch_size` +controls rows per batch/commit, not the total source size. + +Negative row counts now raise `ValueError`, rather than acting like empty or +unbounded requests. Boolean and other non-integer row counts raise `TypeError`. +In particular, negative Arrow sizes no longer produce an empty result. + +Native fetch buffers have a cumulative **256 MiB** allocation budget, including +Arrow data buffers and native rowset storage. Native parameter-array buffers have +a separate cumulative **256 MiB** budget across all columns in one array-bound +`executemany()` call. A row count within the allowed range can still exceed its +byte budget when columns are wide. Parameter-array storage grows with the number +of parameter sets times the sum of each column's native element width and length +indicator; text buffers also include terminators, and wide text uses two-byte +UTF-16 code units. NULL rows still reserve array slots. + +For allocation-limit errors, reduce `fetchmany()`/Arrow batch sizes or split +`executemany()` inputs into smaller calls, and avoid unnecessarily large +`setinputsizes()` hints. Non-numeric size hints must be non-negative integers no +greater than 268,435,456 in the SQL type's size units; this does not guarantee that +the corresponding byte allocation fits. Numeric hints describe precision/scale +instead of buffer length and retain their normalization to supported SQL Server +precision (up to 38). Boolean or negative size/scale hints raise `ValueError`. +Formatted DECIMAL/NUMERIC parameter text exceeding 4000 UTF-16 code units is +rejected before native binding. + +These budgets bound the native buffers checked by the driver, **not total process +memory or total result memory**. Python objects, driver/core allocations, input +collections, and retained result batches consume additional memory. `fetchall()` +and `arrow()` still materialize all remaining rows; process bounded +`fetchmany()` batches or iterate `arrow_reader()` without retaining every batch +when the complete result may not fit in memory. + ## Getting Started Examples Connect to SQL Server and execute a simple query: diff --git a/mssql_python/cursor.py b/mssql_python/cursor.py index 8e9f75ffe..c0ba8320e 100644 --- a/mssql_python/cursor.py +++ b/mssql_python/cursor.py @@ -1209,6 +1209,7 @@ def _validate_native_row_count(value: int, name: str, allow_zero: bool) -> int: @property def arraysize(self) -> int: + """Default fetchmany row count: an integer from 1 to 1,000,000, not bool.""" return self._arraysize @arraysize.setter @@ -1253,6 +1254,13 @@ def setinputsizes(self, sizes: List[Union[int, tuple]]) -> None: Args: sizes: A sequence of tuples, one for each parameter. Each tuple contains (sql_type, size, decimal_digits) where size and decimal_digits are optional. + Non-numeric sizes must be integers from 0 to 268,435,456 in the SQL + type's size units; booleans and negative sizes/scales are rejected. + Numeric precision/scale hints retain normalization to SQL Server's + supported precision (up to 38), rather than being byte lengths. + Array-bound executemany buffers across all columns must fit within + 256 MiB, including NULL slots and length indicators. Reduce batch + sizes or excessive size hints if this native allocation limit is hit. """ # Get valid SQL types from centralized constants @@ -2529,8 +2537,13 @@ def executemany( # pylint: disable=too-many-locals,too-many-branches,too-many-s Args: operation: SQL query or command. seq_of_parameters: Sequence of sequences or mappings of parameters. + Native array binding supports at most 1,000,000 parameter sets and + 256 MiB of parameter buffers per call, depending on column widths + and length indicators (including NULL slots). Split larger inputs + into smaller calls. These limits do not bound total process memory. Raises: Error: If the operation fails. + ValueError: If formatted DECIMAL/NUMERIC text exceeds 4000 UTF-16 code units. """ logger.debug( "executemany: Starting - operation_length=%d, batch_count=%d", @@ -2882,6 +2895,8 @@ def executemany( # pylint: disable=too-many-locals,too-many-branches,too-many-s conversion_failed = True if conversion_failed: raise ValueError(err_msg) from None + if _utf16_length_exceeds(processed_row[i], MAX_INLINE_CHAR): + raise ValueError("numeric parameter text is too large") processed_parameters.append(processed_row) # Now transpose the processed parameters @@ -3023,7 +3038,13 @@ def fetchmany(self, size: Optional[int] = None) -> List[Row]: Fetch the next set of rows of a query result. Args: - size: Number of rows to fetch at a time. + size: Integer from 0 to 1,000,000, not bool. None uses arraysize. + Zero returns an empty list without fetching. Native fetch buffers + must also fit within 256 MiB; reduce size for wide columns. + + Raises: + TypeError: If size is a boolean or another non-integer. + ValueError: If size is negative or exceeds 1,000,000. Returns: List of Row objects. @@ -3184,7 +3205,10 @@ def arrow_batch(self, batch_size: int = 8192) -> "pyarrow.RecordBatch": query result set. Args: - batch_size: Maximum number of rows to fetch in the Record Batch. + batch_size: Maximum rows, an integer from 0 to 1,000,000, not bool. + Zero returns schema only, without advancing the result set. + Native fetch buffers must fit within 256 MiB; reduce batch_size + for wide columns. This does not bound total process memory. Returns: A pyarrow RecordBatch object containing up to batch_size rows. @@ -3225,7 +3249,10 @@ def arrow(self, batch_size: int = 8192) -> "pyarrow.Table": Fetch the entire result as a pyarrow Table. Args: - batch_size: Size of the Record Batches which make up the Table. + batch_size: Rows per Record Batch, from 0 to 1,000,000, not bool. + Zero returns schema only. Each batch is subject to the native + fetch buffer budget described by arrow_batch; the complete + Table's memory consumption is not bounded by that budget. Returns: A pyarrow Table containing all remaining rows from the result set. @@ -3264,7 +3291,10 @@ def arrow_reader(self, batch_size: int = 8192) -> "_ArrowReader": is **not** closed and can be re-executed. ``close()`` is idempotent. Args: - batch_size: Size of the Record Batches produced by the reader. + batch_size: Rows per Record Batch, from 0 to 1,000,000, not bool. + Zero yields no data batches. Validation occurs when reading. + Each batch uses the native fetch budget described by arrow_batch; + retaining batches can still grow total process memory. Returns: A pyarrow-compatible RecordBatchReader for the result set. @@ -3654,7 +3684,9 @@ def bulkcopy( - The number of values in each row must match the number of columns in the target table - batch_size: Number of rows to send per batch. Default 0 uses server optimal. + batch_size: Rows per batch/commit, an integer from 0 to 1,000,000, + not bool. Default 0 uses native core defaults. This does not + limit the total number of input rows or total process memory. timeout: Operation timeout in seconds. Default is 30. 0 disables the bulk copy operation timeout. @@ -3843,7 +3875,9 @@ def bulkcopy_arrow( Args: table_name: Target table name (may include schema, e.g. 'dbo.MyTable'). source: Arrow source (see above). - batch_size: Rows per TDS commit. Default 0 uses server optimal. + batch_size: Rows per TDS commit, an integer from 0 to 1,000,000, + not bool. Default 0 uses native core defaults. This does not + limit the total source rows or total process memory. timeout: Operation timeout in seconds. Default 30. 0 disables the bulk copy operation timeout. column_mappings: Same two formats as :meth:`bulkcopy`. When omitted, diff --git a/tests/test_023_execute_path_parity.py b/tests/test_023_execute_path_parity.py index 4833de9d3..6eae78c7d 100644 --- a/tests/test_023_execute_path_parity.py +++ b/tests/test_023_execute_path_parity.py @@ -375,6 +375,30 @@ def test_executemany_auto_detection_scans_utf16_lengths_for_dae(cursor): assert [tuple(row) for row in cursor.fetchall()] == [(astral_value,), (bmp_value,)] +def test_executemany_auto_text_uses_wide_array_binding(cursor, monkeypatch): + """UTF-8 byte length does not govern automatically detected wide parameters.""" + from mssql_python import ddbc_bindings + + values = ["x" * 4000, "\u0800" * 2000] + expected = [] + for value in values: + cursor.execute("SELECT CONVERT(NVARCHAR(MAX), CONVERT(VARCHAR(MAX), ?))", [value]) + expected.append(cursor.fetchone()[0]) + cursor.execute("CREATE TABLE #auto_wide_array (value NVARCHAR(MAX))") + array_calls = [] + execute_many = ddbc_bindings.SQLExecuteMany + + def capture_array_binding(*args): + array_calls.append((args[3][0].paramCType, args[3][0].columnSize)) + return execute_many(*args) + + monkeypatch.setattr(ddbc_bindings, "SQLExecuteMany", capture_array_binding) + cursor.executemany("INSERT INTO #auto_wide_array (value) VALUES (?)", [(v,) for v in values]) + assert array_calls == [(ddbc_sql_const.SQL_C_WCHAR.value, 4000)] + cursor.execute("SELECT value FROM #auto_wide_array ORDER BY DATALENGTH(value) DESC") + assert [row[0] for row in cursor.fetchall()] == expected + + def test_executemany_ignores_str_subclass_encode_override(cursor): class HostileStr(str): def __len__(self): @@ -449,6 +473,71 @@ def test_setinputsizes_rejects_oversized_numeric_text(cursor): cursor.setinputsizes(None) +@pytest.mark.parametrize( + "sql_type", [ddbc_sql_const.SQL_DECIMAL.value, ddbc_sql_const.SQL_NUMERIC.value] +) +@pytest.mark.parametrize( + "value", + [ + decimal.Decimal("1e100000"), + decimal.Decimal("1e-100000"), + "1e100000", + "1e-100000", + "9" * 4001, + decimal.Decimal("0." + "0" * 3998 + "1"), + ], + ids=[ + "decimal-positive-exponent", + "decimal-negative-exponent", + "text-positive-exponent", + "text-negative-exponent", + "long-text", + "exactly-4001-units", + ], +) +def test_executemany_setinputsizes_rejects_oversized_numeric_text( + cursor, monkeypatch, sql_type, value +): + """Reject formatted numeric text before the native array binder encodes it.""" + from mssql_python import ddbc_bindings + + def unexpected_array_binding(*args): + pytest.fail("Oversized numeric text reached the native array binder") + + monkeypatch.setattr(ddbc_bindings, "SQLExecuteMany", unexpected_array_binding) + cursor.setinputsizes([(sql_type, 18, 0)]) + try: + with pytest.raises(ValueError, match="numeric parameter text is too large"): + cursor.executemany("SELECT ?", [(decimal.Decimal("1"),), (value,)]) + finally: + cursor.setinputsizes(None) + + +@pytest.mark.parametrize("length", [3999, 4000]) +def test_executemany_numeric_text_at_limit_reaches_array_binding(cursor, monkeypatch, length): + from mssql_python import ddbc_bindings + + value = decimal.Decimal("0." + "0" * (length - 3) + "1") + array_calls = [] + + class ReachedArrayBinding(Exception): + pass + + def capture_array_binding(*args): + array_calls.append(args[2]) + raise ReachedArrayBinding + + monkeypatch.setattr(ddbc_bindings, "SQLExecuteMany", capture_array_binding) + cursor.setinputsizes([(ddbc_sql_const.SQL_DECIMAL.value, 18, 0)]) + try: + with pytest.raises(ReachedArrayBinding): + cursor.executemany("SELECT ?", [(value,)]) + assert array_calls == [[[format(value, "f")]]] + assert len(array_calls[0][0][0]) == length + finally: + cursor.setinputsizes(None) + + def test_setinputsizes_binary_dae(cursor): """Declared binary sizes over 8000 stream through the native DAE path.""" value = b"\xab" * 10000 From e3449ca3f2a8b75a837921b14fd84201896429bd Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Tue, 20 Oct 2026 14:28:48 +0530 Subject: [PATCH 38/39] FIX: Isolate native parameter rejection tests Use dedicated context-managed cursors so expected binding exceptions cannot leave input-size hints on the shared module cursor. Verified the reported binary-rejection then numeric-test ordering: eight cases pass. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- tests/test_004_cursor.py | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/tests/test_004_cursor.py b/tests/test_004_cursor.py index cf8574d00..313f5eac8 100644 --- a/tests/test_004_cursor.py +++ b/tests/test_004_cursor.py @@ -1727,16 +1727,20 @@ def test_setinputsizes_rejects_excessive_column_size(): cursor.setinputsizes([(mssql_python.SQL_VARCHAR, 1 << 40, 0)]) -def test_executemany_rejects_excessive_cumulative_parameter_buffers(cursor): - cursor.setinputsizes([(mssql_python.SQL_WVARCHAR, 50_000_000, 0)] * 3) - with pytest.raises(RuntimeError, match="Parameter buffers exceed the 256 MiB allocation limit"): - cursor.executemany("SELECT ?, ?, ?", [(None, None, None)]) +def test_executemany_rejects_excessive_cumulative_parameter_buffers(db_connection): + with db_connection.cursor() as cursor: + cursor.setinputsizes([(mssql_python.SQL_WVARCHAR, 50_000_000, 0)] * 3) + with pytest.raises( + RuntimeError, match="Parameter buffers exceed the 256 MiB allocation limit" + ): + cursor.executemany("SELECT ?, ?, ?", [(None, None, None)]) -def test_executemany_rejects_text_for_binary_parameter(cursor): - cursor.setinputsizes([(mssql_python.SQL_VARBINARY, 10, 0)]) - with pytest.raises(RuntimeError, match="object type does not match"): - cursor.executemany("SELECT CAST(? AS VARBINARY(10))", [("text",)]) +def test_executemany_rejects_text_for_binary_parameter(db_connection): + with db_connection.cursor() as cursor: + cursor.setinputsizes([(mssql_python.SQL_VARBINARY, 10, 0)]) + with pytest.raises(RuntimeError, match="object type does not match"): + cursor.executemany("SELECT CAST(? AS VARBINARY(10))", [("text",)]) def test_fetchmany_rejects_excessive_native_buffer(cursor): From 44017f9e31d9288ee1dde2c318d8d2a0ea0fc811 Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Tue, 20 Oct 2026 16:01:01 +0530 Subject: [PATCH 39/39] FIX: Preserve Arrow rows across native buffer limits Use budget-aware amortized growth and lossless partial Arrow batches with statement-owned pending rows and resumable LOB reads. Preserve pending state across mixed fetch, raw fetch/GetData, scroll/skip, and metadata-only attribute changes; clear it on real result invalidation. Document short-batch semantics and add isolated ordinary/LOB and lifecycle regressions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- CHANGELOG.md | 12 +- README.md | 17 + mssql_python/cursor.py | 31 +- mssql_python/pybind/connection/connection.cpp | 20 +- mssql_python/pybind/connection/connection.h | 2 +- mssql_python/pybind/ddbc_bindings.cpp | 743 ++++++++++-------- mssql_python/pybind/ddbc_bindings.h | 51 ++ tests/test_004_cursor.py | 40 + tests/test_004_cursor_arrow.py | 581 ++++++++++++++ 9 files changed, 1156 insertions(+), 341 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f9f2eb4d2..edac2ed50 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -72,12 +72,22 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), process-memory or total-result limits; see the README's cursor-size guidance. - Oversized formatted DECIMAL/NUMERIC text is rejected with `ValueError` before native binding on both scalar and array-bound execution paths. +- Arrow buffer growth is budget-aware. When a batch reaches the native budget, + completed rows are returned and overflow/unread fetched rows are retained for + subsequent Arrow or row-wise fetches. Short batches no longer imply exhaustion; + continue until empty. `arrow()` and `arrow_reader()` handle this automatically. + A single over-budget row raises without being skipped and may be recovered + row-wise if its source buffers fit. Pending rows are discarded on re-execution, + result-set changes, cancellation, close, connection reset, or transaction/ + autocommit changes, but not metadata-only connection-attribute changes. + Total retained-result memory is still not limited by the per-fetch budget. - Fetches reuse owned native metadata for stable columns within a result set; `fetchmany()` avoids the Python metadata-dictionary roundtrip. Re-execution, result transitions and statement/connection cleanup invalidate this metadata. Public descriptions stay fresh, name-validation timing is preserved, and `sql_variant` columns retain per-row descriptions and per-value probes. - Fetch buffers, decoding settings and converted values are not cached. + Apart from pending rows retained by budget-limited Arrow fetches, fetch buffers, + decoding settings and converted values are not cached. - DATE, TIME, and TIMESTAMP fetch conversion uses checked CPython constructors for the standard datetime types, while preserving cached substitute constructors, their positional arguments and exceptions, and fractional-second truncation. diff --git a/README.md b/README.md index ab1387325..c9b700769 100644 --- a/README.md +++ b/README.md @@ -199,6 +199,23 @@ of parameter sets times the sum of each column's native element width and length indicator; text buffers also include terminators, and wide text uses two-byte UTF-16 code units. NULL rows still reserve array slots. +For Arrow fetches, `batch_size` is an upper bound, not a guaranteed batch length. +When the next row would exceed the native budget, the driver returns completed +rows and retains the overflow row and any unread fetched rows for the next fetch. +A short batch does **not** indicate end of results: continue until an empty batch. +`arrow()` and `arrow_reader()` do this automatically. Row-wise fetch methods can +also consume these pending rows; `fetchmany()` may return a short batch while +draining them, so continue until an empty list. + +If a single row cannot coexist with its Arrow output within the budget, the +driver raises an allocation-limit error and retains the row. It can be recovered +with a row-wise fetch if its source buffers fit. If the source buffers themselves +exceed the budget, retries continue to fail rather than skipping the row. +Executing another statement, moving to the next result set, cancelling, or closing +abandons pending rows. Connection resets, transaction completion, and autocommit +changes invalidate them with the result state. Non-resetting connection attributes +preserve pending rows, and schema-only Arrow requests do not consume them. + For allocation-limit errors, reduce `fetchmany()`/Arrow batch sizes or split `executemany()` inputs into smaller calls, and avoid unnecessarily large `setinputsizes()` hints. Non-numeric size hints must be non-negative integers no diff --git a/mssql_python/cursor.py b/mssql_python/cursor.py index c0ba8320e..75dc24f15 100644 --- a/mssql_python/cursor.py +++ b/mssql_python/cursor.py @@ -3041,6 +3041,8 @@ def fetchmany(self, size: Optional[int] = None) -> List[Row]: size: Integer from 0 to 1,000,000, not bool. None uses arraysize. Zero returns an empty list without fetching. Native fetch buffers must also fit within 256 MiB; reduce size for wide columns. + When recovering pending Arrow rows, a call may return fewer than + size rows before the result is exhausted. Continue until empty. Raises: TypeError: If size is a boolean or another non-integer. @@ -3201,14 +3203,17 @@ def fetchall(self) -> List[Row]: def arrow_batch(self, batch_size: int = 8192) -> "pyarrow.RecordBatch": """ - Fetch a single pyarrow Record Batch of the specified size from the + Fetch a single pyarrow Record Batch of up to the specified size from the query result set. Args: batch_size: Maximum rows, an integer from 0 to 1,000,000, not bool. Zero returns schema only, without advancing the result set. - Native fetch buffers must fit within 256 MiB; reduce batch_size - for wide columns. This does not bound total process memory. + Native fetch buffers must fit within 256 MiB. A batch may be + shorter to stay within this budget; subsequent fetches resume + without losing rows. This does not bound total process memory. + If one row cannot fit, an error leaves it pending for recovery + by a row-wise fetch when its source buffers fit the budget. Returns: A pyarrow RecordBatch object containing up to batch_size rows. @@ -3249,7 +3254,7 @@ def arrow(self, batch_size: int = 8192) -> "pyarrow.Table": Fetch the entire result as a pyarrow Table. Args: - batch_size: Rows per Record Batch, from 0 to 1,000,000, not bool. + batch_size: Maximum rows per Record Batch, from 0 to 1,000,000, not bool. Zero returns schema only. Each batch is subject to the native fetch buffer budget described by arrow_batch; the complete Table's memory consumption is not bounded by that budget. @@ -3263,7 +3268,7 @@ def arrow(self, batch_size: int = 8192) -> "pyarrow.Table": batches: list["pyarrow.RecordBatch"] = [] while True: batch = self.arrow_batch(batch_size) - if batch.num_rows < batch_size or batch_size <= 0: + if batch.num_rows == 0 or batch_size <= 0: if not batches or batch.num_rows > 0: batches.append(batch) break @@ -3273,7 +3278,7 @@ def arrow(self, batch_size: int = 8192) -> "pyarrow.Table": def arrow_reader(self, batch_size: int = 8192) -> "_ArrowReader": """ Fetch the result as a pyarrow-compatible RecordBatchReader, which - yields Record Batches of the specified size until the current result + yields Record Batches of up to the specified size until the current result set is exhausted. The returned object is an ``_ArrowReader`` wrapper that behaves like @@ -3291,7 +3296,7 @@ def arrow_reader(self, batch_size: int = 8192) -> "_ArrowReader": is **not** closed and can be re-executed. ``close()`` is idempotent. Args: - batch_size: Rows per Record Batch, from 0 to 1,000,000, not bool. + batch_size: Maximum rows per Record Batch, from 0 to 1,000,000, not bool. Zero yields no data batches. Validation occurs when reading. Each batch uses the native fetch budget described by arrow_batch; retaining batches can still grow total process memory. @@ -4111,7 +4116,7 @@ def scroll( self, value: int, mode: str = "relative" ) -> None: # pylint: disable=too-many-branches """ - Scroll using SQLFetchScroll only, matching test semantics: + Advance a forward-only cursor without materializing skipped rows: - relative(N>0): consume N rows; rownumber = previous + N; next fetch returns the following row. - absolute(-1): before first (rownumber = -1), no data consumed. @@ -4158,8 +4163,6 @@ def scroll( f"Cannot move backward by {value} rows on a forward-only cursor", ) - row_data: list = [] - # Absolute positioning not supported with forward-only cursors if mode == "absolute": raise NotSupportedError( @@ -4172,12 +4175,10 @@ def scroll( if value == 0: return - # For forward-only cursors, use multiple SQL_FETCH_NEXT calls - # This matches pyodbc's approach for skip operations + # Native fetch also consumes rows cached by a quota-limited Arrow batch. for i in range(value): - ret = ddbc_bindings.DDBCSQLFetchScroll( - self.hstmt, ddbc_sql_const.SQL_FETCH_NEXT.value, 0, row_data - ) + ret = ddbc_bindings.DDBCSQLFetch(self.hstmt) + check_error(ddbc_sql_const.SQL_HANDLE_STMT.value, self.hstmt, ret) if ret == ddbc_sql_const.SQL_NO_DATA.value: raise IndexError( "Cannot scroll to specified position: end of result set reached" diff --git a/mssql_python/pybind/connection/connection.cpp b/mssql_python/pybind/connection/connection.cpp index 9c4304d2f..9ebf9eac4 100644 --- a/mssql_python/pybind/connection/connection.cpp +++ b/mssql_python/pybind/connection/connection.cpp @@ -115,7 +115,7 @@ void Connection::connect(const py::dict& attrs_before) { void Connection::disconnect(bool rollbackBeforeDisconnect) { PERF_TIMER("Connection::disconnect"); - clearResultMetadata(); + clearResultMetadata(true); // Determine GIL state once, up front. disconnect() runs both from // pybind11-bound methods (GIL held) and from GIL-less destructor / shutdown // paths: Connection::~Connection() dropping the last shared_ptr, or teardown @@ -266,7 +266,7 @@ void Connection::checkError(SQLRETURN ret) const { } } -void Connection::clearResultMetadata() { +void Connection::clearResultMetadata(bool discardPendingRows) { std::vector handles; { std::lock_guard lock(_childHandlesMutex); @@ -280,7 +280,11 @@ void Connection::clearResultMetadata() { // Releasing the last handle can acquire the connection cleanup gate. // Keep that destruction outside the child-list lock. for (const auto& handle : handles) { - handle->resultMetadata.clear(); + if (discardPendingRows) { + handle->clearResult(); + } else { + handle->resultMetadata.clear(); + } } } @@ -291,7 +295,7 @@ void Connection::commit() { } updateLastUsed(); LOG("Committing transaction"); - clearResultMetadata(); + clearResultMetadata(true); SQLRETURN ret; { // Release the GIL during the blocking SQLEndTran network round-trip. @@ -308,7 +312,7 @@ void Connection::rollback() { } updateLastUsed(); LOG("Rolling back transaction"); - clearResultMetadata(); + clearResultMetadata(true); SQLRETURN ret; { // Release the GIL during the blocking SQLEndTran network round-trip. @@ -323,7 +327,7 @@ void Connection::setAutocommit(bool enable) { if (!_dbcHandle) { ThrowStdException("Connection handle not allocated"); } - clearResultMetadata(); + clearResultMetadata(true); SQLINTEGER value = enable ? SQL_AUTOCOMMIT_ON : SQL_AUTOCOMMIT_OFF; LOG("Setting autocommit=%d", enable); SQLRETURN ret; @@ -417,7 +421,7 @@ SqlHandlePtr Connection::allocStatementHandle() { } SQLRETURN Connection::setAttribute(SQLINTEGER attribute, py::object value) { - clearResultMetadata(); + clearResultMetadata(attribute == SQL_ATTR_RESET_CONNECTION); LOG("Setting SQL attribute=%d", attribute); // SQLPOINTER ptr = nullptr; // SQLINTEGER length = 0; @@ -604,7 +608,7 @@ bool Connection::reset() { if (!_dbcHandle) { ThrowStdException("Connection handle not allocated"); } - clearResultMetadata(); + clearResultMetadata(true); LOG("Resetting connection via SQL_ATTR_RESET_CONNECTION"); // NOTE: SQL_ATTR_RESET_CONNECTION is a pool-checkin reset: it asks the // driver to wipe per-session state (temp tables, open cursors, SET diff --git a/mssql_python/pybind/connection/connection.h b/mssql_python/pybind/connection/connection.h index f43613842..caee75038 100644 --- a/mssql_python/pybind/connection/connection.h +++ b/mssql_python/pybind/connection/connection.h @@ -103,7 +103,7 @@ class Connection { void allocateDbcHandle(); void checkError(SQLRETURN ret) const; void applyAttrsBefore(const py::dict& attrs_before); - void clearResultMetadata(); + void clearResultMetadata(bool discardPendingRows = false); std::u16string _connStr; bool _fromPool = false; diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index 51fe81414..a6bcdf45f 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -370,6 +370,12 @@ constexpr int MAX_NATIVE_ROW_COUNT = 1000000; constexpr size_t MAX_NATIVE_FETCH_BYTES = 256ULL * 1024 * 1024; constexpr size_t MAX_NATIVE_PARAMETER_BYTES = 256ULL * 1024 * 1024; +class NativeFetchLimit : public std::runtime_error { + public: + NativeFetchLimit() + : std::runtime_error("Native fetch buffers exceed the 256 MiB allocation limit") {} +}; + void ValidateNativeRowCount(int value, const char* name, bool allowZero) { const int minimum = allowZero ? 0 : 1; if (value < minimum || value > MAX_NATIVE_ROW_COUNT) { @@ -475,7 +481,7 @@ void ReserveNativeFetchBytes(size_t& reservedBytes, size_t count, size_t element const size_t totalBytes = CheckedAddSize(reservedBytes, allocationBytes, "Native fetch buffer size is too large"); if (totalBytes > MAX_NATIVE_FETCH_BYTES) { - ThrowStdException("Native fetch buffers exceed the 256 MiB allocation limit"); + throw NativeFetchLimit(); } reservedBytes = totalBytes; } @@ -489,13 +495,14 @@ void ResizeNativeFetchBuffer(std::vector& buffer, size_t count, const size_t oldBytes = CheckedMultiplySize( buffer.size(), sizeof(ElementType), "Native fetch buffer size is too large"); - ReserveNativeFetchBytes(reservedBytes, count, sizeof(ElementType)); + size_t newReservedBytes = reservedBytes; + ReserveNativeFetchBytes(newReservedBytes, count, sizeof(ElementType)); { std::vector resizedBuffer(count); std::copy(buffer.begin(), buffer.end(), resizedBuffer.begin()); buffer.swap(resizedBuffer); } - reservedBytes -= oldBytes; + reservedBytes = newReservedBytes - oldBytes; } template @@ -508,6 +515,11 @@ void EnsureNativeFetchBufferSize(std::vector& buffer, size_t requir while (newSize < requiredSize) { newSize = CheckedMultiplySize(newSize, 2, "Native fetch buffer size is too large"); } + const size_t available = (MAX_NATIVE_FETCH_BYTES - reservedBytes) / sizeof(ElementType); + newSize = std::min(newSize, available); + if (newSize < requiredSize) { + throw NativeFetchLimit(); + } ResizeNativeFetchBuffer(buffer, newSize, reservedBytes); } @@ -1940,7 +1952,7 @@ void SqlHandle::markImplicitlyFreed() { _type); return; // Refuse to mark - let normal free() handle it } - resultMetadata.clear(); + clearResult(); _implicitly_freed = true; } @@ -1957,7 +1969,7 @@ void SqlHandle::free() { SQLRETURN SqlHandle::freeHandle() { PERF_TIMER("SqlHandle::free"); - resultMetadata.clear(); + clearResult(); bool pythonShuttingDown = is_python_finalizing(); bool skipDuringShutdown = _type == SQL_HANDLE_STMT || _type == SQL_HANDLE_DBC; #ifdef _WIN32 @@ -1998,7 +2010,7 @@ SQLRETURN SqlHandle::freeHandle() { } void SqlHandle::close_cursor() { - resultMetadata.clear(); + clearResult(); if (is_python_finalizing()) { return; } @@ -2026,7 +2038,7 @@ void SqlHandle::close_cursor() { } void SqlHandle::cancel() { - resultMetadata.clear(); + clearResult(); if (is_python_finalizing()) { return; } @@ -2061,7 +2073,7 @@ SQLRETURN SQLResetStmt_wrap(SqlHandlePtr statementHandle) { if (statementHandle->isImplicitlyFreed()) { return SQL_INVALID_HANDLE; } - statementHandle->resultMetadata.clear(); + statementHandle->clearResult(); if (!SQLFreeStmt_ptr) { DriverLoader::getInstance().loadDriver(); } @@ -2083,7 +2095,7 @@ SQLRETURN SQLResetStmt_wrap(SqlHandlePtr statementHandle) { SQLRETURN SQLGetTypeInfo_Wrapper(SqlHandlePtr StatementHandle, SQLSMALLINT DataType) { PERF_TIMER("SQLGetTypeInfo_Wrapper"); - StatementHandle->resultMetadata.clear(); + StatementHandle->clearResult(); if (!SQLGetTypeInfo_ptr) { ThrowStdException("SQLGetTypeInfo function not loaded"); } @@ -2096,7 +2108,7 @@ SQLRETURN SQLGetTypeInfo_Wrapper(SqlHandlePtr StatementHandle, SQLSMALLINT DataT SQLRETURN SQLProcedures_wrap(SqlHandlePtr StatementHandle, const py::object& catalogObj, const py::object& schemaObj, const py::object& procedureObj) { PERF_TIMER("SQLProcedures_wrap"); - StatementHandle->resultMetadata.clear(); + StatementHandle->clearResult(); if (!SQLProcedures_ptr) { ThrowStdException("SQLProcedures function not loaded"); } @@ -2121,7 +2133,7 @@ SQLRETURN SQLForeignKeys_wrap(SqlHandlePtr StatementHandle, const py::object& pk const py::object& fkCatalogObj, const py::object& fkSchemaObj, const py::object& fkTableObj) { PERF_TIMER("SQLForeignKeys_wrap"); - StatementHandle->resultMetadata.clear(); + StatementHandle->clearResult(); if (!SQLForeignKeys_ptr) { ThrowStdException("SQLForeignKeys function not loaded"); } @@ -2154,7 +2166,7 @@ SQLRETURN SQLForeignKeys_wrap(SqlHandlePtr StatementHandle, const py::object& pk SQLRETURN SQLPrimaryKeys_wrap(SqlHandlePtr StatementHandle, const py::object& catalogObj, const py::object& schemaObj, const std::u16string& table) { PERF_TIMER("SQLPrimaryKeys_wrap"); - StatementHandle->resultMetadata.clear(); + StatementHandle->clearResult(); if (!SQLPrimaryKeys_ptr) { ThrowStdException("SQLPrimaryKeys function not loaded"); } @@ -2177,7 +2189,7 @@ SQLRETURN SQLStatistics_wrap(SqlHandlePtr StatementHandle, const py::object& cat const py::object& schemaObj, const std::u16string& table, SQLUSMALLINT unique, SQLUSMALLINT reserved) { PERF_TIMER("SQLStatistics_wrap"); - StatementHandle->resultMetadata.clear(); + StatementHandle->clearResult(); if (!SQLStatistics_ptr) { ThrowStdException("SQLStatistics function not loaded"); } @@ -2200,7 +2212,7 @@ SQLRETURN SQLColumns_wrap(SqlHandlePtr StatementHandle, const py::object& catalo const py::object& schemaObj, const py::object& tableObj, const py::object& columnObj) { PERF_TIMER("SQLColumns_wrap"); - StatementHandle->resultMetadata.clear(); + StatementHandle->clearResult(); if (!SQLColumns_ptr) { ThrowStdException("SQLColumns function not loaded"); } @@ -2355,7 +2367,7 @@ static void CheckFetchError(const SqlHandlePtr& handle, SQLRETURN ret) { // Wrap SQLExecDirect SQLRETURN SQLExecDirect_wrap(SqlHandlePtr StatementHandle, const std::u16string& Query) { PERF_TIMER("SQLExecDirect_wrap"); - StatementHandle->resultMetadata.clear(); + StatementHandle->clearResult(); LOG("SQLExecDirect: Executing query directly - statement_handle=%p, " "query_length=%zu chars", (void*)StatementHandle->get(), Query.length()); @@ -2392,7 +2404,7 @@ SQLRETURN SQLTables_wrap(SqlHandlePtr StatementHandle, const std::u16string& cat const std::u16string& schema, const std::u16string& table, const std::u16string& tableType) { PERF_TIMER("SQLTables_wrap"); - StatementHandle->resultMetadata.clear(); + StatementHandle->clearResult(); if (!SQLTables_ptr) { LOG("SQLTables: Function pointer not initialized, loading driver"); DriverLoader::getInstance().loadDriver(); @@ -2440,7 +2452,7 @@ SQLRETURN SQLExecute_wrap(const SqlHandlePtr statementHandle, return SQL_INVALID_HANDLE; } - statementHandle->resultMetadata.clear(); + statementHandle->clearResult(); SQLHANDLE hStmt = statementHandle->get(); // Configure forward-only / read-only cursor (matches slow path semantics). @@ -3266,7 +3278,7 @@ SQLRETURN SQLExecuteMany_wrap(const SqlHandlePtr statementHandle, const std::u16 if (!statementHandle || !statementHandle->get()) { return SQL_INVALID_HANDLE; } - statementHandle->resultMetadata.clear(); + statementHandle->clearResult(); LOG("SQLExecuteMany: Starting batch execution - param_count=%zu, " "param_set_size=%zu", columnwise_params.size(), paramSetSize); @@ -3554,7 +3566,7 @@ SQLRETURN SQLSpecialColumns_wrap(SqlHandlePtr StatementHandle, SQLSMALLINT ident const std::u16string& table, SQLSMALLINT scope, SQLSMALLINT nullable) { PERF_TIMER("SQLSpecialColumns_wrap"); - StatementHandle->resultMetadata.clear(); + StatementHandle->clearResult(); if (!SQLSpecialColumns_ptr) { ThrowStdException("SQLSpecialColumns function not loaded"); } @@ -3573,6 +3585,22 @@ SQLRETURN SQLSpecialColumns_wrap(SqlHandlePtr StatementHandle, SQLSMALLINT ident table.empty() ? 0 : SQL_NTS, scope, nullable); } +std::shared_ptr PendingRowsForNextFetch(SqlHandlePtr handle) { + auto pending = handle->pendingArrowRows(); + if (pending) { + pending->hasSelectedRow = false; + if (pending->next == pending->count) { + handle->setPendingArrowRows(nullptr); + return nullptr; + } + } + return pending; +} + +SQLRETURN ReadSelectedArrowRow(SqlHandlePtr handle, PendingArrowRows& pending, + SQLUSMALLINT columnCount, py::list& row, + const std::string& charEncoding, py::handle messages); + // Wrap SQLFetch to retrieve rows SQLRETURN SQLFetch_wrap(SqlHandlePtr StatementHandle) { PERF_TIMER("SQLFetch_wrap"); @@ -3581,6 +3609,11 @@ SQLRETURN SQLFetch_wrap(SqlHandlePtr StatementHandle) { LOG("SQLFetch: Function pointer not initialized, loading driver"); DriverLoader::getInstance().loadDriver(); // Load the driver } + if (auto pending = PendingRowsForNextFetch(StatementHandle)) { + pending->selectedRow = pending->next++; + pending->hasSelectedRow = true; + return SQL_SUCCESS; + } // Release the GIL during the blocking ODBC call SQLRETURN ret = SQL_ERROR; @@ -3795,6 +3828,14 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p // Note: wcharEncoding parameter is reserved for future use // Currently WCHAR data always uses UTF-16LE for Windows compatibility (void)wcharEncoding; // Suppress unused parameter warning + if (auto pending = StatementHandle->pendingArrowRows()) { + if (!pending->hasSelectedRow) { + ThrowStdException("Pending Arrow data requires a successful DDBCSQLFetch or " + "DDBCSQLFetchScroll before DDBCSQLGetData"); + } + return ReadSelectedArrowRow(StatementHandle, *pending, colCount, row, + charEncoding, messages); + } LOG("SQLGetData: Getting data from %d columns for statement_handle=%p (charCtype=%d)", colCount, (void*)StatementHandle->get(), charCtype); @@ -4583,7 +4624,7 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p } } if (!complete) { - StatementHandle->resultMetadata.clear(); + StatementHandle->clearResult(); } else if (pending && pending->columns.size() == colCount) { StatementHandle->resultMetadata.publish(snapshot.generation, std::move(pending)); } @@ -4600,10 +4641,22 @@ SQLRETURN SQLFetchScroll_wrap(SqlHandlePtr StatementHandle, SQLSMALLINT FetchOri "the driver."); DriverLoader::getInstance().loadDriver(); // Load the driver } + if (StatementHandle->pendingArrowRows()) { + if (FetchOrientation != SQL_FETCH_NEXT) { + ThrowStdException("Only SQL_FETCH_NEXT is supported while Arrow rows remain pending"); + } + SQLRETURN ret = SQLFetch_wrap(StatementHandle); + CheckFetchError(StatementHandle, ret); + if (SQL_SUCCEEDED(ret) && row_data.size() == 0) { + const SQLSMALLINT colCount = SQLNumResultCols_wrap(StatementHandle); + ret = SQLGetData_wrap(StatementHandle, colCount, row_data); + } + return ret; + } // Unbind any columns from previous fetch operations to avoid memory // corruption - SQLFreeStmt_ptr(StatementHandle->get(), SQL_UNBIND); + CheckFetchError(StatementHandle, SQLFreeStmt_ptr(StatementHandle->get(), SQL_UNBIND)); // Perform scroll operation SQLRETURN ret = SQL_ERROR; @@ -4613,6 +4666,7 @@ SQLRETURN SQLFetchScroll_wrap(SqlHandlePtr StatementHandle, SQLSMALLINT FetchOri py::gil_scoped_release release; ret = SQLFetchScroll_ptr(StatementHandle->get(), FetchOrientation, FetchOffset); } + CheckFetchError(StatementHandle, ret); // If successful and caller wants data, retrieve it if (SQL_SUCCEEDED(ret) && row_data.size() == 0) { @@ -4836,32 +4890,12 @@ SQLRETURN SQLBindColums(SQLHSTMT hStmt, ColumnBuffers& buffers, const Metadata& // Fetch rows in batches // TODO: Move to anonymous namespace, since it is not used outside this file template -SQLRETURN FetchBatchData(SQLHSTMT hStmt, ColumnBuffers& buffers, const Metadata& columnNames, +SQLRETURN ConvertBufferedRows(SQLHSTMT hStmt, ColumnBuffers& buffers, const Metadata& columnNames, py::list& rows, SQLUSMALLINT numCols, SQLULEN& numRowsFetched, const std::vector& lobColumns, 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; - { - numRowsFetched = 0; - // Release the GIL during the blocking ODBC fetch - py::gil_scoped_release release; - 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; - } - if (!SQL_SUCCEEDED(ret)) { - LOG("FetchBatchData: Error while fetching rows in batches - " - "SQLRETURN=%d", - ret); - return ret; - } + py::handle messages = {}, SQLULEN firstRow = 0, + int bufferRows = 0) { // Pre-cache column metadata to avoid repeated dictionary lookups. // The vectors below are consumed later by construct_rows, so they are // declared at function scope; only the population work is wrapped in the @@ -4919,6 +4953,18 @@ SQLRETURN FetchBatchData(SQLHSTMT hStmt, ColumnBuffers& buffers, const Metadata& columnInfos[col].processedColumnSize + 1; // +1 for null terminator #endif } + if (bufferRows > 0) { + const bool wide = (isCharType && useWideChar) || dt == SQL_WCHAR || + dt == SQL_WVARCHAR || dt == SQL_WLONGVARCHAR || + dt == SQL_SS_XML; + columnInfos[col].fetchBufferSize = + wide ? buffers.wcharBuffers[col].size() / bufferRows + : buffers.charBuffers[col].size() / bufferRows; + if (dt == SQL_BINARY || dt == SQL_VARBINARY || + dt == SQL_LONGVARBINARY || dt == SQL_SS_UDT) { + columnInfos[col].processedColumnSize = columnInfos[col].fetchBufferSize; + } + } } for (SQLUSMALLINT col = 0; col < numCols; col++) { @@ -4970,6 +5016,7 @@ SQLRETURN FetchBatchData(SQLHSTMT hStmt, ColumnBuffers& buffers, const Metadata& case SQL_WCHAR: case SQL_WVARCHAR: case SQL_WLONGVARCHAR: + case SQL_SS_XML: columnProcessors[col] = ColumnProcessors::ProcessWChar; break; case SQL_SS_UDT: @@ -5007,7 +5054,7 @@ SQLRETURN FetchBatchData(SQLHSTMT hStmt, ColumnBuffers& buffers, const Metadata& void release() { released = true; } }; - for (SQLULEN i = 0; i < numRowsFetched; i++) { + for (SQLULEN i = firstRow; i < numRowsFetched; i++) { // Create row and immediately fill it (atomic operation per row) // This eliminates the two-phase pattern that could leave garbage rows // on exception @@ -5195,6 +5242,30 @@ SQLRETURN FetchBatchData(SQLHSTMT hStmt, ColumnBuffers& buffers, const Metadata& guard.release(); Py_DECREF(row); } + return SQL_SUCCESS; +} + +template +SQLRETURN FetchBatchData(SQLHSTMT hStmt, ColumnBuffers& buffers, const Metadata& columnNames, + py::list& rows, SQLUSMALLINT numCols, SQLULEN& numRowsFetched, + const std::vector& lobColumns, + 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; + { + numRowsFetched = 0; + py::gil_scoped_release release; + PERF_TIMER("FetchBatchData::SQLFetchScroll_call"); + ret = SQLFetchScroll_ptr(hStmt, SQL_FETCH_NEXT, 0); + } + CaptureFetchDiagnostics(hStmt, ret, messages); + if (!SQL_SUCCEEDED(ret)) { + return ret; + } + ConvertBufferedRows(hStmt, buffers, columnNames, rows, numCols, numRowsFetched, + lobColumns, charEncoding, charCtype, messages); return ret; } @@ -5362,6 +5433,9 @@ struct FetchStateGuard { } }; +SQLRETURN DrainPendingArrowRows(SqlHandlePtr handle, py::list& rows, SQLULEN limit, + const std::string& charEncoding, py::handle messages); + // FetchMany_wrap - Fetches multiple rows of data from the result set. // // @param StatementHandle: Handle to the statement from which data is to be @@ -5389,6 +5463,9 @@ SQLRETURN FetchMany_wrap(SqlHandlePtr StatementHandle, py::list& rows, py::handl // 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. charCtype = EffectiveCharCtypeForFetch(charCtype, charEncoding); + if (auto pending = PendingRowsForNextFetch(StatementHandle)) { + return DrainPendingArrowRows(StatementHandle, rows, fetchSize, charEncoding, messages); + } SQLRETURN ret = SQL_ERROR; ResultMetadataFailureGuard metadataFailure(StatementHandle->resultMetadata, ret); SQLHSTMT hStmt = StatementHandle->get(); @@ -5518,9 +5595,9 @@ SQLRETURN FetchMany_wrap(SqlHandlePtr StatementHandle, py::list& rows, py::handl template SQLRETURN GetDataVar(SQLHSTMT hStmt, SQLUSMALLINT colNumber, SQLSMALLINT cType, std::vector& dataVec, SQLLEN* indicator, size_t& reservedBytes, - py::handle messages) { - size_t start = 0; - size_t end = 0; + py::handle messages, PendingArrowRows& pending) { + size_t& start = pending.lobStart; + size_t& end = pending.lobRequired; // Determine null terminator size based on data type size_t sizeNullTerminator = 0; @@ -5542,6 +5619,11 @@ SQLRETURN GetDataVar(SQLHSTMT hStmt, SQLUSMALLINT colNumber, SQLSMALLINT cType, } while (true) { + // Save continuation before allocating: a quota-limited batch can return now + // and resume the same SQLGetData stream in the next fetch call. + if (end > dataVec.size()) { + ResizeNativeFetchBuffer(dataVec, end, reservedBytes); + } SQLLEN localInd = 0; SQLRETURN ret = SQLGetData_ptr( hStmt, colNumber, cType, reinterpret_cast(dataVec.data() + start), @@ -5552,6 +5634,7 @@ SQLRETURN GetDataVar(SQLHSTMT hStmt, SQLUSMALLINT colNumber, SQLSMALLINT cType, // Handle NULL data if (localInd == SQL_NULL_DATA) { *indicator = SQL_NULL_DATA; + start = end = 0; return SQL_SUCCESS; } @@ -5567,6 +5650,7 @@ SQLRETURN GetDataVar(SQLHSTMT hStmt, SQLUSMALLINT colNumber, SQLSMALLINT cType, } else { *indicator = localInd; // Preserve SQL_NO_TOTAL or other negative values } + start = end = 0; break; } @@ -5589,8 +5673,6 @@ SQLRETURN GetDataVar(SQLHSTMT hStmt, SQLUSMALLINT colNumber, SQLSMALLINT cType, // The next read starts where the null terminator would have been placed start = dataVec.size() - sizeNullTerminator; - // Resize buffer for next iteration - ResizeNativeFetchBuffer(dataVec, end, reservedBytes); } else { // Unexpected return code return ret; @@ -5600,6 +5682,175 @@ SQLRETURN GetDataVar(SQLHSTMT hStmt, SQLUSMALLINT colNumber, SQLSMALLINT cType, return SQL_SUCCESS; } +template +SQLRETURN ReadArrowScalar(SQLHSTMT stmt, SQLUSMALLINT column, SQLSMALLINT cType, + std::vector& buffer, SQLLEN* indicator, size_t& reservedBytes, + py::handle messages, size_t count = 1) { + ResizeNativeFetchBuffer(buffer, count, reservedBytes); + const SQLRETURN ret = SQLGetData_ptr(stmt, column, cType, buffer.data(), + count * sizeof(T), indicator); + CaptureFetchDiagnostics(stmt, ret, messages); + return ret; +} + +SQLRETURN ReadArrowLobColumn(SQLHSTMT stmt, PendingArrowRows& pending, + SQLSMALLINT type, SQLUSMALLINT column, + size_t& reservedBytes, py::handle messages) { + struct ScratchBudget { + PendingArrowRows& pending; + size_t& total; + size_t before; + ~ScratchBudget() { pending.reservedBytes += total - before; } + } budget{pending, reservedBytes, reservedBytes}; + auto& buffers = pending.buffers; + SQLLEN* indicator = buffers.indicators[column].data(); + const SQLUSMALLINT ordinal = column + 1; + switch (type) { + case SQL_SS_UDT: + case SQL_BINARY: + case SQL_VARBINARY: + case SQL_LONGVARBINARY: + return GetDataVar(stmt, ordinal, SQL_C_BINARY, buffers.charBuffers[column], + indicator, reservedBytes, messages, pending); + case SQL_CHAR: + case SQL_VARCHAR: + case SQL_LONGVARCHAR: + if (pending.charCtype == SQL_C_CHAR) { + return GetDataVar(stmt, ordinal, SQL_C_CHAR, buffers.charBuffers[column], + indicator, reservedBytes, messages, pending); + } + [[fallthrough]]; + case SQL_SS_XML: + case SQL_WCHAR: + case SQL_WVARCHAR: + case SQL_WLONGVARCHAR: + return GetDataVar(stmt, ordinal, SQL_C_WCHAR, buffers.wcharBuffers[column], + indicator, reservedBytes, messages, pending); + case SQL_INTEGER: + return ReadArrowScalar(stmt, ordinal, SQL_C_SLONG, buffers.intBuffers[column], + indicator, reservedBytes, messages); + case SQL_SMALLINT: + return ReadArrowScalar(stmt, ordinal, SQL_C_SSHORT, buffers.smallIntBuffers[column], + indicator, reservedBytes, messages); + case SQL_TINYINT: + case SQL_BIT: + return ReadArrowScalar(stmt, ordinal, type == SQL_BIT ? SQL_C_BIT : SQL_C_TINYINT, + buffers.charBuffers[column], indicator, reservedBytes, messages); + case SQL_REAL: + return ReadArrowScalar(stmt, ordinal, SQL_C_FLOAT, buffers.realBuffers[column], + indicator, reservedBytes, messages); + case SQL_DECIMAL: + case SQL_NUMERIC: + return ReadArrowScalar(stmt, ordinal, SQL_C_CHAR, buffers.charBuffers[column], + indicator, reservedBytes, messages, MAX_DIGITS_IN_NUMERIC); + case SQL_DOUBLE: + case SQL_FLOAT: + return ReadArrowScalar(stmt, ordinal, SQL_C_DOUBLE, buffers.doubleBuffers[column], + indicator, reservedBytes, messages); + case SQL_TIMESTAMP: + case SQL_TYPE_TIMESTAMP: + case SQL_DATETIME: + return ReadArrowScalar(stmt, ordinal, SQL_C_TYPE_TIMESTAMP, + buffers.timestampBuffers[column], indicator, + reservedBytes, messages); + case SQL_BIGINT: + return ReadArrowScalar(stmt, ordinal, SQL_C_SBIGINT, buffers.bigIntBuffers[column], + indicator, reservedBytes, messages); + case SQL_TYPE_DATE: + return ReadArrowScalar(stmt, ordinal, SQL_C_TYPE_DATE, buffers.dateBuffers[column], + indicator, reservedBytes, messages); + case SQL_SS_TIME2: + return ReadArrowScalar(stmt, ordinal, SQL_C_SS_TIME2, buffers.timeBuffers[column], + indicator, reservedBytes, messages); + case SQL_GUID: + return ReadArrowScalar(stmt, ordinal, SQL_C_GUID, buffers.guidBuffers[column], + indicator, reservedBytes, messages); + case SQL_SS_TIMESTAMPOFFSET: + return ReadArrowScalar(stmt, ordinal, SQL_C_SS_TIMESTAMPOFFSET, + buffers.datetimeoffsetBuffers[column], indicator, + reservedBytes, messages); + default: + ThrowStdException("Unsupported data type while recovering a pending Arrow row"); + } + return SQL_ERROR; +} + +SQLRETURN CompletePendingArrowLobs(SqlHandlePtr handle, PendingArrowRows& pending, + py::list& columns, SQLUSMALLINT columnCount, + py::handle messages) { + if (pending.hasLobs) { + size_t reservedBytes = pending.reservedBytes; + for (; pending.lobColumn < columnCount; ++pending.lobColumn) { + auto column = columns[pending.lobColumn].cast(); + const SQLRETURN ret = ReadArrowLobColumn(handle->get(), pending, + column["DataType"].cast(), + pending.lobColumn, reservedBytes, messages); + CheckFetchError(handle, ret); + if (!SQL_SUCCEEDED(ret)) { + return ret; + } + } + } + return SQL_SUCCESS; +} + +SQLRETURN ReadSelectedArrowRow(SqlHandlePtr handle, PendingArrowRows& pending, + SQLUSMALLINT columnCount, py::list& row, + const std::string& charEncoding, py::handle messages) { + py::list columns; + SQLRETURN ret = SQLDescribeCol_wrap(handle, columns, messages); + if (!SQL_SUCCEEDED(ret)) { + return ret; + } + if (columnCount > columns.size()) { + ThrowStdException("Column count exceeds the pending result column count"); + } + ret = CompletePendingArrowLobs(handle, pending, columns, columnCount, messages); + if (!SQL_SUCCEEDED(ret)) { + return ret; + } + py::list rows; + SQLULEN end = pending.selectedRow + 1; + ret = ConvertBufferedRows(handle->get(), pending.buffers, columns, rows, columnCount, end, + {}, charEncoding, pending.charCtype, messages, + pending.selectedRow, pending.fetchSize); + if (SQL_SUCCEEDED(ret)) { + for (auto value : rows[0].cast()) { + row.append(value); + } + } + return ret; +} + +SQLRETURN DrainPendingArrowRows(SqlHandlePtr handle, py::list& rows, SQLULEN limit, + const std::string& charEncoding, py::handle messages) { + auto pending = handle->pendingArrowRows(); + if (!pending || pending->next == pending->count) { + return SQL_SUCCESS; + } + py::list columns; + SQLRETURN ret = SQLDescribeCol_wrap(handle, columns, messages); + if (!SQL_SUCCEEDED(ret)) { + return ret; + } + const auto numCols = static_cast(columns.size()); + ret = CompletePendingArrowLobs(handle, *pending, columns, numCols, messages); + if (!SQL_SUCCEEDED(ret)) { + return ret; + } + SQLULEN end = pending->next + std::min(limit, pending->count - pending->next); + ret = ConvertBufferedRows(handle->get(), pending->buffers, columns, rows, numCols, end, + {}, charEncoding, pending->charCtype, messages, + pending->next, pending->fetchSize); + if (SQL_SUCCEEDED(ret)) { + pending->next = end; + if (pending->next == pending->count) { + handle->setPendingArrowRows(nullptr); + } + } + return ret; +} + 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) @@ -5625,9 +5876,17 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, CheckedAddSize(batchSize, 7, "Arrow batch size is too large") / 8; // Fetch narrow char data as SQL_C_CHAR if on Linux/macOS and configured by the user charCtype = EffectiveCharCtypeForFetch(charCtype, "utf-8"); + if (arrowBatchSize > 0) { + PendingRowsForNextFetch(StatementHandle); + } + auto [pendingGeneration, pending] = StatementHandle->pendingArrowSnapshot(); + if (pending) { + charCtype = pending->charCtype; + } // An overly large fetch size doesn't seem to help performance. - int fetchSize = arrowBatchSize > 0 ? std::min(arrowBatchSize, 64) : 1; + int fetchSize = pending ? pending->fetchSize + : arrowBatchSize > 0 ? std::min(arrowBatchSize, 64) : 1; SQLRETURN ret = SQL_ERROR; ResultMetadataFailureGuard metadataFailure(StatementHandle->resultMetadata, ret); @@ -5652,7 +5911,7 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, std::vector columnNullable(numCols); std::vector columnVarLen(numCols, false); std::vector nullCounts(numCols, 0); - size_t reservedBytes = 0; + size_t reservedBytes = pending ? pending->reservedBytes : 0; std::vector> arrowArrayPrivateData(numCols); std::vector> arrowSchemaPrivateData(numCols); @@ -5698,7 +5957,7 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, arrowColumnProducer->varVal = AllocateArrowArray(offsetCount, reservedBytes, "Arrow offset buffer is too large"); - ResizeNativeFetchBuffer(arrowColumnProducer->varData, initialVarDataSize, + EnsureNativeFetchBufferSize(arrowColumnProducer->varData, initialVarDataSize, reservedBytes); columnVarLen[i] = true; // start at offset 0 @@ -5713,7 +5972,7 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, arrowColumnProducer->varVal = AllocateArrowArray(offsetCount, reservedBytes, "Arrow offset buffer is too large"); - ResizeNativeFetchBuffer(arrowColumnProducer->varData, initialVarDataSize, + EnsureNativeFetchBufferSize(arrowColumnProducer->varData, initialVarDataSize, reservedBytes); columnVarLen[i] = true; // start at offset 0 @@ -5841,7 +6100,7 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, std::memset(arrowColumnProducer->valid.get(), 0xFF, bitmapSize); } - if (arrowBatchSize > 0 && !hasLobColumns) { + if (arrowBatchSize > 0 && !hasLobColumns && !pending) { const size_t rowSize = calculateRowSize(columnNames, numCols, charCtype); const size_t remainingBytes = MAX_NATIVE_FETCH_BYTES - reservedBytes; const size_t rowsWithinBudget = remainingBytes / rowSize; @@ -5852,7 +6111,7 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, } const int bufferFetchSize = arrowBatchSize > 0 ? fetchSize : 0; - if (bufferFetchSize > 0) { + if (!pending && bufferFetchSize > 0) { ReserveNativeFetchBytes( reservedBytes, CheckedMultiplySize(static_cast(numCols), @@ -5861,279 +6120,117 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, sizeof(SQLLEN)); } // Initialize column buffers - ColumnBuffers buffers(numCols, bufferFetchSize); - SQLULEN numRowsFetched = 0; + if (!pending) { + pending = std::make_shared(numCols, bufferFetchSize, charCtype, + hasLobColumns); + pending->reservedBytes = static_cast(numCols) * bufferFetchSize * sizeof(SQLLEN); + } + auto& buffers = pending->buffers; + SQLULEN& numRowsFetched = pending->count; FetchStateGuard fetchStateGuard(StatementHandle, messages); if (!hasLobColumns && bufferFetchSize > 0) { + const size_t before = reservedBytes; ret = SQLBindColums(hStmt, buffers, columnNames, numCols, fetchSize, reservedBytes, charCtype, messages); + pending->reservedBytes += reservedBytes - before; if (!SQL_SUCCEEDED(ret)) { LOG("Error when binding columns"); return ret; } } + if (arrowBatchSize > 0) { + if (!StatementHandle->publishPendingArrowRows(pendingGeneration, pending)) { + ThrowStdException("Result was reset while preparing an Arrow fetch"); + } + } fetchStateGuard.configure(&numRowsFetched, fetchSize); int idxRowArrow = 0; while (idxRowArrow < arrowBatchSize) { - int spaceLeftInArrowBatch = arrowBatchSize - idxRowArrow; - if (fetchSize > spaceLeftInArrowBatch) { - // Adjust fetch size for final batch to avoid overfetching - fetchStateGuard.setRowArraySize(spaceLeftInArrowBatch); - } - { - // Release GIL during the blocking ODBC fetch - py::gil_scoped_release release; - ret = SQLFetch_ptr(hStmt); - } - fetchStateGuard.observe(ret); - if (ret == SQL_NO_DATA) { - ret = SQL_SUCCESS; // Normal completion - break; - } - if (!SQL_SUCCEEDED(ret)) { - LOG("Error while fetching rows in batches"); - return ret; + if (pending->next == pending->count) { + fetchStateGuard.setRowArraySize(std::min(fetchSize, arrowBatchSize - idxRowArrow)); + { + py::gil_scoped_release release; + ret = SQLFetch_ptr(hStmt); + } + if (!SQL_SUCCEEDED(ret) && ret != SQL_NO_DATA) { + StatementHandle->setPendingArrowRows(nullptr); + } + fetchStateGuard.observe(ret); + if (ret == SQL_NO_DATA) { + pending->next = pending->count = 0; + ret = SQL_SUCCESS; + break; + } + if (!SQL_SUCCEEDED(ret)) { + return ret; + } + pending->next = 0; + pending->lobColumn = 0; + pending->lobStart = pending->lobRequired = 0; } - // numRowsFetched is the SQL_ATTR_ROWS_FETCHED_PTR attribute. - // It'll be populated by SQLFetch - assert(numRowsFetched + idxRowArrow <= static_cast(arrowBatchSize)); - for (SQLULEN idxRowSql = 0; idxRowSql < numRowsFetched; idxRowSql++) { - for (SQLUSMALLINT idxCol = 0; idxCol < numCols; idxCol++) { - auto& arrowColumnProducer = arrowArrayPrivateData[idxCol]; - auto dataType = dataTypes[idxCol]; - + for (SQLULEN idxRowSql = pending->next; + idxRowSql < numRowsFetched && idxRowArrow < arrowBatchSize; idxRowSql++) { + try { if (hasLobColumns) { - assert(idxRowSql == 0 && "GetData only works one row at a time"); - - switch (dataType) { - case SQL_SS_UDT: - case SQL_BINARY: - case SQL_VARBINARY: - case SQL_LONGVARBINARY: { - ret = GetDataVar(hStmt, idxCol + 1, SQL_C_BINARY, - buffers.charBuffers[idxCol], - buffers.indicators[idxCol].data(), reservedBytes, messages); - CheckFetchError(StatementHandle, ret); - if (!SQL_SUCCEEDED(ret)) { - LOG("Error fetching BINARY LOB for column %d", idxCol + 1); - return ret; - } - break; - } - case SQL_CHAR: - case SQL_VARCHAR: - case SQL_LONGVARCHAR: { - if (charCtype == SQL_C_CHAR) { - ret = GetDataVar(hStmt, idxCol + 1, SQL_C_CHAR, - buffers.charBuffers[idxCol], - buffers.indicators[idxCol].data(), reservedBytes, messages); - CheckFetchError(StatementHandle, ret); - if (!SQL_SUCCEEDED(ret)) { - LOG("Error fetching CHAR LOB data for column %d", idxCol + 1); - return ret; - } - break; - } - // else fall through to SQL_C_WCHAR case - } - case SQL_SS_XML: - case SQL_WCHAR: - case SQL_WVARCHAR: - case SQL_WLONGVARCHAR: { - ret = GetDataVar(hStmt, idxCol + 1, SQL_C_WCHAR, - buffers.wcharBuffers[idxCol], - buffers.indicators[idxCol].data(), reservedBytes, messages); - CheckFetchError(StatementHandle, ret); - if (!SQL_SUCCEEDED(ret)) { - LOG("Error fetching WCHAR LOB data for column %d", idxCol + 1); - return ret; - } - break; - } - case SQL_INTEGER: { - buffers.intBuffers[idxCol].resize(1); - 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; - } - break; - } - case SQL_SMALLINT: { - buffers.smallIntBuffers[idxCol].resize(1); - ret = SQLGetData_ptr(hStmt, idxCol + 1, SQL_C_SSHORT, - 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; - } - break; - } - case SQL_TINYINT: { - buffers.charBuffers[idxCol].resize(1); - ret = - 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; - } - break; - } - case SQL_BIT: { - buffers.charBuffers[idxCol].resize(1); - 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; - } - break; - } - case SQL_REAL: { - buffers.realBuffers[idxCol].resize(1); - 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; - } - break; - } - case SQL_DECIMAL: - case SQL_NUMERIC: { - buffers.charBuffers[idxCol].resize(MAX_DIGITS_IN_NUMERIC); - ret = SQLGetData_ptr(hStmt, idxCol + 1, SQL_C_CHAR, - 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; - } - break; - } - case SQL_DOUBLE: - case SQL_FLOAT: { - buffers.doubleBuffers[idxCol].resize(1); - ret = SQLGetData_ptr(hStmt, idxCol + 1, SQL_C_DOUBLE, - 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; - } - break; - } - case SQL_TIMESTAMP: - case SQL_TYPE_TIMESTAMP: - case SQL_DATETIME: { - buffers.timestampBuffers[idxCol].resize(1); - ret = SQLGetData_ptr(hStmt, idxCol + 1, SQL_C_TYPE_TIMESTAMP, - 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; - } - break; - } - case SQL_BIGINT: { - buffers.bigIntBuffers[idxCol].resize(1); - ret = SQLGetData_ptr(hStmt, idxCol + 1, SQL_C_SBIGINT, - 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; - } - break; - } - case SQL_TYPE_DATE: { - buffers.dateBuffers[idxCol].resize(1); - ret = SQLGetData_ptr(hStmt, idxCol + 1, SQL_C_TYPE_DATE, - 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; - } - break; - } - case SQL_SS_TIME2: { - buffers.timeBuffers[idxCol].resize(1); - ret = SQLGetData_ptr(hStmt, idxCol + 1, SQL_C_SS_TIME2, - 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; - } - break; + for (; pending->lobColumn < numCols; ++pending->lobColumn) { + ret = ReadArrowLobColumn(hStmt, *pending, dataTypes[pending->lobColumn], + pending->lobColumn, reservedBytes, messages); + CheckFetchError(StatementHandle, ret); + if (!SQL_SUCCEEDED(ret)) { + return ret; } - case SQL_GUID: { - buffers.guidBuffers[idxCol].resize(1); - 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; - } - break; - } - case SQL_SS_TIMESTAMPOFFSET: { - buffers.datetimeoffsetBuffers[idxCol].resize(1); - ret = SQLGetData_ptr(hStmt, idxCol + 1, SQL_C_SS_TIMESTAMPOFFSET, - 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); - return ret; - } - break; - } - default: { - std::ostringstream errorString; - errorString << "Unsupported data type for column ID - " << (idxCol + 1) - << ", Type - " << dataType; - LOG("SQLGetData: %s", errorString.str().c_str()); - ThrowStdException(errorString.str()); - break; + } + } + // Reserve the whole row before modifying any Arrow values or offsets. + // Keep the raw row (and the rest of its rowset) pending on quota exhaustion. + for (SQLUSMALLINT col = 0; col < numCols; ++col) { + if (!columnVarLen[col] || buffers.indicators[col][idxRowSql] == SQL_NULL_DATA) { + continue; + } + const SQLLEN indicator = buffers.indicators[col][idxRowSql]; + if (indicator < 0) { + ThrowStdException("Unexpected negative data length."); + } + size_t bytes = static_cast(indicator); + const auto type = dataTypes[col]; + if (type == SQL_GUID) { + bytes = 37; + } else if (type == SQL_WCHAR || type == SQL_WVARCHAR || + type == SQL_WLONGVARCHAR || type == SQL_SS_XML || + ((type == SQL_CHAR || type == SQL_VARCHAR || + type == SQL_LONGVARCHAR) && charCtype == SQL_C_WCHAR)) { + if (bytes % sizeof(SQLWCHAR) != 0) { + ThrowStdException("Wide-character data has an invalid byte length"); } + const auto& source = buffers.wcharBuffers[col]; + const size_t stride = source.size() / pending->fetchSize; + const size_t offset = + CheckedArrowSourceOffset(source, idxRowSql, stride, bytes); + bytes = simdutf::utf8_length_from_utf16le_with_replacement( + reinterpret_cast(source.data() + offset), + bytes / sizeof(SQLWCHAR)).count; } + auto& producer = arrowArrayPrivateData[col]; + EnsureNativeFetchBufferSize( + producer->varData, + CheckedAddSize(producer->varVal[idxRowArrow], bytes, + "Arrow value buffer is too large"), reservedBytes); + } + } catch (const NativeFetchLimit&) { + if (idxRowArrow == 0) { + ThrowStdException("One Arrow row cannot fit within the 256 MiB native " + "fetch budget. The row remains pending; use a row-wise " + "fetch to recover it when its source buffers fit."); } + goto arrow_batch_complete; + } + for (SQLUSMALLINT idxCol = 0; idxCol < numCols; idxCol++) { + auto& arrowColumnProducer = arrowArrayPrivateData[idxCol]; + auto dataType = dataTypes[idxCol]; SQLLEN indicator = buffers.indicators[idxCol][idxRowSql]; @@ -6250,14 +6347,6 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, static_assert(sizeof(SQLWCHAR) == sizeof(char16_t)); static_assert(alignof(SQLWCHAR) == alignof(char16_t)); const auto* utf16Source = reinterpret_cast(wcharSource); - size_t maxUtf8Size = - CheckedMultiplySize(dataLenW, 3, "Arrow value buffer is too large"); - EnsureNativeFetchBufferSize( - *target_vec, - CheckedAddSize(start, maxUtf8Size, - "Arrow value buffer is too large"), - reservedBytes); - size_t bytesWritten = simdutf::convert_utf16le_to_utf8_with_replacement( utf16Source, dataLenW, reinterpret_cast(target_vec->data() + start)); @@ -6414,10 +6503,15 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, } } idxRowArrow++; + pending->next++; } } +arrow_batch_complete: fetchStateGuard.close(); + if (arrowBatchSize > 0 && pending->next == pending->count) { + StatementHandle->setPendingArrowRows(nullptr); + } // Transfer ownership of buffers to batch ArrowSchema // First, allocate memory for the necessary structures @@ -6629,6 +6723,12 @@ SQLRETURN FetchAll_wrap(SqlHandlePtr StatementHandle, py::list& rows, // 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. charCtype = EffectiveCharCtypeForFetch(charCtype, charEncoding); + PendingRowsForNextFetch(StatementHandle); + SQLRETURN pendingRet = DrainPendingArrowRows( + StatementHandle, rows, std::numeric_limits::max(), charEncoding, messages); + if (!SQL_SUCCEEDED(pendingRet)) { + return pendingRet; + } SQLRETURN ret = SQL_ERROR; ResultMetadataFailureGuard metadataFailure(StatementHandle->resultMetadata, ret); SQLHSTMT hStmt = StatementHandle->get(); @@ -6800,6 +6900,17 @@ SQLRETURN FetchOne_wrap(SqlHandlePtr StatementHandle, py::list& row, // 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. charCtype = EffectiveCharCtypeForFetch(charCtype, charEncoding); + if (auto pending = PendingRowsForNextFetch(StatementHandle)) { + py::list rows; + const SQLRETURN ret = DrainPendingArrowRows(StatementHandle, rows, 1, + charEncoding, messages); + if (SQL_SUCCEEDED(ret)) { + for (auto value : rows[0].cast()) { + row.append(value); + } + } + return ret; + } SQLRETURN ret = SQL_ERROR; ResultMetadataFailureGuard metadataFailure(StatementHandle->resultMetadata, ret); SQLHSTMT hStmt = StatementHandle->get(); @@ -6837,7 +6948,7 @@ SQLRETURN FetchOne_wrap(SqlHandlePtr StatementHandle, py::list& row, // Wrap SQLMoreResults SQLRETURN SQLMoreResults_wrap(SqlHandlePtr StatementHandle) { PERF_TIMER("SQLMoreResults_wrap"); - StatementHandle->resultMetadata.clear(); + StatementHandle->clearResult(); LOG("SQLMoreResults_wrap: Check for more results"); if (!SQLMoreResults_ptr) { LOG("SQLMoreResults_wrap: Function pointer not initialized. Loading " diff --git a/mssql_python/pybind/ddbc_bindings.h b/mssql_python/pybind/ddbc_bindings.h index 748ecc5ca..0d4f07212 100644 --- a/mssql_python/pybind/ddbc_bindings.h +++ b/mssql_python/pybind/ddbc_bindings.h @@ -293,6 +293,8 @@ struct ConnectionCleanupState { bool disconnected = false; // Protected by mutex, shared with every child. }; +struct PendingArrowRows; + class SqlHandle { public: SqlHandle(SQLSMALLINT type, SQLHANDLE rawHandle, @@ -334,8 +336,37 @@ class SqlHandle { std::unordered_map describeCache; void clearDescribeCache() { describeCache.clear(); } ResultMetadataCache resultMetadata; + std::shared_ptr pendingArrowRows() const { + std::lock_guard lock(_pendingRowsMutex); + return _pendingArrowRows; + } + std::pair> pendingArrowSnapshot() const { + std::lock_guard lock(_pendingRowsMutex); + return {_pendingRowsGeneration, _pendingArrowRows}; + } + bool publishPendingArrowRows(uint64_t generation, std::shared_ptr rows) { + std::lock_guard lock(_pendingRowsMutex); + if (generation != _pendingRowsGeneration) { + return false; + } + _pendingArrowRows = std::move(rows); + return true; + } + void setPendingArrowRows(std::shared_ptr rows) { + std::lock_guard lock(_pendingRowsMutex); + _pendingArrowRows = std::move(rows); + } + void clearResult() { + resultMetadata.clear(); + std::lock_guard lock(_pendingRowsMutex); + ++_pendingRowsGeneration; + _pendingArrowRows.reset(); + } private: + mutable std::mutex _pendingRowsMutex; + uint64_t _pendingRowsGeneration = 0; + std::shared_ptr _pendingArrowRows; // The caller must release the GIL before waiting for native cleanup. std::unique_lock lockForCleanup() const; SQLSMALLINT _type; @@ -449,6 +480,26 @@ struct ColumnBuffers { indicators(numCols, std::vector(fetchSize)) {} }; +// No Python objects: cancellation and connection cleanup can release this without the GIL. +struct PendingArrowRows { + ColumnBuffers buffers; + int fetchSize; + int charCtype; + bool hasLobs; + SQLULEN count = 0; + SQLULEN next = 0; + SQLULEN selectedRow = 0; + bool hasSelectedRow = false; + size_t reservedBytes = 0; + SQLUSMALLINT lobColumn = 0; + size_t lobStart = 0; + size_t lobRequired = 0; + + PendingArrowRows(SQLSMALLINT columns, int fetchSize, int charCtype, bool hasLobs) + : buffers(columns, fetchSize), fetchSize(fetchSize), charCtype(charCtype), + hasLobs(hasLobs) {} +}; + // Performance: Column processor function type for fast type conversion // Using function pointers eliminates switch statement overhead in the hot loop typedef void (*ColumnProcessor)(PyObject* row, ColumnBuffers& buffers, const void* colInfo, diff --git a/tests/test_004_cursor.py b/tests/test_004_cursor.py index 313f5eac8..a19d21e61 100644 --- a/tests/test_004_cursor.py +++ b/tests/test_004_cursor.py @@ -21,6 +21,9 @@ import mssql_python import uuid import re +import subprocess +import sys +import textwrap from unittest.mock import call, patch from conftest import is_azure_sql_connection @@ -1755,6 +1758,43 @@ def test_fetchall_clamps_wide_result_batch_to_native_buffer_budget(cursor): assert [tuple(row) for row in cursor.fetchall()] == [("x",) * 34] +@pytest.mark.parametrize("sql_type", ["NVARCHAR(4000)", "NVARCHAR(MAX)"]) +def test_connection_attribute_preserves_pending_arrow_rows(conn_str, sql_type): + pytest.importorskip("pyarrow") + script = textwrap.dedent(""" + import os + import sys + import mssql_python + + with mssql_python.connect(os.environ["DB_CONNECTION_STRING"]) as connection: + with connection.cursor() as cursor: + columns = ", ".join( + f"CAST(REPLICATE(N'x', 4000) AS {sys.argv[1]}) AS value{i}" + for i in range(9) + ) + cursor.execute( + "SELECT TOP (8192) " + "ROW_NUMBER() OVER (ORDER BY a.object_id, b.object_id) AS id, " + + columns + + " FROM sys.all_objects AS a CROSS JOIN sys.all_objects AS b ORDER BY id" + ) + batch = cursor.arrow_batch(8192) + assert 0 < batch.num_rows < 8192 + next_id = batch.column(0)[-1].as_py() + 1 + del batch + connection.set_attr(mssql_python.SQL_ATTR_CONNECTION_TIMEOUT, 30) + assert cursor.fetchone()[0] == next_id + """) + result = subprocess.run( + [sys.executable, "-c", script, sql_type], + cwd=Path(__file__).resolve().parents[1], + capture_output=True, + text=True, + timeout=120, + ) + assert result.returncode == 0, result.stdout + result.stderr + + def test_description(cursor): """Test description""" cursor.execute("SELECT * FROM #pytest_all_data_types WHERE id = 1") diff --git a/tests/test_004_cursor_arrow.py b/tests/test_004_cursor_arrow.py index 859f7a630..2a4a87e90 100644 --- a/tests/test_004_cursor_arrow.py +++ b/tests/test_004_cursor_arrow.py @@ -7,6 +7,11 @@ import pytest import decimal import io +import os +from pathlib import Path +import subprocess +import sys +import textwrap from datetime import datetime, date, time, timezone import mssql_python @@ -22,6 +27,582 @@ pytestmark = pytest.mark.skipif(pa is None, reason="pyarrow is not installed") +_ARROW_BUDGET_SCRIPT = """ +import gc +import os +import sys +from datetime import date +from decimal import Decimal +from uuid import UUID +import mssql_python +import pyarrow.compute as pc + +def wide_query(sql_type, columns, count=8192, identifiers=False, extra=""): + payload = ", ".join( + f"CAST(REPLICATE(N'A', 4000) AS {sql_type}) AS c{i}" for i in range(columns) + ) + return f''' + WITH digits(d) AS ( + SELECT d FROM (VALUES (0),(1),(2),(3),(4),(5),(6),(7),(8),(9)) v(d) + ), numbers AS ( + SELECT TOP ({count}) + ROW_NUMBER() OVER (ORDER BY a.d,b.d,c.d,d.d) AS id + FROM digits a CROSS JOIN digits b CROSS JOIN digits c CROSS JOIN digits d + ) + SELECT {"id," if identifiers else ""} {payload} {extra} + FROM numbers ORDER BY id + ''' + +def check_payload(value, columns, offset=0): + for column in range(offset, offset + columns): + assert pc.all(pc.equal(value.column(column), "A" * 4000)).as_py() + +""" + + +def _run_arrow_budget_script(conn_str, script, *args): + env = os.environ.copy() + env["DB_CONNECTION_STRING"] = str(conn_str) + result = subprocess.run( + [sys.executable, "-c", _ARROW_BUDGET_SCRIPT + textwrap.dedent(script), *args], + cwd=Path(__file__).resolve().parents[1], + env=env, + capture_output=True, + text=True, + timeout=180, + ) + assert result.returncode == 0, result.stdout + result.stderr + + +@pytest.mark.parametrize("sql_type", ["nvarchar(4000)", "nvarchar(max)"]) +@pytest.mark.parametrize("method", ["arrow", "arrow_batch"]) +@pytest.mark.parametrize("columns", [7, 9]) +def test_arrow_native_budget_lossless_batches(conn_str, sql_type, method, columns): + """Fitting results stay whole; oversized batches return every row without loss.""" + _run_arrow_budget_script( + conn_str, + """ + sql_type, method, columns = sys.argv[1], sys.argv[2], int(sys.argv[3]) + with mssql_python.connect(os.environ["DB_CONNECTION_STRING"]) as conn: + with conn.cursor() as cur: + cur.execute(wide_query(sql_type, columns)) + if method == "arrow": + value = cur.arrow() + assert value.num_rows == 8192 + check_payload(value, columns) + if columns == 7: + assert value.nbytes == 229841920 + else: + assert value.nbytes > 256 * 1024 * 1024 + assert len(value.column(0).chunks) > 1 + del value + else: + total = 0 + total_bytes = 0 + first = True + while True: + value = cur.arrow_batch(8192) + count = value.num_rows + if first: + assert count == 8192 if columns == 7 else 0 < count < 8192 + first = False + if not count: + break + check_payload(value, columns) + total += count + total_bytes += value.nbytes + assert value.nbytes <= 256 * 1024 * 1024 + assert cur.rownumber == total - 1 + del value + assert total == 8192 + if columns == 7: + assert total_bytes == 229841920 + else: + assert total_bytes > 256 * 1024 * 1024 + assert cur.fetchone() is None + """, + sql_type, + method, + str(columns), + ) + + +@pytest.mark.parametrize("sql_type", ["nvarchar(4000)", "nvarchar(max)"]) +def test_arrow_native_budget_pending_mixed_fetches(conn_str, sql_type): + """A quota boundary preserves the overflow row and the rest of its native rowset.""" + _run_arrow_budget_script( + conn_str, + """ + sql_type = sys.argv[1] + extra = ''' + , CAST('123.45' AS DECIMAL(10,2)) amount + , CAST('2025-01-02' AS DATE) day + , CAST('58185E0D-3A91-44D8-BC46-7107217E0A6D' AS UNIQUEIDENTIFIER) guid + , CAST(0x010203 AS VARBINARY(3)) binary_value + , CAST(NULL AS INT) null_value + , CAST(1 AS BIT) flag + ''' + def check_row(row, expected): + assert row[0] == expected + assert all(value == "A" * 4000 for value in row[1:10]) + assert row[10] == Decimal("123.45") + assert row[11] == date(2025, 1, 2) + assert row[12] == UUID("58185E0D-3A91-44D8-BC46-7107217E0A6D") + assert row[13] == b"\\x01\\x02\\x03" + assert row[14] is None + assert row[15] is True + + with mssql_python.connect(os.environ["DB_CONNECTION_STRING"]) as conn: + for remaining in ("mixed", "fetchall", "arrow_reader"): + with conn.cursor() as cur: + cur.execute(wide_query(sql_type, 9, identifiers=True, extra=extra)) + first = cur.arrow_batch(8192) + n = first.num_rows + assert 0 < n < 8192 + assert first.column(0).to_pylist() == list(range(1, n + 1)) + del first + assert cur.arrow_batch(0).num_rows == 0 + from mssql_python import ddbc_bindings + from mssql_python.constants import ConstantsDDBC + # Updating the timeout invalidates metadata, not pending result rows. + assert ddbc_bindings.DDBCSQLSetStmtAttr( + cur.hstmt, ConstantsDDBC.SQL_ATTR_QUERY_TIMEOUT.value, 0 + ) in (0, 1) + if remaining == "mixed": + check_row(cur.fetchone(), n + 1) + n += 1 + rows = cur.fetchmany(5) + assert len(rows) == 5 + for i, row in enumerate(rows, n + 1): + check_row(row, i) + n += 5 + small = cur.arrow_batch(1) + assert small.column(0).to_pylist() == [n + 1] + n += 1 + del small, rows + rows = cur.fetchmany(100) + assert 0 < len(rows) <= 100 + for i, row in enumerate(rows, n + 1): + check_row(row, i) + n += len(rows) + del rows + rest = cur.arrow() + assert rest.column(0).to_pylist() == list(range(n + 1, 8193)) + del rest + elif remaining == "fetchall": + rows = cur.fetchall() + assert len(rows) == 8192 - n + for i, row in enumerate(rows, n + 1): + check_row(row, i) + del rows + else: + with cur.arrow_reader() as reader: + for batch in reader: + count = batch.num_rows + assert batch.column(0).to_pylist() == list( + range(n + 1, n + count + 1) + ) + n += count + del batch + assert n == 8192 + continue + assert cur.rownumber == 8191 + assert cur.fetchone() is None + gc.collect() + """, + sql_type, + ) + + +@pytest.mark.parametrize("sql_type", ["nvarchar(4000)", "nvarchar(max)"]) +def test_arrow_native_budget_pending_scroll_skip(conn_str, sql_type): + """Forward skips consume logical pending rows before advancing the driver.""" + _run_arrow_budget_script( + conn_str, + """ + sql_type = sys.argv[1] + with mssql_python.connect(os.environ["DB_CONNECTION_STRING"]) as conn: + with conn.cursor() as cur: + cur.execute(wide_query(sql_type, 9, identifiers=True)) + first = cur.arrow_batch(8192) + position = first.num_rows + assert 0 < position < 8192 + del first + cur.scroll(0) + cur.skip(0) + assert cur.rownumber == position - 1 + for value, mode in ((-1, "relative"), (0, "absolute"), (1, "absolute")): + try: + cur.scroll(value, mode) + except mssql_python.NotSupportedError: + pass + else: + raise AssertionError("Unsupported Python scroll was accepted") + assert cur.rownumber == position - 1 + cur.scroll(1) + position += 1 + assert cur.fetchone()[0] == position + 1 + position += 1 + cur.skip(3) + position += 3 + batch = cur.arrow_batch(1) + assert batch.column(0).to_pylist() == [position + 1] + position += 1 + del batch + # Cross the retained 64-row bound rowset (or the one-row LOB state). + offset = 67 if sql_type == "nvarchar(4000)" else 3 + cur.scroll(offset) + position += offset + rows = cur.fetchmany(2) + assert [row[0] for row in rows] == [position + 1, position + 2] + position += 2 + del rows + assert cur.rownumber == position - 1 + cur.skip(8192 - position) + assert cur.rownumber == 8191 + assert cur.fetchone() is None + try: + cur.scroll(1) + except IndexError: + pass + else: + raise AssertionError("Skipping beyond EOF did not raise") + assert cur.fetchone() is None + assert cur.execute("SELECT 42").fetchone()[0] == 42 + """, + sql_type, + ) + + +@pytest.mark.parametrize("sql_type", ["nvarchar(4000)", "nvarchar(max)"]) +def test_arrow_native_budget_pending_raw_fetches(conn_str, sql_type): + """Raw Fetch/FetchScroll select rows; GetData reads the selected cached row.""" + _run_arrow_budget_script( + conn_str, + """ + from mssql_python import ddbc_bindings as ddbc + from mssql_python.constants import ConstantsDDBC as C + sql_type = sys.argv[1] + unsupported = ( + (C.SQL_FETCH_RELATIVE.value, 0), (C.SQL_FETCH_RELATIVE.value, 1), + (C.SQL_FETCH_RELATIVE.value, 67), (C.SQL_FETCH_RELATIVE.value, -1), + (C.SQL_FETCH_ABSOLUTE.value, 0), (C.SQL_FETCH_ABSOLUTE.value, 1), + (C.SQL_FETCH_ABSOLUTE.value, -1), (C.SQL_FETCH_FIRST.value, 0), + (C.SQL_FETCH_LAST.value, 0), (C.SQL_FETCH_PRIOR.value, 0), + ) + def get_data(cur, count): + row = [] + assert ddbc.DDBCSQLGetData( + cur.hstmt, count, row, "utf-16le", "utf-16le", C.SQL_WCHAR.value + ) in (0, 1) + return row + with mssql_python.connect(os.environ["DB_CONNECTION_STRING"]) as conn: + with conn.cursor() as cur: + cur.execute(wide_query(sql_type, 9, identifiers=True)) + first = cur.arrow_batch(8192) + position = first.num_rows + assert 0 < position < 8192 + del first + try: + get_data(cur, 1) + except RuntimeError as exc: + assert "requires a successful DDBCSQLFetch" in str(exc) + else: + raise AssertionError("GetData consumed an unselected pending row") + for orientation, offset in unsupported: + try: + ddbc.DDBCSQLFetchScroll(cur.hstmt, orientation, offset, []) + except RuntimeError as exc: + assert "Only SQL_FETCH_NEXT" in str(exc) + else: + raise AssertionError("Unsupported scroll advanced pending data") + assert ddbc.DDBCSQLFetch(cur.hstmt) in (0, 1) + position += 1 + cur.scroll(0) + cur.skip(0) + assert get_data(cur, 0) == [] + assert get_data(cur, 1) == [position] + for orientation, offset in unsupported: + try: + ddbc.DDBCSQLFetchScroll(cur.hstmt, orientation, offset, []) + except RuntimeError as exc: + assert "Only SQL_FETCH_NEXT" in str(exc) + else: + raise AssertionError("Unsupported scroll discarded the selected row") + assert get_data(cur, 1) == [position] + assert cur.arrow_batch(0).num_rows == 0 + assert ddbc.DDBCSQLSetStmtAttr( + cur.hstmt, C.SQL_ATTR_QUERY_TIMEOUT.value, 0 + ) in (0, 1) + assert get_data(cur, 3) == [position, "A" * 4000, "A" * 4000] + try: + get_data(cur, 11) + except RuntimeError as exc: + assert "Column count exceeds" in str(exc) + else: + raise AssertionError("Invalid prefix count was accepted") + row = get_data(cur, 10) + assert row[0] == position + assert row[1:] == ["A" * 4000] * 9 + del row + row = [] + assert ddbc.DDBCSQLFetchScroll( + cur.hstmt, C.SQL_FETCH_NEXT.value, 0, row + ) in (0, 1) + position += 1 + assert row[0] == position + assert row[1:] == ["A" * 4000] * 9 + del row + assert cur.fetchone()[0] == position + 1 + position += 1 + assert ddbc.DDBCSQLFetch(cur.hstmt) in (0, 1) + position += 1 + assert get_data(cur, 1) == [position] + rows = cur.fetchmany(2) + assert [row[0] for row in rows] == [position + 1, position + 2] + position += 2 + del rows + # A raw fetch without GetData still consumes exactly one logical row. + assert ddbc.DDBCSQLFetch(cur.hstmt) in (0, 1) + position += 1 + batch = cur.arrow_batch(1) + assert batch.column(0).to_pylist() == [position + 1] + position += 1 + del batch + row = [] + assert ddbc.DDBCSQLFetchScroll( + cur.hstmt, C.SQL_FETCH_NEXT.value, 0, row + ) in (0, 1) + position += 1 + assert row[0] == position + assert row[1:] == ["A" * 4000] * 9 + del row + # NEXT ignores its offset, and a nonempty output suppresses materialization. + for offset in (0, 1, 67, -1, 123): + row = ["unchanged"] + assert ddbc.DDBCSQLFetchScroll( + cur.hstmt, C.SQL_FETCH_NEXT.value, offset, row + ) in (0, 1) + position += 1 + assert row == ["unchanged"] + assert get_data(cur, 1) == [position] + while True: + if position % 2 == 0: + ret = ddbc.DDBCSQLFetchScroll( + cur.hstmt, C.SQL_FETCH_NEXT.value, 0, ["no materialization"] + ) + else: + ret = ddbc.DDBCSQLFetch(cur.hstmt) + if ret == C.SQL_NO_DATA.value: + break + assert ret in (0, 1) + position += 1 + assert get_data(cur, 1) == [position] + assert position == 8192 + row = [] + assert ddbc.DDBCSQLFetchScroll( + cur.hstmt, C.SQL_FETCH_NEXT.value, 0, row + ) == C.SQL_NO_DATA.value + assert row == [] + assert cur.fetchone() is None + cur.execute(wide_query(sql_type, 9, identifiers=True)) + first = cur.arrow_batch(8192) + position = first.num_rows + del first + assert ddbc.DDBCSQLFetch(cur.hstmt) in (0, 1) + position += 1 + assert get_data(cur, 1) == [position] + rows = cur.fetchall() + assert [row[0] for row in rows] == list(range(position + 1, 8193)) + assert all(all(value == "A" * 4000 for value in row[1:]) for row in rows) + del rows + assert cur.fetchone() is None + """, + sql_type, + ) + + +def test_arrow_native_budget_native_scroll_orientations(conn_str): + """Driver-rejected forward-only orientations raise and leave the result untouched.""" + _run_arrow_budget_script( + conn_str, + """ + from mssql_python import ddbc_bindings as ddbc + from mssql_python.constants import ConstantsDDBC as C + unsupported = ( + (C.SQL_FETCH_RELATIVE.value, 0), (C.SQL_FETCH_RELATIVE.value, 1), + (C.SQL_FETCH_RELATIVE.value, 67), (C.SQL_FETCH_RELATIVE.value, -1), + (C.SQL_FETCH_ABSOLUTE.value, 0), (C.SQL_FETCH_ABSOLUTE.value, 1), + (C.SQL_FETCH_ABSOLUTE.value, -1), (C.SQL_FETCH_FIRST.value, 0), + (C.SQL_FETCH_LAST.value, 0), (C.SQL_FETCH_PRIOR.value, 0), + ) + with mssql_python.connect(os.environ["DB_CONNECTION_STRING"]) as conn: + for orientation, offset in unsupported: + with conn.cursor() as cur: + cur.execute("SELECT id FROM (VALUES (1),(2),(3)) v(id) ORDER BY id") + row = [] + try: + ddbc.DDBCSQLFetchScroll(cur.hstmt, orientation, offset, row) + except mssql_python.Error: + pass + else: + raise AssertionError("Unsupported native scroll did not raise") + assert row == [] + assert cur.fetchone()[0] == 1 + for offset in (0, 1, 67, -1): + with conn.cursor() as cur: + cur.execute("SELECT id FROM (VALUES (1),(2),(3)) v(id) ORDER BY id") + row = [] + assert ddbc.DDBCSQLFetchScroll( + cur.hstmt, C.SQL_FETCH_NEXT.value, offset, row + ) in (0, 1) + assert row == [1] + assert cur.fetchone()[0] == 2 + """, + ) + + +@pytest.mark.parametrize("recovery", ["arrow_batch", "fetchone", "raw_fetch"]) +def test_arrow_native_budget_resumes_partial_lob(conn_str, recovery): + """Resume SQLGetData when source growth, rather than Arrow growth, hits the budget.""" + _run_arrow_budget_script( + conn_str, + """ + recovery = sys.argv[1] + query = wide_query("nvarchar(max)", 7, count=6002, identifiers=True) + query = query.replace( + "CAST(REPLICATE(N'A', 4000) AS nvarchar(max)) AS c0", + "CASE WHEN id = 6001 THEN REPLICATE(CAST(N'A' AS nvarchar(max)), " + "40 * 1024 * 1024) ELSE CAST(REPLICATE(N'A', 4000) AS nvarchar(max)) END AS c0" + ) + with mssql_python.connect(os.environ["DB_CONNECTION_STRING"]) as conn: + with conn.cursor() as cur: + cur.execute(query) + first = cur.arrow_batch(8192) + assert first.num_rows == 6000 + del first + if recovery != "arrow_batch": + if recovery == "raw_fetch": + from mssql_python import ddbc_bindings as ddbc + from mssql_python.constants import ConstantsDDBC as C + assert ddbc.DDBCSQLFetch(cur.hstmt) in (0, 1) + prefix = [] + assert ddbc.DDBCSQLGetData( + cur.hstmt, 1, prefix, "utf-16le", "utf-16le", C.SQL_WCHAR.value + ) in (0, 1) + assert prefix == [6001] + row = [] + assert ddbc.DDBCSQLGetData( + cur.hstmt, 8, row, "utf-16le", "utf-16le", C.SQL_WCHAR.value + ) in (0, 1) + else: + row = cur.fetchone() + assert row[0] == 6001 + assert len(row[1]) == 40 * 1024 * 1024 + assert row[1].count("A") == len(row[1]) + assert all(value == "A" * 4000 for value in row[2:]) + del row + assert cur.fetchone()[0] == 6002 + else: + batch = cur.arrow_batch(8192) + assert batch.column(0).to_pylist() == [6001, 6002] + value = batch.column(1)[0].as_py() + assert len(value) == 40 * 1024 * 1024 + assert value.count("A") == len(value) + del value, batch + assert cur.fetchone() is None + """, + recovery, + ) + + +def test_arrow_native_budget_single_row_recovery(conn_str): + """A row that cannot coexist with its Arrow output stays available to row-wise fetches.""" + _run_arrow_budget_script( + conn_str, + """ + with mssql_python.connect(os.environ["DB_CONNECTION_STRING"]) as conn: + with conn.cursor() as cur: + cur.execute(''' + SELECT CASE WHEN id = 1 + THEN REPLICATE(CAST(NCHAR(20013) AS NVARCHAR(MAX)), 56 * 1024 * 1024) + ELSE N'end' END AS value + FROM (VALUES (1),(2)) v(id) ORDER BY id + ''') + initial_position = cur.rownumber + for _ in range(2): + try: + cur.arrow_batch(8192) + except RuntimeError as exc: + assert "row remains pending" in str(exc) + else: + raise AssertionError("Oversized Arrow row did not raise") + assert cur.rownumber == initial_position + assert cur.arrow_batch(0).num_rows == 0 + row = cur.fetchone() + assert len(row[0]) == 56 * 1024 * 1024 + assert row[0].count(chr(20013)) == len(row[0]) + del row + assert cur.fetchmany(2)[0][0] == "end" + assert cur.fetchone() is None + cur.execute(''' + SELECT CASE WHEN id = 1 + THEN REPLICATE(CAST(N'Z' AS NVARCHAR(MAX)), 140 * 1024 * 1024) + ELSE N'end' END AS value + FROM (VALUES (1),(2)) v(id) ORDER BY id + ''') + for method in ("arrow_batch", "fetchone", "arrow_batch"): + try: + getattr(cur, method)() + except RuntimeError as exc: + assert "256 MiB" in str(exc) + else: + raise AssertionError("A row exceeding the source budget was dropped") + cur.skip(1) + assert cur.fetchone()[0] == "end" + assert cur.fetchone() is None + assert cur.execute("SELECT 42").fetchone()[0] == 42 + """, + ) + + +@pytest.mark.parametrize("sql_type", ["nvarchar(4000)", "nvarchar(max)"]) +def test_arrow_native_budget_pending_reset_cleanup(conn_str, sql_type): + _run_arrow_budget_script( + conn_str, + """ + sql_type = sys.argv[1] + for operation in ("execute", "nextset", "close", "reader_close", "connection_close"): + with mssql_python.connect(os.environ["DB_CONNECTION_STRING"]) as conn: + with conn.cursor() as cur: + cur.execute(wide_query(sql_type, 9) + "; SELECT 42 AS answer") + first = cur.arrow_batch(8192) + assert 0 < first.num_rows < 8192 + del first + from mssql_python import ddbc_bindings + assert ddbc_bindings.DDBCSQLFetch(cur.hstmt) in (0, 1) + if operation == "nextset": + assert cur.nextset() + assert cur.fetchone()[0] == 42 + assert cur.fetchone() is None + elif operation == "execute": + assert cur.execute("SELECT 42").fetchone()[0] == 42 + elif operation == "reader_close": + reader = cur.arrow_reader() + reader.close() + assert cur.execute("SELECT 42").fetchone()[0] == 42 + elif operation == "close": + cur.close() + with conn.cursor() as fresh: + assert fresh.execute("SELECT 42").fetchone()[0] == 42 + else: + conn.close() + gc.collect() + """, + sql_type, + ) + + def get_arrow_test_data(include_lobs: bool, batch_length: int): arrow_test_data = [ (pa.uint8(), "tinyint", [1, 2, None, 4, 5, 0, 2**8 - 1]),