FEAT: Streaming support in fetchone, fetchmany, fetchall for varbinarymax data type. (#232)

### Work Item / Issue Reference <!-- IMPORTANT: Please follow the PR template guidelines below. For mssql-python maintainers: Insert your ADO Work Item ID below (e.g. AB#37452) For external contributors: Insert Github Issue number below (e.g. #149) Only one reference is required - either GitHub issue OR ADO Work Item. --> <!-- mssql-python maintainers: ADO Work Item --> > [AB#33395](https://sqlclientdrivers.visualstudio.com/c6d89619-62de-46a0-8b46-70b92a84d85e/_workitems/edit/33395) <!-- External contributors: GitHub Issue --> > GitHub Issue: #<ISSUE_NUMBER> ------------------------------------------------------------------- ### Summary <!-- Insert your summary of changes below. Minimum 10 characters required. --> This pull request significantly improves support for streaming and fetching large binary (VARBINARY(MAX)) and large text columns in the MSSQL Python driver. The main changes include robust chunked retrieval of large objects (LOBs), correct handling of edge cases (such as empty or null values), and enhanced test coverage for these scenarios. **LOB Streaming and Fetching Enhancements:** * Added a new helper function `FetchLobColumnData` to efficiently stream and assemble large binary/text columns (LOBs) from the database, handling chunking, null/empty values, and correct type conversion for both binary and (wide/narrow) string columns. * Updated the logic in `SQLGetData_wrap`, `FetchBatchData`, `FetchMany_wrap`, and `FetchAll_wrap` to detect LOB columns and use the new streaming path for fetching them, including proper fallback to row-by-row fetching when LOBs are present. [[1]](diffhunk://#diff-dde2297345718ec449a14e7dff91b7bb2342b008ecc071f562233646d71144a1L2062-R2184) [[2]](diffhunk://#diff-dde2297345718ec449a14e7dff91b7bb2342b008ecc071f562233646d71144a1L2542-R2631) [[3]](diffhunk://#diff-dde2297345718ec449a14e7dff91b7bb2342b008ecc071f562233646d71144a1R2760-R2788) [[4]](diffhunk://#diff-dde2297345718ec449a14e7dff91b7bb2342b008ecc071f562233646d71144a1L2699-R2803) [[5]](diffhunk://#diff-dde2297345718ec449a14e7dff91b7bb2342b008ecc071f562233646d71144a1R2882-R2910) [[6]](diffhunk://#diff-dde2297345718ec449a14e7dff91b7bb2342b008ecc071f562233646d71144a1L2792-R2925) * Modified the batch fetch function signatures and logic to propagate LOB column information and ensure correct handling during bulk fetches. **Testing Improvements:** * Replaced and expanded the test for large binary data with a new, comprehensive test (`test_varbinarymax_insert_fetch`) that verifies insertion and retrieval of empty, small, and large VARBINARY(MAX) values (including edge cases around the 8000-byte threshold) using `fetchone`, `fetchall`, and `fetchmany`. These changes ensure that the driver can reliably handle large binary and text columns in all fetch scenarios, improving correctness and robustness for users working with LOB data. <!-- ### PR Title Guide > For feature requests FEAT: (short-description) > For non-feature requests like test case updates, config updates , dependency updates etc CHORE: (short-description) > For Fix requests FIX: (short-description) > For doc update requests DOC: (short-description) > For Formatting, indentation, or styling update STYLE: (short-description) > For Refactor, without any feature changes REFACTOR: (short-description) > For release related changes, without any feature changes RELEASE: #<RELEASE_VERSION> (short-description) ### Contribution Guidelines External contributors: - Create a GitHub issue first: https://github.com/microsoft/mssql-python/issues/new - Link the GitHub issue in the "GitHub Issue" section above - Follow the PR title format and provide a meaningful summary mssql-python maintainers: - Create an ADO Work Item following internal processes - Link the ADO Work Item in the "ADO Work Item" section above - Follow the PR title format and provide a meaningful summary -->

Author: gargsaumya

mssql_python/pybind/ddbc_bindings.cpp

@@ -2155,45 +2155,40 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p
  case SQL_BINARY:
  case SQL_VARBINARY:
  case SQL_LONGVARBINARY: {
- // TODO: revisit
- HandleZeroColumnSizeAtFetch(columnSize);
- std::unique_ptr<SQLCHAR[]> dataBuffer(new SQLCHAR[columnSize]);
- SQLLEN dataLen;
- ret = SQLGetData_ptr(hStmt, i, SQL_C_BINARY, dataBuffer.get(), columnSize, &dataLen);
+ // Use streaming for large VARBINARY (columnSize unknown or > 8000)
+ if (columnSize == SQL_NO_TOTAL || columnSize == 0 || columnSize > 8000) {
+ LOG("Streaming LOB for column {} (VARBINARY)", i);
+ row.append(FetchLobColumnData(hStmt, i, SQL_C_BINARY, false, true));
+ } else {
+ // Small VARBINARY, fetch directly
+ std::vector<SQLCHAR> dataBuffer(columnSize);
+ SQLLEN dataLen;
+ ret = SQLGetData_ptr(hStmt, i, SQL_C_BINARY, dataBuffer.data(), columnSize, &dataLen);
 
- if (SQL_SUCCEEDED(ret)) {
- // TODO: Refactor these if's across other switches to avoid code duplication
- if (dataLen > 0) {
-if (static_cast<size_t>(dataLen) <= columnSize) {
- row.append(py::bytes(reinterpret_cast<const char*>(
- dataBuffer.get()), dataLen));
-} else {
- // In this case, buffer size is smaller, and data to be retrieved is longer
- // TODO: Revisit
+ if (SQL_SUCCEEDED(ret)) {
+ if (dataLen > 0) {
+ if (static_cast<size_t>(dataLen) <= columnSize) {
+ row.append(py::bytes(reinterpret_cast<const char*>(dataBuffer.data()), dataLen));
+ } else {
+ LOG("VARBINARY column {} data truncated, using streaming LOB", i);
+ row.append(FetchLobColumnData(hStmt, i, SQL_C_BINARY, false, true));
+ }
+ } else if (dataLen == SQL_NULL_DATA) {
+ row.append(py::none());
+ } else if (dataLen == 0) {
+ row.append(py::bytes(""));
+ } else {
  std::ostringstream oss;
- oss << "Buffer length for fetch (" << columnSize << ") is smaller, & data "
- << "to be retrieved is longer (" << dataLen << "). ColumnID - "
-  << i << ", datatype - " << dataType;
+ oss << "Unexpected negative length (" << dataLen << ") returned by SQLGetData. ColumnID=" 
+ << i << ", dataType=" << dataType << ", bufferSize=" << columnSize;
+ LOG("Error: {}", oss.str());
  ThrowStdException(oss.str());
  }
- } else if (dataLen == SQL_NULL_DATA) {
- row.append(py::none());
- } else if (dataLen == 0) {
- // Empty bytes
- row.append(py::bytes(""));
- } else if (dataLen < 0) {
- // This is unexpected
- LOG("SQLGetData returned an unexpected negative data length. "
- "Raising exception. Column ID - {}, Data Type - {}, Data Length - {}",
- i, dataType, dataLen);
- ThrowStdException("SQLGetData returned an unexpected negative data length");
+ } else {
+ LOG("Error retrieving VARBINARY data for column {}. SQLGetData rc = {}", i, ret);
+ row.append(py::none());
  }
-} else {
-LOG("Error retrieving data for column - {}, data type - {}, SQLGetData return "
-"code - {}. Returning NULL value instead",
-i, dataType, ret);
-row.append(py::none());
-}
+ }
  break;
  }
  case SQL_TINYINT: {

tests/test_004_cursor.py

@@ -6113,34 +6113,71 @@ def test_binary_data_over_8000_bytes(cursor, db_connection):
  drop_table_if_exists(cursor, "#pytest_small_binary")
  db_connection.commit()
 
-def test_binary_data_large(cursor, db_connection):
- """Test insertion of binary data larger than 8000 bytes with streaming support."""
+def test_varbinarymax_insert_fetch(cursor, db_connection):
+ """Test for VARBINARY(MAX) insert and fetch (streaming support) using execute per row"""
  try:
- drop_table_if_exists(cursor, "#pytest_large_binary")
+ # Create test table
+ drop_table_if_exists(cursor, "#pytest_varbinarymax")
  cursor.execute("""
- CREATE TABLE #pytest_large_binary (
- id INT PRIMARY KEY,
- large_binary VARBINARY(MAX)
+ CREATE TABLE #pytest_varbinarymax (
+ id INT,
+ binary_data VARBINARY(MAX)
  )
  """)
- 
- # Large binary data > 8000 bytes
- large_data = b'A' * 10000 # 10 KB
- cursor.execute("INSERT INTO #pytest_large_binary (id, large_binary) VALUES (?, ?)", (1, large_data))
+
+ # Prepare test data
+ test_data = [
+ (2, b''), # Empty bytes
+ (3, b'1234567890'), # Small binary
+ (4, b'A' * 9000), # Large binary > 8000 (streaming)
+ (5, b'B' * 20000), # Large binary > 8000 (streaming)
+ (6, b'C' * 8000), # Edge case: exactly 8000 bytes
+ (7, b'D' * 8001), # Edge case: just over 8000 bytes
+ ]
+
+ # Insert each row using execute
+ for row_id, binary in test_data:
+ cursor.execute("INSERT INTO #pytest_varbinarymax VALUES (?, ?)", (row_id, binary))
  db_connection.commit()
- print("Inserted large binary data (>8000 bytes) successfully.")
- 
- # commented out for now
- # cursor.execute("SELECT large_binary FROM #pytest_large_binary WHERE id=1")
- # result = cursor.fetchone()
- # assert result[0] == large_data, f"Large binary data mismatch, got {len(result[0])} bytes"
- 
- # print("Large binary data (>8000 bytes) inserted and verified successfully.")
- 
+
+ # ---------- FETCHONE TEST (multi-column) ----------
+ cursor.execute("SELECT id, binary_data FROM #pytest_varbinarymax ORDER BY id")
+ rows = []
+ while True:
+ row = cursor.fetchone()
+ if row is None:
+ break
+ rows.append(row)
+
+ assert len(rows) == len(test_data), f"Expected {len(test_data)} rows, got {len(rows)}"
+
+ # Validate each row
+ for i, (expected_id, expected_data) in enumerate(test_data):
+ fetched_id, fetched_data = rows[i]
+ assert fetched_id == expected_id, f"Row {i+1} ID mismatch: expected {expected_id}, got {fetched_id}"
+ assert isinstance(fetched_data, bytes), f"Row {i+1} expected bytes, got {type(fetched_data)}"
+ assert fetched_data == expected_data, f"Row {i+1} data mismatch"
+
+ # ---------- FETCHALL TEST ----------
+ cursor.execute("SELECT id, binary_data FROM #pytest_varbinarymax ORDER BY id")
+ all_rows = cursor.fetchall()
+ assert len(all_rows) == len(test_data)
+
+ # ---------- FETCHMANY TEST ----------
+ cursor.execute("SELECT id, binary_data FROM #pytest_varbinarymax ORDER BY id")
+ batch_size = 2
+ batches = []
+ while True:
+ batch = cursor.fetchmany(batch_size)
+ if not batch:
+ break
+ batches.extend(batch)
+ assert len(batches) == len(test_data)
+
  except Exception as e:
- pytest.fail(f"Large binary data insertion test failed: {e}")
+ pytest.fail(f"VARBINARY(MAX) insert/fetch test failed: {e}")
  finally:
- drop_table_if_exists(cursor, "#pytest_large_binary")
+ drop_table_if_exists(cursor, "#pytest_varbinarymax")
  db_connection.commit()
 
 
@@ -6303,6 +6340,40 @@ def test_binary_mostly_small_one_large(cursor, db_connection):
  drop_table_if_exists(cursor, "#pytest_mixed_size_binary")
  db_connection.commit()
 
+def test_varbinarymax_insert_fetch_null(cursor, db_connection):
+ """Test insertion and retrieval of NULL value in VARBINARY(MAX) column."""
+ try:
+ drop_table_if_exists(cursor, "#pytest_varbinarymax_null")
+ cursor.execute("""
+ CREATE TABLE #pytest_varbinarymax_null (
+ id INT,
+ binary_data VARBINARY(MAX)
+ )
+ """)
+
+ # Insert a row with NULL for binary_data
+ cursor.execute(
+ "INSERT INTO #pytest_varbinarymax_null VALUES (?, CAST(NULL AS VARBINARY(MAX)))",
+ (1,)
+ )
+ db_connection.commit()
+
+ # Fetch the row
+ cursor.execute("SELECT id, binary_data FROM #pytest_varbinarymax_null")
+ row = cursor.fetchone()
+
+ assert row is not None, "No row fetched"
+ fetched_id, fetched_data = row
+ assert fetched_id == 1, "ID mismatch"
+ assert fetched_data is None, "Expected NULL for binary_data"
+
+ except Exception as e:
+ pytest.fail(f"VARBINARY(MAX) NULL insert/fetch test failed: {e}")
+
+ finally:
+ drop_table_if_exists(cursor, "#pytest_varbinarymax_null")
+ db_connection.commit()
+
 def test_only_null_and_empty_binary(cursor, db_connection):
  """Test table with only NULL and empty binary values to ensure fallback doesn't produce size=0"""
  try: