Skip to content

Commit 48a2f23

Browse files
gargsaumyaCopilot
andcommitted
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>
1 parent f533230 commit 48a2f23

3 files changed

Lines changed: 167 additions & 13 deletions

File tree

‎mssql_python/pybind/ddbc_bindings.cpp‎

Lines changed: 43 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -511,6 +511,29 @@ void EnsureNativeFetchBufferSize(std::vector<ElementType>& buffer, size_t requir
511511
ResizeNativeFetchBuffer(buffer, newSize, reservedBytes);
512512
}
513513

514+
size_t ExtendNativeFetchBuffer(std::vector<char>& buffer, size_t currentSize, size_t appendSize,
515+
size_t& reservedBytes) {
516+
const size_t requiredSize =
517+
CheckedAddSize(currentSize, appendSize, "Native fetch buffer size is too large");
518+
EnsureNativeFetchBufferSize(buffer, requiredSize, reservedBytes);
519+
return requiredSize;
520+
}
521+
522+
py::tuple TestLobBufferGrowth(size_t chunkCount, size_t chunkSize) {
523+
std::vector<char> buffer;
524+
size_t dataSize = 0;
525+
size_t reservedBytes = 0;
526+
size_t growthCount = 0;
527+
for (size_t chunk = 0; chunk < chunkCount; ++chunk) {
528+
const size_t previousAllocation = buffer.size();
529+
dataSize = ExtendNativeFetchBuffer(buffer, dataSize, chunkSize, reservedBytes);
530+
if (buffer.size() != previousAllocation) {
531+
++growthCount;
532+
}
533+
}
534+
return py::make_tuple(growthCount, dataSize, buffer.size(), reservedBytes);
535+
}
536+
514537
template <typename ElementType>
515538
std::unique_ptr<ElementType[]> AllocateArrowArray(size_t count, size_t& reservedBytes,
516539
const char* errorMessage) {
@@ -2396,6 +2419,7 @@ SQLRETURN SQLExecute_wrap(const SqlHandlePtr statementHandle,
23962419
bool use_prepare,
23972420
const py::dict& encoding_settings) {
23982421
PERF_TIMER("SQLExecute_wrap");
2422+
ValidateInputSizes(input_sizes.ptr());
23992423
if (!statementHandle || !statementHandle->get()) {
24002424
return SQL_INVALID_HANDLE;
24012425
}
@@ -3218,9 +3242,14 @@ SQLRETURN BindParameterArray(SqlHandle& handle, SQLHANDLE hStmt, const py::list&
32183242

32193243
SQLRETURN SQLExecuteMany_wrap(const SqlHandlePtr statementHandle, const std::u16string& query,
32203244
const py::list& columnwise_params,
3221-
std::vector<ParamInfo>& paramInfos, size_t paramSetSize,
3245+
std::vector<ParamInfo>& paramInfos, py::handle paramSetSizeArg,
32223246
const py::dict& encodingSettings) {
32233247
PERF_TIMER("SQLExecuteMany_wrap");
3248+
const size_t paramSetSize = static_cast<size_t>(
3249+
ValidateNativeRowCountArgument(paramSetSizeArg, "Parameter set size", false));
3250+
if (!statementHandle || !statementHandle->get()) {
3251+
return SQL_INVALID_HANDLE;
3252+
}
32243253
statementHandle->resultMetadata.clear();
32253254
LOG("SQLExecuteMany: Starting batch execution - param_count=%zu, "
32263255
"param_set_size=%zu",
@@ -3552,6 +3581,7 @@ py::object FetchLobColumnData(SQLHSTMT hStmt, SQLUSMALLINT colIndex, SQLSMALLINT
35523581
py::handle messages) {
35533582
PERF_TIMER("FetchLobColumnData");
35543583
std::vector<char> buffer;
3584+
size_t dataSize = 0;
35553585
size_t reservedBytes = 0;
35563586
ReserveNativeFetchBytes(reservedBytes, DAE_CHUNK_SIZE, sizeof(char));
35573587
SQLRETURN ret = SQL_SUCCESS_WITH_INFO;
@@ -3623,10 +3653,8 @@ py::object FetchLobColumnData(SQLHSTMT hStmt, SQLUSMALLINT colIndex, SQLSMALLINT
36233653
}
36243654
}
36253655
if (bytesRead > 0) {
3626-
const size_t previousSize = buffer.size();
3627-
const size_t requiredSize =
3628-
CheckedAddSize(previousSize, bytesRead, "LOB fetch buffer is too large");
3629-
ResizeNativeFetchBuffer(buffer, requiredSize, reservedBytes);
3656+
const size_t previousSize = dataSize;
3657+
dataSize = ExtendNativeFetchBuffer(buffer, dataSize, bytesRead, reservedBytes);
36303658
std::copy_n(chunk.data(), bytesRead, buffer.data() + previousSize);
36313659
LOG("FetchLobColumnData: Appended %zu bytes at loop %d", bytesRead, loopCount);
36323660
}
@@ -3635,21 +3663,21 @@ py::object FetchLobColumnData(SQLHSTMT hStmt, SQLUSMALLINT colIndex, SQLSMALLINT
36353663
break;
36363664
}
36373665
}
3638-
LOG("FetchLobColumnData: Total bytes collected=%zu for column %d", buffer.size(), colIndex);
3666+
LOG("FetchLobColumnData: Total bytes collected=%zu for column %d", dataSize, colIndex);
36393667

3640-
if (buffer.empty()) {
3668+
if (dataSize == 0) {
36413669
if (isBinary) {
36423670
return py::bytes("");
36433671
}
36443672
return py::str("");
36453673
}
36463674
if (isWideChar) {
3647-
if (buffer.size() % sizeof(SQLWCHAR) != 0) {
3675+
if (dataSize % sizeof(SQLWCHAR) != 0) {
36483676
ThrowStdException("Wide-character LOB data has an invalid byte length");
36493677
}
36503678
int byteOrder = -1;
36513679
PyObject* decoded =
3652-
PyUnicode_DecodeUTF16(buffer.data(), static_cast<Py_ssize_t>(buffer.size()),
3680+
PyUnicode_DecodeUTF16(buffer.data(), static_cast<Py_ssize_t>(dataSize),
36533681
"strict", &byteOrder);
36543682
if (decoded == nullptr) {
36553683
throw py::error_already_set();
@@ -3659,18 +3687,18 @@ py::object FetchLobColumnData(SQLHSTMT hStmt, SQLUSMALLINT colIndex, SQLSMALLINT
36593687
if (isBinary) {
36603688
LOG("FetchLobColumnData: Returning binary data - %zu bytes for column "
36613689
"%d",
3662-
buffer.size(), colIndex);
3663-
return py::bytes(buffer.data(), buffer.size());
3690+
dataSize, colIndex);
3691+
return py::bytes(buffer.data(), dataSize);
36643692
}
36653693

36663694
// For SQL_C_CHAR data, decode using the appropriate encoding.
36673695
const std::string effectiveCharEncoding = GetEffectiveCharDecoding(charEncoding);
3668-
py::bytes raw_bytes(buffer.data(), buffer.size());
3696+
py::bytes raw_bytes(buffer.data(), dataSize);
36693697
try {
36703698
py::object decoded = raw_bytes.attr("decode")(effectiveCharEncoding, "strict");
36713699
LOG("FetchLobColumnData: Decoded narrow string with '%s' - %zu bytes -> %zu chars for "
36723700
"column %d",
3673-
effectiveCharEncoding.c_str(), buffer.size(), py::len(decoded), colIndex);
3701+
effectiveCharEncoding.c_str(), dataSize, py::len(decoded), colIndex);
36743702
return decoded;
36753703
} catch (const py::error_already_set& e) {
36763704
LOG_ERROR("FetchLobColumnData: Failed to decode with '%s' for column %d: %s",
@@ -7000,6 +7028,8 @@ PYBIND11_MODULE(ddbc_bindings, m) {
70007028
"Fetch an arrow batch of given length from the result set", py::arg("StatementHandle"),
70017029
py::arg("capsules"), py::arg("arrowBatchSize"), py::arg("charCtype"),
70027030
py::arg("messages") = py::none());
7031+
m.def("_test_lob_buffer_growth", &TestLobBufferGrowth, py::arg("chunk_count"),
7032+
py::arg("chunk_size"));
70037033
m.def("DDBCSQLFreeHandle", &SQLFreeHandle_wrap, "Free a handle");
70047034
m.def("DDBCSQLResetStmt", &SQLResetStmt_wrap,
70057035
"Close cursor and unbind params without freeing HSTMT");

‎mssql_python/pybind/param_detect.hpp‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -274,6 +274,89 @@ inline void NormalizeTimeParam(PyObject* params, Py_ssize_t index, SQLULEN& colu
274274
}
275275
}
276276

277+
inline long long ValidatedInputSizeInteger(PyObject* value, const char* fieldName) {
278+
if (!PyLong_Check(value) || PyBool_Check(value)) {
279+
throw py::type_error(std::string(fieldName) + " must be an integer");
280+
}
281+
int overflow = 0;
282+
const long long result = PyLong_AsLongLongAndOverflow(value, &overflow);
283+
if ((result == -1 && PyErr_Occurred()) || overflow != 0) {
284+
PyErr_Clear();
285+
throw py::value_error(std::string(fieldName) + " is out of range");
286+
}
287+
return result;
288+
}
289+
290+
inline void ValidateInputSizes(PyObject* inputSizes) {
291+
if (inputSizes == Py_None) {
292+
return;
293+
}
294+
if (!PyList_Check(inputSizes)) {
295+
throw py::type_error("inputSizes must be None or a list");
296+
}
297+
298+
const Py_ssize_t count = PyList_GET_SIZE(inputSizes);
299+
for (Py_ssize_t index = 0; index < count; ++index) {
300+
PyObject* entry = PyList_GET_ITEM(inputSizes, index);
301+
if (!PyTuple_Check(entry) || PyTuple_GET_SIZE(entry) != 4) {
302+
throw py::type_error("each inputSizes entry must be a four-item tuple");
303+
}
304+
305+
const long long sqlType =
306+
ValidatedInputSizeInteger(PyTuple_GET_ITEM(entry, 0), "SQL type");
307+
const long long cType =
308+
ValidatedInputSizeInteger(PyTuple_GET_ITEM(entry, 1), "C type");
309+
PyObject* columnSize = PyTuple_GET_ITEM(entry, 2);
310+
PyObject* decimalDigits = PyTuple_GET_ITEM(entry, 3);
311+
if (!PyLong_Check(columnSize) || PyBool_Check(columnSize) ||
312+
!PyLong_Check(decimalDigits) || PyBool_Check(decimalDigits)) {
313+
throw py::type_error("column size and decimal digits must be integers");
314+
}
315+
316+
if (sqlType < std::numeric_limits<SQLSMALLINT>::min() ||
317+
sqlType > std::numeric_limits<SQLSMALLINT>::max() ||
318+
cType < std::numeric_limits<SQLSMALLINT>::min() ||
319+
cType > std::numeric_limits<SQLSMALLINT>::max()) {
320+
throw py::value_error("SQL and C types must fit in SQLSMALLINT");
321+
}
322+
const bool isNumeric = sqlType == SQL_DECIMAL || sqlType == SQL_NUMERIC;
323+
py::int_ zero(0);
324+
const int negativeColumnSize =
325+
PyObject_RichCompareBool(columnSize, zero.ptr(), Py_LT);
326+
const int negativeDecimalDigits =
327+
PyObject_RichCompareBool(decimalDigits, zero.ptr(), Py_LT);
328+
if (negativeColumnSize == -1 || negativeDecimalDigits == -1) {
329+
throw py::error_already_set();
330+
}
331+
if (negativeColumnSize == 1) {
332+
throw py::value_error("column size must be non-negative");
333+
}
334+
if (negativeDecimalDigits == 1) {
335+
throw py::value_error("decimal digits must be non-negative");
336+
}
337+
if (!isNumeric) {
338+
const unsigned long long requestedSize = PyLong_AsUnsignedLongLong(columnSize);
339+
if (requestedSize == static_cast<unsigned long long>(-1) && PyErr_Occurred()) {
340+
PyErr_Clear();
341+
throw py::value_error("column size is out of range");
342+
}
343+
if (requestedSize > std::numeric_limits<SQLULEN>::max()) {
344+
throw py::value_error("column size is out of range");
345+
}
346+
const unsigned long long requestedDigits =
347+
PyLong_AsUnsignedLongLong(decimalDigits);
348+
if (requestedDigits == static_cast<unsigned long long>(-1) && PyErr_Occurred()) {
349+
PyErr_Clear();
350+
throw py::value_error("decimal digits are out of range");
351+
}
352+
if (requestedDigits >
353+
static_cast<unsigned long long>(std::numeric_limits<SQLSMALLINT>::max())) {
354+
throw py::value_error("decimal digits are out of range");
355+
}
356+
}
357+
}
358+
}
359+
277360
inline void ApplyInputSizeOverride(PyObject* params, PyObject* inputSize, Py_ssize_t index,
278361
ParamInfo& info, const std::string& charEncoding) {
279362
py::tuple values = borrow<py::tuple>(inputSize);
@@ -391,6 +474,11 @@ inline std::vector<ParamInfo> DetectParamTypes(PyObject* params, PyObject* input
391474
const std::string& charEncoding = "utf-8") {
392475
PyTypeCache::initialize();
393476

477+
if (!PyList_Check(params)) {
478+
throw py::type_error("params must be a list");
479+
}
480+
ValidateInputSizes(inputSizes);
481+
394482
const Py_ssize_t n = PyList_GET_SIZE(params);
395483
const Py_ssize_t inputSizeCount = inputSizes == Py_None ? 0 : PyList_GET_SIZE(inputSizes);
396484
std::vector<ParamInfo> infos(n);

‎tests/test_010_pybind_functions.py‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,42 @@ def test_arrow_batch_rejects_unsafe_size_before_handle_access(batch_size):
5151
ddbc.DDBCSQLFetchArrowBatch(None, [], batch_size, 0)
5252

5353

54+
@pytest.mark.skipif(not DDBC_AVAILABLE, reason="ddbc_bindings not available")
55+
@pytest.mark.parametrize("param_set_size", [-1, 0, 1_000_001, True])
56+
def test_executemany_rejects_unsafe_size_before_handle_access(param_set_size):
57+
with pytest.raises(RuntimeError, match="Parameter set size"):
58+
ddbc.SQLExecuteMany(None, "", [], [], param_set_size, {})
59+
60+
61+
@pytest.mark.skipif(not DDBC_AVAILABLE, reason="ddbc_bindings not available")
62+
@pytest.mark.parametrize(
63+
("input_sizes", "message"),
64+
[
65+
((12, -8, 1, 0), "inputSizes must be None or a list"),
66+
([None], "four-item tuple"),
67+
([(12, -8, 1)], "four-item tuple"),
68+
([(12, -8, True, 0)], "must be integers"),
69+
([(12, -8, -1, 0)], "column size must be non-negative"),
70+
([(12, -8, 2**128, 0)], "column size is out of range"),
71+
([(12, -8, 1, 2**128)], "decimal digits are out of range"),
72+
],
73+
)
74+
def test_execute_rejects_malformed_input_sizes_before_handle_access(input_sizes, message):
75+
with pytest.raises((TypeError, ValueError), match=message):
76+
ddbc.DDBCSQLExecute(None, "", [], input_sizes, [False], True, {})
77+
78+
79+
@pytest.mark.skipif(not DDBC_AVAILABLE, reason="ddbc_bindings not available")
80+
def test_lob_buffer_growth_is_geometric():
81+
growth_count, data_size, allocation_size, reserved_bytes = ddbc._test_lob_buffer_growth(
82+
64, 64 * 1024
83+
)
84+
assert data_size == 4 * 1024 * 1024
85+
assert allocation_size == data_size
86+
assert reserved_bytes == allocation_size
87+
assert growth_count == 7
88+
89+
5490
@pytest.mark.skipif(not DDBC_AVAILABLE, reason="ddbc_bindings not available")
5591
class TestPybindModuleInfo:
5692
"""Test module information and architecture detection."""

0 commit comments

Comments
 (0)