diff --git a/CHANGELOG.md b/CHANGELOG.md index 2cac01d4d..edac2ed50 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -57,12 +57,37 @@ 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. +- 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 40fc409e7..c9b700769 100644 --- a/README.md +++ b/README.md @@ -170,6 +170,69 @@ 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 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 +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 d4b85e56b..75dc24f15 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 @@ -47,10 +48,37 @@ 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") 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 + + +def _encoded_length_exceeds(value: str, encoding: str, limit: int) -> bool: + encoder = codecs.getincrementalencoder(encoding)(errors="strict") + total = 0 + chunk_size = 4096 + 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 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) @@ -406,9 +434,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 @@ -962,7 +988,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, @@ -1170,6 +1196,26 @@ 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: + """Default fetchmany row count: an integer from 1 to 1,000,000, not bool.""" + 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. @@ -1208,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 @@ -1230,7 +1283,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." ) @@ -1238,12 +1295,29 @@ 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 + 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." + f"Invalid column size: {column_size}. Must be a non-negative integer " + f"no greater than {MAX_NATIVE_PARAMETER_SIZE}." ) - 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." @@ -1262,7 +1336,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." ) @@ -2459,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", @@ -2535,6 +2618,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: @@ -2549,6 +2633,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__") @@ -2570,12 +2655,47 @@ 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 - if sample_value is not None: - if isinstance(sample_value, str) and column_size > MAX_INLINE_CHAR: - is_dae = True - elif isinstance(sample_value, (bytes, bytearray)) and column_size > 8000: - is_dae = True + # 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, + ) + 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, + ) + 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)) + ] + 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( + _utf16_length_exceeds(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 + ) + 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 ( @@ -2590,6 +2710,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 @@ -2599,14 +2720,18 @@ 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 > MAX_INLINE_BINARY: + 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 # 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 @@ -2671,6 +2796,20 @@ 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 _utf16_length_exceeds(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, @@ -2685,7 +2824,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 @@ -2694,7 +2833,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: @@ -2756,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 @@ -2764,9 +2905,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 @@ -2900,7 +3038,15 @@ 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. + 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. + ValueError: If size is negative or exceeds 1,000,000. Returns: List of Row objects. @@ -2911,8 +3057,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 [] if self._cached_decoding_generation != self._connection._decoding_generation: @@ -3055,17 +3203,24 @@ 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 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. 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. """ 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() @@ -3074,7 +3229,7 @@ def arrow_batch(self, batch_size: int = 8192) -> "pyarrow.RecordBatch": char_decoding = self._get_decoding_settings(ddbc_sql_const.SQL_CHAR.value) char_c_type = char_decoding.get("ctype", ddbc_sql_const.SQL_WCHAR.value) ret = ddbc_bindings.DDBCSQLFetchArrowBatch( - self.hstmt, capsules, max(batch_size, 0), char_c_type, self.messages + self.hstmt, capsules, batch_size, char_c_type, self.messages ) check_error(ddbc_sql_const.SQL_HANDLE_STMT.value, self.hstmt, ret) @@ -3099,7 +3254,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: 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. Returns: A pyarrow Table containing all remaining rows from the result set. @@ -3110,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 @@ -3120,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 @@ -3138,7 +3296,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: 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. Returns: A pyarrow-compatible RecordBatchReader for the result set. @@ -3461,12 +3622,14 @@ 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__}" ) - 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__}") @@ -3526,7 +3689,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. @@ -3715,7 +3880,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, @@ -3949,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. @@ -3996,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( @@ -4010,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 4e4523955..a6bcdf45f 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -23,6 +23,7 @@ #include // For std::memcpy #include #include +#include #include // std::forward #include // CPython datetime API (PyDateTime_IMPORT, PyDateTime_GET_*, etc.) @@ -325,12 +326,242 @@ 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; +} + +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; + +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) { + ThrowStdException(std::string(name) + " must be between " + std::to_string(minimum) + + " and " + std::to_string(MAX_NATIVE_ROW_COUNT)); + } +} + +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"); + 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"); + } + return 0; +} + +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)) { + ThrowStdException(errorMessage); + } + 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) { + throw NativeFetchLimit(); + } + reservedBytes = totalBytes; +} + +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"); + 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 = newReservedBytes - oldBytes; +} + +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"); + } + 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); +} + +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(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; + 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) { + 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) + "'"; @@ -347,6 +578,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); @@ -357,6 +589,94 @@ 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 put_data_fn(nullptr, 0); + } + 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(); + 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)) { + 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"); + return SQL_ERROR; +} + // 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) { @@ -502,6 +822,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"); + } + SQLRETURN describeRc = PreResolveUdtTypes(hStmt, paramInfos); if (!SQL_SUCCEEDED(describeRc)) return describeRc; // GH-627: resolve unknown NULL param SQL types before binding any param. @@ -509,6 +833,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, @@ -541,7 +868,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", @@ -581,8 +909,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) { @@ -616,8 +943,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) { @@ -1626,7 +1952,7 @@ void SqlHandle::markImplicitlyFreed() { _type); return; // Refuse to mark - let normal free() handle it } - resultMetadata.clear(); + clearResult(); _implicitly_freed = true; } @@ -1643,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 @@ -1684,7 +2010,7 @@ SQLRETURN SqlHandle::freeHandle() { } void SqlHandle::close_cursor() { - resultMetadata.clear(); + clearResult(); if (is_python_finalizing()) { return; } @@ -1712,7 +2038,7 @@ void SqlHandle::close_cursor() { } void SqlHandle::cancel() { - resultMetadata.clear(); + clearResult(); if (is_python_finalizing()) { return; } @@ -1747,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(); } @@ -1769,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"); } @@ -1782,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"); } @@ -1807,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"); } @@ -1840,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"); } @@ -1863,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"); } @@ -1886,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"); } @@ -2041,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()); @@ -2078,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(); @@ -2121,11 +2447,12 @@ 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; } - statementHandle->resultMetadata.clear(); + statementHandle->clearResult(); SQLHANDLE hStmt = statementHandle->get(); // Configure forward-only / read-only cursor (matches slow path semantics). @@ -2156,7 +2483,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(); @@ -2195,10 +2523,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; @@ -2217,58 +2541,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; } @@ -2290,6 +2564,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; @@ -2298,6 +2575,12 @@ SQLRETURN BindParameterArray(SqlHandle& handle, SQLHANDLE hStmt, const py::list& if (!SQL_SUCCEEDED(describeRc)) return describeRc; // 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]; @@ -2359,8 +2642,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()) { @@ -2368,16 +2661,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 = + EncodeUnicodeValue(columnValues[i].ptr(), "utf-16-le"); + 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(encodedData, static_cast(encodedSize), + reinterpret_cast(destination)); + destination[wcharLength] = 0; strLenOrIndArray[i] = SQL_NTS; } } @@ -2385,7 +2703,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: @@ -2456,8 +2774,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()) { @@ -2465,19 +2790,30 @@ SQLRETURN BindParameterArray(SqlHandle& handle, SQLHANDLE hStmt, const py::list& std::memset(charArray + i * (info.columnSize + 1), 0, info.columnSize + 1); } else { - std::string encodedStr; - + if (info.paramCType == SQL_C_BINARY && + !py::isinstance(columnValues[i]) && + !py::isinstance(columnValues[i])) { + ThrowStdException(MakeParamMismatchErrorStr(info.paramCType, + paramIndex)); + } + 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 = - columnValues[i].attr("encode")(charEncoding, "strict"); - encodedStr = encoded.cast(); + encoded = + EncodeUnicodeValue(columnValues[i].ptr(), charEncoding); + 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 " @@ -2488,29 +2824,41 @@ 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 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 { - // bytes/bytearray - use as-is (already encoded) - encodedStr = columnValues[i].cast(); + ThrowStdException( + MakeParamMismatchErrorStr(info.paramCType, paramIndex)); } - 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 - " "param_index=%d", paramIndex); dataPtr = charArray; - bufferLength = info.columnSize + 1; + bufferLength = static_cast(elementWidth); break; } case SQL_C_BIT: { @@ -2922,10 +3270,15 @@ 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"); - statementHandle->resultMetadata.clear(); + const size_t paramSetSize = static_cast( + ValidateNativeRowCountArgument(paramSetSizeArg, "Parameter set size", false)); + if (!statementHandle || !statementHandle->get()) { + return SQL_INVALID_HANDLE; + } + statementHandle->clearResult(); LOG("SQLExecuteMany: Starting batch execution - param_count=%zu, " "param_set_size=%zu", columnwise_params.size(), paramSetSize); @@ -2992,14 +3345,34 @@ 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 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); @@ -3015,7 +3388,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; @@ -3030,48 +3403,18 @@ SQLRETURN SQLExecuteMany_wrap(const SqlHandlePtr statementHandle, const std::u16 dae_chunk_count, rc); 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; + if (rc != SQL_NEED_DATA) break; + + const ParamInfo* matchedInfo = reinterpret_cast(token); + const ParamInfo* first = rowParamInfos.data(); + const ParamInfo* last = first + rowParamInfos.size(); + if (matchedInfo < first || matchedInfo >= last) { + ThrowStdException( + "SQLExecuteMany: unrecognized paramToken from SQLParamData"); } + 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, " @@ -3082,6 +3425,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", @@ -3213,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"); } @@ -3232,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"); @@ -3240,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; @@ -3257,6 +3631,9 @@ 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; int loopCount = 0; @@ -3326,7 +3703,9 @@ py::object FetchLobColumnData(SQLHSTMT hStmt, SQLUSMALLINT colIndex, SQLSMALLINT } } if (bytesRead > 0) { - buffer.insert(buffer.end(), chunk.begin(), chunk.begin() + bytesRead); + 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); } if (ret == SQL_SUCCESS) { @@ -3334,35 +3713,42 @@ 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) { - size_t wcharCount = buffer.size() / sizeof(SQLWCHAR); - std::vector alignedBuf(wcharCount); - std::memcpy(alignedBuf.data(), buffer.data(), buffer.size()); - return py::cast(dupeSqlWCharAsUtf16Le(alignedBuf.data(), wcharCount)); + 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(dataSize), + "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 " "%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", @@ -3442,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); @@ -4230,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)); } @@ -4247,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; @@ -4260,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) { @@ -4277,9 +4684,18 @@ SQLRETURN SQLFetchScroll_wrap(SqlHandlePtr StatementHandle, SQLSMALLINT FetchOri // TODO: Move to anonymous namespace, since it is not used outside this file template SQLRETURN SQLBindColums(SQLHSTMT hStmt, ColumnBuffers& buffers, const Metadata& columnNames, - SQLUSMALLINT numCols, int fetchSize, int charCtype = SQL_C_WCHAR, - py::handle messages = {}) { + SQLUSMALLINT numCols, int fetchSize, size_t& reservedBytes, + int charCtype = SQL_C_WCHAR, py::handle messages = {}) { 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 @@ -4293,25 +4709,39 @@ 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*/; - buffers.wcharBuffers[col - 1].resize(fetchSize * fetchBufferSize); + 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 - 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()); + CheckedFetchBufferLength(fetchBufferSize, sizeof(SQLCHAR)), + buffers.indicators[col - 1].data()); } break; } @@ -4321,49 +4751,57 @@ 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*/; - buffers.wcharBuffers[col - 1].resize(fetchSize * fetchBufferSize); + 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; } 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()); @@ -4371,50 +4809,58 @@ SQLRETURN SQLBindColums(SQLHSTMT hStmt, ColumnBuffers& buffers, const Metadata& 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; 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); - buffers.charBuffers[col - 1].resize(fetchSize * columnSize); + const size_t fetchBufferSize = CheckedFetchColumnSize(columnSize); + ResizeNativeFetchBuffer(buffers.charBuffers[col - 1], + 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: - 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, + sizeof(DateTimeOffset), buffers.indicators[col - 1].data()); break; default: @@ -4437,37 +4883,19 @@ SQLRETURN SQLBindColums(SQLHSTMT hStmt, ColumnBuffers& buffers, const Metadata& } CaptureFetchDiagnostics(hStmt, ret, messages); } + unbindOnFailure.active = false; return ret; } // 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; - { - // 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 @@ -4525,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++) { @@ -4576,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: @@ -4613,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 @@ -4801,83 +5242,126 @@ 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; } // 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(); @@ -4888,6 +5372,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; } @@ -4948,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 @@ -4965,14 +5453,19 @@ 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"); + 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); + 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(); @@ -5055,11 +5548,18 @@ 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); FetchStateGuard fetchStateGuard(StatementHandle, messages); // Bind columns - ret = SQLBindColums(hStmt, buffers, columnNames, numCols, fetchSize, charCtype, messages); + ret = SQLBindColums(hStmt, buffers, columnNames, numCols, fetchSize, reservedBytes, charCtype, + messages); if (!SQL_SUCCEEDED(ret)) { LOG("FetchMany_wrap: Error when binding columns - SQLRETURN=%d", ret); return ret; @@ -5094,9 +5594,10 @@ 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, py::handle messages) { - size_t start = 0; - size_t end = 0; + std::vector& dataVec, SQLLEN* indicator, size_t& reservedBytes, + 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; @@ -5114,10 +5615,15 @@ 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) { + // 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), @@ -5128,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; } @@ -5143,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; } @@ -5151,18 +5659,20 @@ 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 - dataVec.resize(end); } else { // Unexpected return code return ret; @@ -5172,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) @@ -5183,14 +5862,31 @@ 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"); + 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 = + 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"); + 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 = 64; + // An overly large fetch size doesn't seem to help performance. + int fetchSize = pending ? pending->fetchSize + : arrowBatchSize > 0 ? std::min(arrowBatchSize, 64) : 1; SQLRETURN ret = SQL_ERROR; ResultMetadataFailureGuard metadataFailure(StatementHandle->resultMetadata, ret); @@ -5212,10 +5908,10 @@ 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); + size_t reservedBytes = pending ? pending->reservedBytes : 0; std::vector> arrowArrayPrivateData(numCols); std::vector> arrowSchemaPrivateData(numCols); @@ -5230,7 +5926,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 || @@ -5259,8 +5954,11 @@ 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 = + AllocateArrowArray(offsetCount, reservedBytes, + "Arrow offset buffer is too large"); + EnsureNativeFetchBufferSize(arrowColumnProducer->varData, initialVarDataSize, + reservedBytes); columnVarLen[i] = true; // start at offset 0 arrowColumnProducer->varVal[0] = 0; @@ -5271,8 +5969,11 @@ 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 = + AllocateArrowArray(offsetCount, reservedBytes, + "Arrow offset buffer is too large"); + EnsureNativeFetchBufferSize(arrowColumnProducer->varData, initialVarDataSize, + reservedBytes); columnVarLen[i] = true; // start at offset 0 arrowColumnProducer->varVal[0] = 0; @@ -5280,33 +5981,45 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, break; case SQL_TINYINT: format = "C"; - arrowColumnProducer->uint8Val = std::make_unique(arrowBatchSize); + arrowColumnProducer->uint8Val = + AllocateArrowArray(batchSize, reservedBytes, + "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 = + AllocateArrowArray(batchSize, reservedBytes, + "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 = + AllocateArrowArray(batchSize, reservedBytes, + "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 = + AllocateArrowArray(batchSize, reservedBytes, + "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 = + 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 = std::make_unique(arrowBatchSize); + arrowColumnProducer->float64Val = + AllocateArrowArray(batchSize, reservedBytes, + "Arrow value buffer is too large"); arrowColumnProducer->ptrValueBuffer = arrowColumnProducer->float64Val.get(); break; case SQL_DECIMAL: @@ -5319,7 +6032,9 @@ 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 = + AllocateArrowArray(batchSize, reservedBytes, + "Arrow value buffer is too large"); arrowColumnProducer->ptrValueBuffer = arrowColumnProducer->decimalVal.get(); break; } @@ -5327,28 +6042,38 @@ 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 = + 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 = std::make_unique(arrowBatchSize); + arrowColumnProducer->tsMicroVal = + AllocateArrowArray(batchSize, reservedBytes, + "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 = + AllocateArrowArray(batchSize, reservedBytes, + "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 = + AllocateArrowArray(batchSize, reservedBytes, + "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 = + AllocateArrowArray(bitmapSize, reservedBytes, + "Arrow bitmap is too large"); + std::memset(arrowColumnProducer->bitVal.get(), 0, bitmapSize); arrowColumnProducer->ptrValueBuffer = arrowColumnProducer->bitVal.get(); break; default: @@ -5369,285 +6094,143 @@ 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 = + AllocateArrowArray(bitmapSize, reservedBytes, "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); + } + + 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; + 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 (!pending && bufferFetchSize > 0) { + ReserveNativeFetchBytes( + reservedBytes, + CheckedMultiplySize(static_cast(numCols), + static_cast(bufferFetchSize), + "Native fetch indicator buffer is too large"), + sizeof(SQLLEN)); + } // Initialize column buffers - ColumnBuffers buffers(numCols, fetchSize); - 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 && fetchSize > 0) { - ret = SQLBindColums(hStmt, buffers, columnNames, numCols, fetchSize, charCtype, 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]; - auto columnSize = columnSizes[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(), 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(), 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(), 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; + 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_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; - } - 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]; @@ -5695,16 +6278,23 @@ 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]; - while (target_vec->size() < start + dataLen) { - target_vec->resize(target_vec->size() * 2); - } + 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][idxRowSql * fetchBufferSize], - dataLen); + std::memcpy(&(*target_vec)[start], + &buffers.charBuffers[idxCol][sourceOffset], dataLen); + } arrowColumnProducer->varVal[idxRowArrow + 1] = start + dataLen; break; } @@ -5712,20 +6302,25 @@ 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]; - while (target_vec->size() < start + dataLen) { - target_vec->resize(target_vec->size() * 2); + 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][idxRowSql * fetchBufferSize], - dataLen); arrowColumnProducer->varVal[idxRowArrow + 1] = start + dataLen; break; } @@ -5736,21 +6331,22 @@ 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)); 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 bytesWritten = simdutf::convert_utf16le_to_utf8_with_replacement( utf16Source, dataLenW, reinterpret_cast(target_vec->data() + start)); @@ -5766,9 +6362,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]; @@ -5906,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 @@ -6121,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(); @@ -6193,14 +6801,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 @@ -6231,14 +6837,23 @@ 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; + ReserveNativeFetchBytes( + reservedBytes, + CheckedMultiplySize(static_cast(numCols), static_cast(fetchSize), + "Native fetch indicator buffer is too large"), + sizeof(SQLLEN)); ColumnBuffers buffers(numCols, fetchSize); SQLULEN numRowsFetched = 0; FetchStateGuard fetchStateGuard(StatementHandle, messages); // Bind columns - ret = SQLBindColums(hStmt, buffers, columnNames, numCols, fetchSize, charCtype, messages); + ret = SQLBindColums(hStmt, buffers, columnNames, numCols, fetchSize, reservedBytes, charCtype, + messages); if (!SQL_SUCCEEDED(ret)) { LOG("FetchAll_wrap: Error when binding columns - SQLRETURN=%d", ret); return ret; @@ -6285,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(); @@ -6322,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 " @@ -6530,6 +7156,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/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/mssql_python/pybind/param_detect.hpp b/mssql_python/pybind/param_detect.hpp index 172a98b51..7f71125ec 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_CHAR ? 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 actualSize; + default: + return actualSize; + } +} + // SQL Server maximum numeric precision inline constexpr int MAX_NUMERIC_PRECISION = 38; @@ -190,6 +210,45 @@ 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 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_BINARY - total) { + return MAX_INLINE_BINARY + 1; + } + total += encodedSize; + } + if (length == 0) { + py::object encoded = encoder.attr("encode")(py::str(), true); + total = PyBytes_GET_SIZE(encoded.ptr()); + } + return total; +} + 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(); @@ -215,8 +274,91 @@ 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) { + ParamInfo& info, const std::string& charEncoding) { py::tuple values = borrow(inputSize); info.paramSQLType = values[0].cast(); info.paramCType = values[1].cast(); @@ -257,16 +399,51 @@ 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)); + 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"); + } 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 + ? 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 = + (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 (!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 (info.isDAE) { info.dataPtr = borrow(obj); } @@ -296,9 +473,15 @@ 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(); + 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); @@ -312,7 +495,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; } @@ -398,16 +581,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. @@ -439,16 +613,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. - // 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) + // 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 = 0; + 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; @@ -467,7 +639,8 @@ inline std::vector DetectParamTypes(PyObject* params, PyObject* input info.decimalDigits = 0; if (length > MAX_INLINE_BINARY) { info.isDAE = true; - info.columnSize = 0; + info.columnSize = + DAEColumnSize(SQL_VARBINARY, static_cast(length)); info.dataPtr = borrow(obj); } else { info.columnSize = std::max(length, 1); diff --git a/tests/test_004_cursor.py b/tests/test_004_cursor.py index 001440edf..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 @@ -1693,6 +1696,105 @@ def test_arraysize(cursor): assert cursor.arraysize == 5, "Arraysize mismatch after change" +@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"): + 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_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(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(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): + 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_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 [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 37a336789..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]), @@ -225,11 +806,25 @@ 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): - 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 diff --git a/tests/test_010_pybind_functions.py b/tests/test_010_pybind_functions.py index 106b64ca3..41e175af4 100644 --- a/tests/test_010_pybind_functions.py +++ b/tests/test_010_pybind_functions.py @@ -37,6 +37,71 @@ ) +@pytest.mark.skipif(not DDBC_AVAILABLE, reason="ddbc_bindings not available") +@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, 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) + + +@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") +@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.""" diff --git a/tests/test_023_execute_path_parity.py b/tests/test_023_execute_path_parity.py index d9bada323..6eae78c7d 100644 --- a/tests/test_023_execute_path_parity.py +++ b/tests/test_023_execute_path_parity.py @@ -296,6 +296,140 @@ 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"), + [ + ("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), + ], +) +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( + ("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.""" + 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(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): + """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 [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_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): + return 1 + + def __iter__(self): + return iter("x") + + def __getitem__(self, key): + return "x" + + def encode(self, *args, **kwargs): + return b"x" * 100_000 + + 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"),), (large_value,)], + ) + 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): + """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.""" @@ -330,6 +464,80 @@ 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) + + +@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 diff --git a/tests/test_024_bulkcopy_arrow.py b/tests/test_024_bulkcopy_arrow.py index a840853ba..c8718284b 100644 --- a/tests/test_024_bulkcopy_arrow.py +++ b/tests/test_024_bulkcopy_arrow.py @@ -202,6 +202,15 @@ 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"): + _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")