diff --git a/mssql_python/constants.py b/mssql_python/constants.py index 401e434a..d6b70573 100644 --- a/mssql_python/constants.py +++ b/mssql_python/constants.py @@ -533,6 +533,19 @@ def get_attribute_set_timing(attribute): # (with spaces) ODBC only honors "PacketSize" without spaces # internally. "packetsize": "PacketSize", + # Multiple Active Result Sets (MARS) — lets a single connection have more + # than one active statement/result set at a time. Required for concurrent + # execute_async on cursors sharing one Connection. + # + # NOTE: Only the ODBC-standard alias "MARS_Connection" is on the allowlist. + # The SQL Server-native spelling "MultipleActiveResultSets" is silently + # ignored by the bundled msodbcsql18 driver (empirically verified: the + # keyword is accepted but MARS is not actually enabled). Do NOT add + # "multipleactiveresultsets" here without confirming the driver honors it + # first — otherwise users get a MARS-off connection while thinking they + # asked for MARS-on, which leads to unpredictable behavior under + # concurrent async workloads. + "mars_connection": "MARS_Connection", } # Canonical normalized key names produced by _ConnectionStringParser._normalize_params. diff --git a/mssql_python/cursor.py b/mssql_python/cursor.py index 85701a40..9c1d4c36 100644 --- a/mssql_python/cursor.py +++ b/mssql_python/cursor.py @@ -163,6 +163,13 @@ def __init__(self, connection: "Connection", timeout: int = 0) -> None: ) self.messages: List[Tuple[str, str]] = [] # Store diagnostic messages + # Async POC: cursor-busy guard. Set True while an execute_async / + # fetch_async is in flight so overlapping async calls on the same + # cursor raise ProgrammingError instead of racing on the HSTMT. + # DB-API threadsafety=1 already forbids sharing cursors across + # threads; this adds the same protection for asyncio tasks. + self._async_in_flight: bool = False + def _is_unicode_string(self, param: str) -> bool: """ Check if a string contains non-ASCII characters. @@ -3486,3 +3493,479 @@ def setoutputsize(self, size: int, column: Optional[int] = None) -> None: are managed automatically by the underlying driver. """ # This is a no-op - buffer sizes are managed automatically + + # ======================================================================== + # Async POC: execute_async + fetch_async + # ======================================================================== + # These methods live in an isolated block and do NOT share code with the + # sync execute / fetchone / fetchmany / fetchall paths. The prep and + # finalize helpers below are DELIBERATELY DUPLICATED from the sync + # implementations so any bug fix or behavior change on the sync path does + # not need to propagate here (and vice versa). Long-term de-duplication is + # deferred until the async POC stabilizes. + # + # Design notes: + # - Statement-level polling via SQL_ATTR_ASYNC_ENABLE, driven from + # C++ (DDBCSQLExecuteAsync / DDBCSQLFetchOneAsync / + # DDBCSQLFetchManyAsync / DDBCSQLFetchAllAsync). + # - The blocking polling loop runs under loop.run_in_executor so the + # asyncio event loop in the caller's thread stays responsive. + # - Data-At-Execution parameters (large VARBINARY(MAX) / NVARCHAR(MAX)) + # are rejected up-front in C++ for async execute; use sync execute() + # for those. + # - Concurrency: one Cursor == one HSTMT, so async ops on the same + # cursor are serialized via ``_async_in_flight``. Two cursors on the + # same Connection can run concurrently. + # ------------------------------------------------------------------------ + + def _check_async_capable(self) -> None: + """Raise NotSupportedError if the ODBC driver doesn't advertise async support. + + Called at the top of every async method. Result is cached inside + Connection::isAsyncCapable() so the SQLGetInfo probe runs at most once + per connection. + """ + conn = self._connection._conn + if not conn.is_async_capable(): + raise NotSupportedError( + driver_error=( + "Async execution is not supported by this ODBC driver. " + "SQLGetInfo(SQL_ASYNC_MODE) reports SQL_AM_NONE." + ), + ddbc_error="", + ) + + def _acquire_async_slot(self) -> None: + """Take the cursor-busy slot; raise ProgrammingError if already held.""" + if self._async_in_flight: + raise ProgrammingError( + driver_error=( + "Cursor is busy: another async operation is already in flight " + "on this cursor. One Cursor == one HSTMT, so async ops on the " + "same cursor must be serialized. Create a second cursor for " + "concurrent async work." + ), + ddbc_error="", + ) + self._async_in_flight = True + + def _release_async_slot(self) -> None: + """Release the cursor-busy slot. Safe to call multiple times.""" + self._async_in_flight = False + + # ---- Duplicated prep / finalize helpers (mirror the sync execute()) ---- + + def _prepare_execute_state_async( # pylint: disable=too-many-locals,too-many-branches,too-many-statements + self, + operation: str, + parameters, + use_prepare: bool, + reset_cursor: bool, + ) -> Tuple[str, List[Any], List[Any], bool, Any]: + """Async-only mirror of the pre-ODBC prep block from sync execute(). + + DELIBERATE DUPLICATE of the corresponding logic in ``execute()`` — do + NOT refactor to call into the sync path. See the block comment above + for rationale. + + Returns a tuple: (operation, parameters, parameters_type, + effective_use_prepare, encoding_settings). + """ + self._check_closed() + if reset_cursor: + if self.hstmt: + self._soft_reset_cursor() + else: + self._reset_cursor() + else: + if self.hstmt: + logger.debug( + "execute_async: Closing cursor for re-execution (reset_cursor=False)" + ) + self.hstmt._close_cursor() + self._clear_rownumber() + + # Clear any previous messages + self.messages = [] + + # Parameter unwrap (same rules as sync execute). + if parameters: + if isinstance(parameters, tuple) and len(parameters) == 1: + if isinstance(parameters[0], (tuple, list, dict)): + actual_params = parameters[0] + elif isinstance(parameters[0], Row): + actual_params = tuple(parameters[0]) + else: + actual_params = parameters + else: + actual_params = parameters + + if operation == self.last_executed_stmt and isinstance( + actual_params, (tuple, list) + ): + parameters = list(actual_params) + else: + operation, converted_params = detect_and_convert_parameters( + operation, actual_params + ) + parameters = list(converted_params) + else: + parameters = [] + + encoding_settings = self._get_encoding_settings() + + logger.debug("execute_async: Creating parameter type list") + param_info = ddbc_bindings.ParamInfo + parameters_type: List[Any] = [] + + if parameters and self._inputsizes: + if len(self._inputsizes) != len(parameters): + warnings.warn( + f"Number of input sizes ({len(self._inputsizes)}) does not match " + f"number of parameters ({len(parameters)}). " + f"This may lead to unexpected behavior.", + Warning, + ) + + if parameters: + for i, param in enumerate(parameters): + paraminfo = self._create_parameter_types_list( + param, param_info, parameters, i + ) + parameters_type.append(paraminfo) + + same_sql = ( + parameters + and operation == self.last_executed_stmt + and self.is_stmt_prepared[0] + ) + if not same_sql: + self.is_stmt_prepared = [False] + effective_use_prepare = use_prepare and not same_sql + + return ( + operation, + parameters, + parameters_type, + effective_use_prepare, + encoding_settings, + ) + + def _finalize_execute_async(self, ret: int, operation: str) -> None: + """Async-only mirror of the post-ODBC finalize block from sync execute(). + + DELIBERATE DUPLICATE of the corresponding logic in ``execute()``. + """ + try: + check_error(ddbc_sql_const.SQL_HANDLE_STMT.value, self.hstmt, ret) + except Exception as e: # pylint: disable=broad-exception-caught + logger.warning("execute_async failed, resetting cursor: %s", e) + self._reset_cursor() + raise + + self._capture_diagnostics(ret) + self.last_executed_stmt = operation + self.rowcount = ddbc_bindings.DDBCSQLRowCount(self.hstmt) + + column_metadata: List[Any] = [] + try: + ddbc_bindings.DDBCSQLDescribeCol(self.hstmt, column_metadata) + self._initialize_description(column_metadata) + except Exception: # pylint: disable=broad-exception-caught + # If describe fails, it's likely there are no results (e.g. INSERT) + self.description = None + + if self.description: + self.rowcount = -1 + self._reset_rownumber() + self._cached_column_map = { + col_desc[0]: i for i, col_desc in enumerate(self.description) + } + self._cached_column_map_lower = ( + {k.lower(): v for k, v in self._cached_column_map.items()} + if get_settings().lowercase + else None + ) + self._cached_converter_map = self._build_converter_map() + self._uuid_str_indices = self._compute_uuid_str_indices() + else: + self.rowcount = ddbc_bindings.DDBCSQLRowCount(self.hstmt) + self._clear_rownumber() + self._cached_column_map = None + self._cached_column_map_lower = None + self._cached_converter_map = None + self._uuid_str_indices = None + + self._reset_inputsizes() + + def _wrap_row_async(self, row_data: List[Any]) -> Row: + """Async-only mirror of fetchone()'s Row-construction block.""" + column_map, converter_map, column_map_lower = self._get_column_and_converter_maps() + return Row( + row_data, + column_map, + cursor=self, + converter_map=converter_map, + uuid_str_indices=self._uuid_str_indices, + column_map_lower=column_map_lower, + ) + + def _wrap_rows_async(self, rows_data: List[List[Any]]) -> List[Row]: + """Async-only mirror of fetchmany() / fetchall()'s Row-construction block.""" + column_map, converter_map, column_map_lower = self._get_column_and_converter_maps() + uuid_idx = self._uuid_str_indices + return [ + Row( + row_data, + column_map, + cursor=self, + converter_map=converter_map, + uuid_str_indices=uuid_idx, + column_map_lower=column_map_lower, + ) + for row_data in rows_data + ] + + # ---- Public async API --------------------------------------------------- + + async def execute_async( + self, + operation: str, + *parameters, + use_prepare: bool = True, + reset_cursor: bool = True, + poll_initial_ms: float = 0.5, + poll_max_ms: float = 20.0, + ) -> "Cursor": + """Async counterpart of :meth:`execute` (POC). + + Runs the ODBC execute call on a background executor thread with + SQL_ATTR_ASYNC_ENABLE turned on, so the caller's asyncio event loop + stays responsive while the driver polls SQLExecute / SQLExecDirect. + + Args: + operation: SQL query or command. + parameters: Sequence of parameters to bind (same semantics as + :meth:`execute`). Data-At-Execution parameters (very large + strings / bytes) are NOT supported on the async path — use + :meth:`execute` for those. + use_prepare: Whether to use SQLPrepareW (default) or SQLExecDirectW. + reset_cursor: Whether to reset the cursor before execution. + poll_initial_ms: First sleep interval between SQL_STILL_EXECUTING + polls in the C++ polling loop. + poll_max_ms: Capped exponential backoff ceiling (1.5x per iteration). + + Returns: + self, for method chaining. + + Raises: + NotSupportedError: driver does not advertise async support. + ProgrammingError: another async operation is already in flight on + this cursor. + """ + import asyncio # lazy — no cost for sync-only users + + self._check_closed() + self._check_async_capable() + + logger.debug( + "execute_async: Starting - operation_length=%d, param_count=%d, use_prepare=%s", + len(operation), + len(parameters), + str(use_prepare), + ) + logger.debug("Executing query (async): %s", operation) + + ( + operation, + parameters, + parameters_type, + effective_use_prepare, + encoding_settings, + ) = self._prepare_execute_state_async( + operation, parameters, use_prepare, reset_cursor + ) + + loop = asyncio.get_running_loop() + self._acquire_async_slot() + try: + ret = await loop.run_in_executor( + None, + ddbc_bindings.DDBCSQLExecuteAsync, + self.hstmt, + operation, + parameters, + parameters_type, + self.is_stmt_prepared, + effective_use_prepare, + encoding_settings, + poll_initial_ms, + poll_max_ms, + ) + finally: + self._release_async_slot() + + self._finalize_execute_async(ret, operation) + return self + + async def fetch_async( + self, + size: Optional[int] = None, + *, + poll_initial_ms: float = 0.5, + poll_max_ms: float = 20.0, + ) -> Union[Row, List[Row], None]: + """Async counterpart of :meth:`fetchone` / :meth:`fetchmany` / :meth:`fetchall` (POC). + + Dispatch table: + + ================= ======================================================= + ``size`` value Behavior + ================= ======================================================= + ``None`` (default) Single ``Row`` or ``None`` (fetchone semantics). + positive int ``List[Row]`` of up to ``size`` rows (fetchmany). + ``-1`` ``List[Row]`` of all remaining rows (fetchall). + ``0`` / other <= 0 Empty list (fetchmany semantics for size <= 0). + ================= ======================================================= + + Args: + size: See dispatch table above. + poll_initial_ms: First sleep interval between SQL_STILL_EXECUTING + polls in the C++ polling loop. + poll_max_ms: Capped exponential backoff ceiling. + + Returns: + A ``Row``, a ``list[Row]``, or ``None`` — see dispatch table. + + Raises: + NotSupportedError: driver does not advertise async support. + ProgrammingError: another async operation is already in flight on + this cursor. + """ + import asyncio # lazy + + self._check_closed() + self._check_async_capable() + + if not self._has_result_set and self.description: + self._reset_rownumber() + + char_decoding = self._get_decoding_settings(ddbc_sql_const.SQL_CHAR.value) + wchar_decoding = self._get_decoding_settings(ddbc_sql_const.SQL_WCHAR.value) + char_enc = char_decoding.get("encoding", "utf-16le") + wchar_enc = wchar_decoding.get("encoding", "utf-16le") + char_ctype = char_decoding.get("ctype", ddbc_sql_const.SQL_WCHAR.value) + + loop = asyncio.get_running_loop() + + if size is None: + # ---- fetchone-async ---- + row_data: List[Any] = [] + self._acquire_async_slot() + try: + ret = await loop.run_in_executor( + None, + ddbc_bindings.DDBCSQLFetchOneAsync, + self.hstmt, + row_data, + char_enc, + wchar_enc, + char_ctype, + poll_initial_ms, + poll_max_ms, + ) + finally: + self._release_async_slot() + + if self.hstmt: + self.messages.extend( + ddbc_bindings.DDBCSQLGetAllDiagRecords(self.hstmt) + ) + + if ret == ddbc_sql_const.SQL_NO_DATA.value: + if self._next_row_index == 0 and self.description is not None: + self.rowcount = 0 + return None + + # rownumber tracking (mirrors sync fetchone) + if self._skip_increment_for_next_fetch: + self._skip_increment_for_next_fetch = False + self._next_row_index += 1 + else: + self._increment_rownumber() + self.rowcount = self._next_row_index + return self._wrap_row_async(row_data) + + if size == -1: + # ---- fetchall-async ---- + rows_data: List[List[Any]] = [] + self._acquire_async_slot() + try: + ret = await loop.run_in_executor( + None, + ddbc_bindings.DDBCSQLFetchAllAsync, + self.hstmt, + rows_data, + char_enc, + wchar_enc, + char_ctype, + poll_initial_ms, + poll_max_ms, + ) + finally: + self._release_async_slot() + + check_error(ddbc_sql_const.SQL_HANDLE_STMT.value, self.hstmt, ret) + + if self.hstmt: + self.messages.extend( + ddbc_bindings.DDBCSQLGetAllDiagRecords(self.hstmt) + ) + + if rows_data and self._has_result_set: + self._next_row_index += len(rows_data) + self._rownumber = self._next_row_index - 1 + + if len(rows_data) == 0 and self._next_row_index == 0: + self.rowcount = 0 + else: + self.rowcount = self._next_row_index + + return self._wrap_rows_async(rows_data) + + # ---- fetchmany-async (size > 0 or size <= 0 sentinel) ---- + if size <= 0: + return [] + rows_data = [] + self._acquire_async_slot() + try: + ret = await loop.run_in_executor( + None, + ddbc_bindings.DDBCSQLFetchManyAsync, + self.hstmt, + rows_data, + size, + char_enc, + wchar_enc, + char_ctype, + poll_initial_ms, + poll_max_ms, + ) + finally: + self._release_async_slot() + + if self.hstmt: + self.messages.extend( + ddbc_bindings.DDBCSQLGetAllDiagRecords(self.hstmt) + ) + + if rows_data and self._has_result_set: + self._next_row_index += len(rows_data) + self._rownumber = self._next_row_index - 1 + + if len(rows_data) == 0 and self._next_row_index == 0: + self.rowcount = 0 + else: + self.rowcount = self._next_row_index + + return self._wrap_rows_async(rows_data) diff --git a/mssql_python/pybind/connection/connection.cpp b/mssql_python/pybind/connection/connection.cpp index 4b366575..de75f082 100644 --- a/mssql_python/pybind/connection/connection.cpp +++ b/mssql_python/pybind/connection/connection.cpp @@ -641,6 +641,45 @@ py::object ConnectionHandle::getInfo(SQLUSMALLINT infoType) const { return _conn->getInfo(infoType); } +// Async POC: probe SQL_ASYNC_MODE and cache the result. Returns true iff the +// driver supports statement-level async (SQL_AM_STATEMENT) or connection-level +// async (SQL_AM_CONNECTION). Called from Cursor.execute_async as a capability +// gate before configuring SQL_ATTR_ASYNC_ENABLE on the statement handle. +bool Connection::isAsyncCapable() const { + int cached = _asyncModeCache.load(std::memory_order_acquire); + if (cached < 0) { + if (!_dbcHandle) { + return false; + } + if (!SQLGetInfo_ptr) { + LOG("isAsyncCapable: SQLGetInfo not initialized, loading driver"); + DriverLoader::getInstance().loadDriver(); + } + SQLUSMALLINT asyncMode = 0; + SQLSMALLINT outLen = 0; + SQLRETURN ret = SQLGetInfo_ptr(_dbcHandle->get(), SQL_ASYNC_MODE, &asyncMode, + sizeof(asyncMode), &outLen); + if (!SQL_SUCCEEDED(ret)) { + LOG("isAsyncCapable: SQLGetInfo(SQL_ASYNC_MODE) failed - SQLRETURN=%d", ret); + // Cache as SQL_AM_NONE so we don't re-probe on every call. + _asyncModeCache.store(static_cast(SQL_AM_NONE), std::memory_order_release); + return false; + } + cached = static_cast(asyncMode); + _asyncModeCache.store(cached, std::memory_order_release); + LOG("isAsyncCapable: SQL_ASYNC_MODE=%d (cached)", cached); + } + return cached == static_cast(SQL_AM_STATEMENT) || + cached == static_cast(SQL_AM_CONNECTION); +} + +bool ConnectionHandle::isAsyncCapable() const { + if (!_conn) { + ThrowStdException("Connection object is not initialized"); + } + return _conn->isAsyncCapable(); +} + void ConnectionHandle::setAttr(int attribute, py::object value) { if (!_conn) { ThrowStdException("Connection not established"); diff --git a/mssql_python/pybind/connection/connection.h b/mssql_python/pybind/connection/connection.h index 981bbc0a..0f16e950 100644 --- a/mssql_python/pybind/connection/connection.h +++ b/mssql_python/pybind/connection/connection.h @@ -3,6 +3,7 @@ #pragma once #include "../ddbc_bindings.h" +#include #include #include #include @@ -54,6 +55,12 @@ class Connection { // Get information about the driver and data source py::object getInfo(SQLUSMALLINT infoType) const; + // Async POC: returns true iff the driver advertises statement-level (or + // higher) async support via SQLGetInfo(SQL_ASYNC_MODE). Result is cached + // per-connection after the first successful call because SQL_ASYNC_MODE + // is a static driver capability. + bool isAsyncCapable() const; + SQLRETURN setAttribute(SQLINTEGER attribute, py::object value); // Add getter for DBC handle for error reporting @@ -92,6 +99,15 @@ class Connection { // Prevents data races between allocStatementHandle() and disconnect(), // or concurrent GC finalizers running from different threads mutable std::mutex _childHandlesMutex; + + // Async POC: cached result of SQLGetInfo(SQL_ASYNC_MODE). + // -1 = uncached + // 0 = SQL_AM_NONE (no async) + // 1 = SQL_AM_CONNECTION + // 2 = SQL_AM_STATEMENT + // Written at most once per connection; atomic to avoid a mutex on the + // read path (called once per execute_async / fetch_async invocation). + mutable std::atomic _asyncModeCache{-1}; }; class ConnectionHandle { @@ -111,6 +127,9 @@ class ConnectionHandle { // Get information about the driver and data source py::object getInfo(SQLUSMALLINT infoType) const; + // Async POC: forwards to Connection::isAsyncCapable(). + bool isAsyncCapable() const; + private: std::shared_ptr _conn; bool _usePool; diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index 3cb00814..404afbf0 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -13,11 +13,13 @@ #include // std::min #include +#include // std::chrono (async polling backoff) #include #include // For std::memcpy #include #include // std::setw, std::setfill #include +#include // std::this_thread::sleep_for (async polling backoff) #include // std::forward @@ -1806,18 +1808,161 @@ SQLRETURN SQLTables_wrap(SqlHandlePtr StatementHandle, const std::u16string& cat return ret; } +// --------------------------------------------------------------------------- +// Async POC helpers (statement-level async via SQL_ATTR_ASYNC_ENABLE + polling) +// --------------------------------------------------------------------------- +// These helpers underpin the upcoming Cursor.execute_async / Cursor.fetch_async +// bindings. They are intentionally usable from the existing sync path with +// asyncMode=false (a no-op path), so we share one code path for both. + +// Backoff parameters for the polling loop. Defaults tuned for typical LAN +// SQL Server round-trips; caller can override per invocation. +struct PollingConfig { + double initial_ms = 0.5; // first sleep interval + double max_ms = 20.0; // capped exponential backoff ceiling + double multiplier = 1.5; // geometric growth factor +}; + +// RAII guard: enables SQL_ATTR_ASYNC_ENABLE on construction, disables on +// destruction. Only enables if `enable` is true AND the driver accepts the +// attribute; otherwise the destructor is a no-op. Ensures the async flag +// never leaks onto a statement handle that a later sync call reuses. +class AsyncEnableGuard { + public: + AsyncEnableGuard(SQLHSTMT hStmt, bool enable) : _hStmt(hStmt), _enabled(false) { + if (enable && _hStmt && SQLSetStmtAttr_ptr) { + SQLRETURN rc = SQLSetStmtAttr_ptr( + _hStmt, SQL_ATTR_ASYNC_ENABLE, + reinterpret_cast(static_cast(SQL_ASYNC_ENABLE_ON)), + 0); + if (SQL_SUCCEEDED(rc)) { + _enabled = true; + } else { + LOG("AsyncEnableGuard: SQLSetStmtAttr(SQL_ATTR_ASYNC_ENABLE, ON) " + "failed - SQLRETURN=%d, hStmt=%p", + rc, (void*)_hStmt); + } + } + } + + ~AsyncEnableGuard() { + if (_enabled && _hStmt && SQLSetStmtAttr_ptr) { + SQLRETURN rc = SQLSetStmtAttr_ptr( + _hStmt, SQL_ATTR_ASYNC_ENABLE, + reinterpret_cast(static_cast(SQL_ASYNC_ENABLE_OFF)), + 0); + if (!SQL_SUCCEEDED(rc)) { + LOG("AsyncEnableGuard: SQLSetStmtAttr(SQL_ATTR_ASYNC_ENABLE, OFF) " + "failed - SQLRETURN=%d, hStmt=%p", + rc, (void*)_hStmt); + } + } + } + + bool enabled() const { return _enabled; } + + AsyncEnableGuard(const AsyncEnableGuard&) = delete; + AsyncEnableGuard& operator=(const AsyncEnableGuard&) = delete; + AsyncEnableGuard(AsyncEnableGuard&&) = delete; + AsyncEnableGuard& operator=(AsyncEnableGuard&&) = delete; + + private: + SQLHSTMT _hStmt; + bool _enabled; +}; + +// Executes the provided ODBC call once (sync) or in a polling loop that keeps +// re-invoking it while it returns SQL_STILL_EXECUTING (async). +// +// IMPORTANT: The caller MUST release the GIL before invoking this helper when +// asyncMode == true, so the sleep_for periods don't block the Python event +// loop running in the calling thread's asyncio task. +// +// `callOnce` is any callable returning SQLRETURN (typically a lambda that +// captures the statement handle and calls SQLExecute / SQLExecDirect / SQLFetch). +template +inline SQLRETURN ExecuteWithPolling(Fn callOnce, bool asyncMode, + const PollingConfig& cfg = PollingConfig{}) { + SQLRETURN ret = callOnce(); + if (!asyncMode) { + return ret; + } + double sleep_ms = cfg.initial_ms; + while (ret == SQL_STILL_EXECUTING) { + std::this_thread::sleep_for( + std::chrono::microseconds(static_cast(sleep_ms * 1000.0))); + sleep_ms = std::min(sleep_ms * cfg.multiplier, cfg.max_ms); + ret = callOnce(); + } + return ret; +} + +// Prepares the statement (if requested) and binds all parameters. Returns +// the SQLRETURN from the last ODBC call; on non-success the caller should +// short-circuit before invoking SQLExecute. +// +// `paramBuffers` is an OUT parameter: it accumulates heap-owned buffers that +// MUST remain in scope until SQLExecute completes, because ODBC keeps raw +// pointers into them across the SQLBindParameter / SQLExecute boundary. +// +// GIL: SQLPrepare is called with the GIL released (matches previous inline +// behavior); BindParameters requires the GIL (inspects py::list contents). +SQLRETURN PrepareAndBind(SqlHandlePtr statementHandle, SQLHANDLE hStmt, SQLWCHAR* queryPtr, + const py::list& params, std::vector& paramInfos, + py::list& isStmtPrepared, bool usePrepare, + std::vector>& paramBuffers, + const std::string& charEncoding) { + // isStmtPrepared is a single-element list carrying a bool by reference + // (Python bools are immutable, so we can't pass the raw bool by ref). + assert(isStmtPrepared.size() == 1); + + SQLRETURN rc = SQL_SUCCESS; + if (usePrepare) { + { + // Release the GIL during the blocking SQLPrepare network call. + py::gil_scoped_release release; + rc = SQLPrepare_ptr(hStmt, queryPtr, SQL_NTS); + } + if (!SQL_SUCCEEDED(rc)) { + LOG("PrepareAndBind: SQLPrepare failed - SQLRETURN=%d, statement_handle=%p", + rc, (void*)hStmt); + return rc; + } + // GH-610: Clear per-handle describe cache (new prepare = new param types) + statementHandle->clearDescribeCache(); + isStmtPrepared[0] = py::cast(true); + } else { + // Caller opted out of preparing; the plan must already exist on hStmt. + bool isStmtPreparedAsBool = isStmtPrepared[0].cast(); + if (!isStmtPreparedAsBool) { + ThrowStdException("Cannot execute unprepared statement"); + } + } + + rc = BindParameters(*statementHandle, hStmt, params, paramInfos, paramBuffers, charEncoding); + return rc; +} + // Executes the provided query. If the query is parametrized, it prepares the // statement and binds the parameters. Otherwise, it executes the query // directly. 'usePrepare' parameter can be used to disable the prepare step for // queries that might already be prepared in a previous call. -SQLRETURN SQLExecute_wrap(const SqlHandlePtr statementHandle, const std::u16string& query, - const py::list& params, std::vector& paramInfos, - py::list& isStmtPrepared, const bool usePrepare, - const py::dict& encodingSettings) { +// +// Internal implementation shared by sync (SQLExecute_wrap) and async +// (SQLExecuteAsync_wrap) entry points. When asyncMode=true, SQL_ATTR_ASYNC_ENABLE +// is turned on via AsyncEnableGuard for the duration of this call, and the +// SQLExecute / SQLExecDirect call is driven through ExecuteWithPolling so the +// caller thread stays responsive between polls. Data-At-Execution parameters +// are rejected up-front in async mode (async DAE is out of scope for the POC). +static SQLRETURN SQLExecute_impl(const SqlHandlePtr statementHandle, + const std::u16string& query, const py::list& params, + std::vector& paramInfos, py::list& isStmtPrepared, + const bool usePrepare, const py::dict& encodingSettings, + bool asyncMode, const PollingConfig& pollCfg) { LOG("SQLExecute: Executing %s query - statement_handle=%p, " - "param_count=%zu, query_length=%zu chars", + "param_count=%zu, query_length=%zu chars, async=%d", (params.size() > 0 ? "parameterized" : "direct"), (void*)statementHandle->get(), - params.size(), query.length()); + params.size(), query.length(), asyncMode ? 1 : 0); if (!SQLPrepare_ptr) { LOG("SQLExecute: Function pointer not initialized, loading driver"); DriverLoader::getInstance().loadDriver(); // Load the driver @@ -1830,6 +1975,21 @@ SQLRETURN SQLExecute_wrap(const SqlHandlePtr statementHandle, const std::u16stri ThrowStdException("Number of parameters and paramInfos do not match"); } + // Async POC: reject Data-At-Execution parameters up-front. DAE requires + // the SQL_NEED_DATA loop below, which is not safe to drive alongside + // SQL_STILL_EXECUTING polling. Callers must fall back to sync execute for + // large VARBINARY(MAX) / NVARCHAR(MAX) parameter values. + if (asyncMode) { + for (const auto& info : paramInfos) { + if (info.isDAE) { + ThrowStdException( + "Data-At-Execution (DAE) parameters are not supported with " + "async execution in this POC. Use sync execute for large " + "VARBINARY(MAX) / NVARCHAR(MAX) parameters."); + } + } + } + RETCODE rc; SQLHANDLE hStmt = statementHandle->get(); if (!statementHandle || !statementHandle->get()) { @@ -1842,6 +2002,11 @@ SQLRETURN SQLExecute_wrap(const SqlHandlePtr statementHandle, const std::u16stri SQLSetStmtAttr_ptr(hStmt, SQL_ATTR_CONCURRENCY, (SQLPOINTER)SQL_CONCUR_READ_ONLY, 0); } + // Async POC: enable SQL_ATTR_ASYNC_ENABLE for the duration of this call. + // When asyncMode=false the guard is a no-op (constructor + destructor skip + // the SQLSetStmtAttr calls entirely). + AsyncEnableGuard asyncGuard(hStmt, asyncMode); + SQLWCHAR* queryPtr = reinterpretU16stringAsSqlWChar(query); if (params.size() == 0) { // Execute statement directly if the statement is not parametrized. This @@ -1849,9 +2014,13 @@ SQLRETURN SQLExecute_wrap(const SqlHandlePtr statementHandle, const std::u16stri // according to DDBC documentation - // https://learn.microsoft.com/en-us/sql/odbc/reference/syntax/sqlexecdirect-function?view=sql-server-ver16 { - // Release the GIL during the blocking ODBC call + // Release the GIL during the blocking ODBC call. In async mode + // ExecuteWithPolling re-invokes SQLExecDirect while it returns + // SQL_STILL_EXECUTING, with backoff sleeps between polls. py::gil_scoped_release release; - rc = SQLExecDirect_ptr(hStmt, queryPtr, SQL_NTS); + rc = ExecuteWithPolling( + [&]() { return SQLExecDirect_ptr(hStmt, queryPtr, SQL_NTS); }, + asyncMode, pollCfg); } if (!SQL_SUCCEEDED(rc) && rc != SQL_NO_DATA) { LOG("SQLExecute: Direct execution failed (non-parameterized query) " @@ -1860,54 +2029,28 @@ SQLRETURN SQLExecute_wrap(const SqlHandlePtr statementHandle, const std::u16stri } return rc; } else { - // isStmtPrepared is a list instead of a bool coz bools in Python are - // immutable. Hence, we can't pass around bools by reference & modify - // them. Therefore, isStmtPrepared must be a list with exactly one bool - // element - assert(isStmtPrepared.size() == 1); - if (usePrepare) { - { - // Release the GIL during the blocking SQLPrepare network call. - py::gil_scoped_release release; - rc = SQLPrepare_ptr(hStmt, queryPtr, SQL_NTS); - } - if (!SQL_SUCCEEDED(rc)) { - LOG("SQLExecute: SQLPrepare failed - SQLRETURN=%d, " - "statement_handle=%p", - rc, (void*)hStmt); - return rc; - } - // GH-610: Clear per-handle describe cache (new prepare = new param types) - statementHandle->clearDescribeCache(); - isStmtPrepared[0] = py::cast(true); - } else { - // Make sure the statement has been prepared earlier if we're not - // preparing now - bool isStmtPreparedAsBool = isStmtPrepared[0].cast(); - if (!isStmtPreparedAsBool) { - // TODO: Print the query - ThrowStdException("Cannot execute unprepared statement"); - } - } - - // This vector manages the heap memory allocated for parameter buffers. - // It must be in scope until SQLExecute is done. // Extract char encoding from encodingSettings dictionary std::string charEncoding = "utf-8"; // default if (encodingSettings.contains("encoding")) { charEncoding = encodingSettings["encoding"].cast(); } + // This vector manages the heap memory allocated for parameter buffers. + // It must remain in scope until SQLExecute (and any DAE loop) completes. std::vector> paramBuffers; - rc = BindParameters(*statementHandle, hStmt, params, paramInfos, paramBuffers, charEncoding); + rc = PrepareAndBind(statementHandle, hStmt, queryPtr, params, paramInfos, + isStmtPrepared, usePrepare, paramBuffers, charEncoding); if (!SQL_SUCCEEDED(rc)) { return rc; } { - // Release the GIL during the blocking SQLExecute network call. + // Release the GIL during the blocking SQLExecute network call. In + // async mode ExecuteWithPolling handles SQL_STILL_EXECUTING with + // backoff sleeps between polls. py::gil_scoped_release release; - rc = SQLExecute_ptr(hStmt); + rc = ExecuteWithPolling([&]() { return SQLExecute_ptr(hStmt); }, + asyncMode, pollCfg); } if (rc == SQL_NEED_DATA) { LOG("SQLExecute: SQL_NEED_DATA received - Starting DAE " @@ -2047,6 +2190,33 @@ SQLRETURN SQLExecute_wrap(const SqlHandlePtr statementHandle, const std::u16stri } } +// Sync wrapper (exposed as DDBCSQLExecute). Existing pybind binding target; +// signature unchanged for source and ABI compatibility with all sync callers. +SQLRETURN SQLExecute_wrap(const SqlHandlePtr statementHandle, const std::u16string& query, + const py::list& params, std::vector& paramInfos, + py::list& isStmtPrepared, const bool usePrepare, + const py::dict& encodingSettings) { + return SQLExecute_impl(statementHandle, query, params, paramInfos, isStmtPrepared, + usePrepare, encodingSettings, + /*asyncMode=*/false, PollingConfig{}); +} + +// Async wrapper (exposed as DDBCSQLExecuteAsync). Called from Cursor.execute_async +// via loop.run_in_executor so the polling loop runs on a background thread with +// the GIL released, keeping the asyncio event loop responsive. +SQLRETURN SQLExecuteAsync_wrap(const SqlHandlePtr statementHandle, const std::u16string& query, + const py::list& params, std::vector& paramInfos, + py::list& isStmtPrepared, const bool usePrepare, + const py::dict& encodingSettings, double poll_initial_ms, + double poll_max_ms) { + PollingConfig cfg; + cfg.initial_ms = poll_initial_ms; + cfg.max_ms = poll_max_ms; + // multiplier stays at the PollingConfig default (1.5) + return SQLExecute_impl(statementHandle, query, params, paramInfos, isStmtPrepared, + usePrepare, encodingSettings, /*asyncMode=*/true, cfg); +} + SQLRETURN BindParameterArray(SqlHandle& handle, SQLHANDLE hStmt, const py::list& columnwise_params, std::vector& paramInfos, size_t paramSetSize, std::vector>& paramBuffers, @@ -4006,17 +4176,27 @@ SQLRETURN SQLBindColums(SQLHSTMT hStmt, ColumnBuffers& buffers, py::list& column // Fetch rows in batches // TODO: Move to anonymous namespace, since it is not used outside this file +// +// Async POC: when asyncMode=true, the SQLFetchScroll network call is driven +// through ExecuteWithPolling so it can return SQL_STILL_EXECUTING. The caller +// is responsible for having installed SQL_ATTR_ASYNC_ENABLE on hStmt (via +// AsyncEnableGuard) before calling this function. SQLRETURN FetchBatchData(SQLHSTMT hStmt, ColumnBuffers& buffers, py::list& columnNames, py::list& rows, SQLUSMALLINT numCols, SQLULEN& numRowsFetched, const std::vector& lobColumns, const std::string& charEncoding = "utf-16le", - int charCtype = SQL_C_WCHAR) { + int charCtype = SQL_C_WCHAR, + bool asyncMode = false, + const PollingConfig& pollCfg = PollingConfig{}) { LOG("FetchBatchData: Fetching data in batches"); SQLRETURN ret; { - // Release the GIL during the blocking ODBC fetch + // Release the GIL during the blocking ODBC fetch. In async mode + // ExecuteWithPolling handles SQL_STILL_EXECUTING with backoff. py::gil_scoped_release release; - ret = SQLFetchScroll_ptr(hStmt, SQL_FETCH_NEXT, 0); + ret = ExecuteWithPolling( + [&]() { return SQLFetchScroll_ptr(hStmt, SQL_FETCH_NEXT, 0); }, + asyncMode, pollCfg); } if (ret == SQL_NO_DATA) { LOG("FetchBatchData: No data to fetch"); @@ -4470,15 +4650,24 @@ size_t calculateRowSize(py::list& columnNames, SQLUSMALLINT numCols) { // 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, - const std::string& charEncoding = "utf-16le", - const std::string& wcharEncoding = "utf-16le", - int charCtype = SQL_C_WCHAR) { +// +// Internal implementation shared by sync (FetchMany_wrap) and async +// (FetchManyAsync_wrap) entry points. In async mode SQL_ATTR_ASYNC_ENABLE is +// installed for the duration of this call and each SQLFetch is polled. +static SQLRETURN FetchMany_impl(SqlHandlePtr StatementHandle, py::list& rows, int fetchSize, + const std::string& charEncoding, + const std::string& wcharEncoding, int charCtype, + bool asyncMode, const PollingConfig& pollCfg) { // 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); SQLRETURN ret; SQLHSTMT hStmt = StatementHandle->get(); + + // Async POC: enable SQL_ATTR_ASYNC_ENABLE for the duration of this call. + // No-op when asyncMode=false. + AsyncEnableGuard asyncGuard(hStmt, asyncMode); + // Retrieve column count SQLSMALLINT numCols = SQLNumResultCols_wrap(StatementHandle); @@ -4510,9 +4699,11 @@ SQLRETURN FetchMany_wrap(SqlHandlePtr StatementHandle, py::list& rows, int fetch lobColumns.size()); while (numRowsFetched < (SQLULEN)fetchSize) { { - // Release GIL during the blocking fetch + // Release GIL during the blocking fetch. In async mode + // ExecuteWithPolling handles SQL_STILL_EXECUTING with backoff. py::gil_scoped_release release; - ret = SQLFetch_ptr(hStmt); + ret = ExecuteWithPolling([&]() { return SQLFetch_ptr(hStmt); }, + asyncMode, pollCfg); } if (ret == SQL_NO_DATA) break; @@ -4542,7 +4733,7 @@ SQLRETURN FetchMany_wrap(SqlHandlePtr StatementHandle, py::list& rows, int fetch SQLSetStmtAttr_ptr(hStmt, SQL_ATTR_ROWS_FETCHED_PTR, &numRowsFetched, 0); ret = FetchBatchData(hStmt, buffers, columnNames, rows, numCols, numRowsFetched, lobColumns, - charEncoding, charCtype); + charEncoding, charCtype, asyncMode, pollCfg); if (!SQL_SUCCEEDED(ret) && ret != SQL_NO_DATA) { LOG("FetchMany_wrap: Error when fetching data - SQLRETURN=%d", ret); return ret; @@ -4558,6 +4749,26 @@ SQLRETURN FetchMany_wrap(SqlHandlePtr StatementHandle, py::list& rows, int fetch return ret; } +// Sync wrapper (exposed as DDBCSQLFetchMany). Existing pybind binding target. +SQLRETURN FetchMany_wrap(SqlHandlePtr StatementHandle, py::list& rows, int fetchSize, + const std::string& charEncoding = "utf-16le", + const std::string& wcharEncoding = "utf-16le", + int charCtype = SQL_C_WCHAR) { + return FetchMany_impl(StatementHandle, rows, fetchSize, charEncoding, wcharEncoding, + charCtype, /*asyncMode=*/false, PollingConfig{}); +} + +// Async wrapper (exposed as DDBCSQLFetchManyAsync). +SQLRETURN FetchManyAsync_wrap(SqlHandlePtr StatementHandle, py::list& rows, int fetchSize, + const std::string& charEncoding, const std::string& wcharEncoding, + int charCtype, double poll_initial_ms, double poll_max_ms) { + PollingConfig cfg; + cfg.initial_ms = poll_initial_ms; + cfg.max_ms = poll_max_ms; + return FetchMany_impl(StatementHandle, rows, fetchSize, charEncoding, wcharEncoding, + charCtype, /*asyncMode=*/true, cfg); +} + // GetDataVar - Progressively fetches variable-length column data using SQLGetData. // // Calls SQLGetData repeatedly, reallocating the buffer as needed, until all data is retrieved. @@ -5589,15 +5800,25 @@ SQLRETURN FetchArrowBatch_wrap(SqlHandlePtr StatementHandle, py::list& capsules, // 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 FetchAll_wrap(SqlHandlePtr StatementHandle, py::list& rows, - const std::string& charEncoding = "utf-16le", - const std::string& wcharEncoding = "utf-16le", - int charCtype = SQL_C_WCHAR) { +// +// Internal implementation shared by sync (FetchAll_wrap) and async +// (FetchAllAsync_wrap) entry points. In async mode SQL_ATTR_ASYNC_ENABLE is +// installed for the duration of this call and each SQLFetch / SQLFetchScroll +// call is polled via ExecuteWithPolling. +static SQLRETURN FetchAll_impl(SqlHandlePtr StatementHandle, py::list& rows, + const std::string& charEncoding, + const std::string& wcharEncoding, int charCtype, + bool asyncMode, const PollingConfig& pollCfg) { // 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); SQLRETURN ret; SQLHSTMT hStmt = StatementHandle->get(); + + // Async POC: enable SQL_ATTR_ASYNC_ENABLE for the duration of this call. + // No-op when asyncMode=false. + AsyncEnableGuard asyncGuard(hStmt, asyncMode); + // Retrieve column count SQLSMALLINT numCols = SQLNumResultCols_wrap(StatementHandle); @@ -5629,9 +5850,11 @@ SQLRETURN FetchAll_wrap(SqlHandlePtr StatementHandle, py::list& rows, lobColumns.size()); while (true) { { - // Release GIL during the blocking fetch + // Release GIL during the blocking fetch. In async mode + // ExecuteWithPolling handles SQL_STILL_EXECUTING with backoff. py::gil_scoped_release release; - ret = SQLFetch_ptr(hStmt); + ret = ExecuteWithPolling([&]() { return SQLFetch_ptr(hStmt); }, + asyncMode, pollCfg); } if (ret == SQL_NO_DATA) break; @@ -5702,7 +5925,7 @@ SQLRETURN FetchAll_wrap(SqlHandlePtr StatementHandle, py::list& rows, while (ret != SQL_NO_DATA) { ret = FetchBatchData(hStmt, buffers, columnNames, rows, numCols, numRowsFetched, lobColumns, - charEncoding, charCtype); + charEncoding, charCtype, asyncMode, pollCfg); if (!SQL_SUCCEEDED(ret) && ret != SQL_NO_DATA) { LOG("FetchAll_wrap: Error when fetching data - SQLRETURN=%d", ret); return ret; @@ -5719,6 +5942,26 @@ SQLRETURN FetchAll_wrap(SqlHandlePtr StatementHandle, py::list& rows, return ret; } +// Sync wrapper (exposed as DDBCSQLFetchAll). Existing pybind binding target. +SQLRETURN FetchAll_wrap(SqlHandlePtr StatementHandle, py::list& rows, + const std::string& charEncoding = "utf-16le", + const std::string& wcharEncoding = "utf-16le", + int charCtype = SQL_C_WCHAR) { + return FetchAll_impl(StatementHandle, rows, charEncoding, wcharEncoding, charCtype, + /*asyncMode=*/false, PollingConfig{}); +} + +// Async wrapper (exposed as DDBCSQLFetchAllAsync). +SQLRETURN FetchAllAsync_wrap(SqlHandlePtr StatementHandle, py::list& rows, + const std::string& charEncoding, const std::string& wcharEncoding, + int charCtype, double poll_initial_ms, double poll_max_ms) { + PollingConfig cfg; + cfg.initial_ms = poll_initial_ms; + cfg.max_ms = poll_max_ms; + return FetchAll_impl(StatementHandle, rows, charEncoding, wcharEncoding, charCtype, + /*asyncMode=*/true, cfg); +} + // FetchOne_wrap - Fetches a single row of data from the result set. // // @param StatementHandle: Handle to the statement from which data is to be @@ -5735,16 +5978,26 @@ SQLRETURN FetchAll_wrap(SqlHandlePtr StatementHandle, py::list& rows, // 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 FetchOne_wrap(SqlHandlePtr StatementHandle, py::list& row, - const std::string& charEncoding = "utf-16le", - const std::string& wcharEncoding = "utf-16le", - int charCtype = SQL_C_WCHAR) { +// +// Internal implementation shared by sync (FetchOne_wrap) and async +// (FetchOneAsync_wrap) entry points. In async mode SQL_ATTR_ASYNC_ENABLE is +// installed via AsyncEnableGuard and the SQLFetch call is polled via +// ExecuteWithPolling. Note: for LOB streams SQLGetData_wrap runs inline and +// does NOT poll — non-LOB fetches are the primary async path today. +static SQLRETURN FetchOne_impl(SqlHandlePtr StatementHandle, py::list& row, + const std::string& charEncoding, + const std::string& wcharEncoding, int charCtype, + bool asyncMode, const PollingConfig& pollCfg) { // 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); SQLRETURN ret; SQLHSTMT hStmt = StatementHandle->get(); + // Async POC: enable SQL_ATTR_ASYNC_ENABLE for the duration of this call. + // No-op when asyncMode=false. + AsyncEnableGuard asyncGuard(hStmt, asyncMode); + // Unbind any columns from previous fetch operations (e.g., fetchmany) // to avoid conflicts with SQLGetData. SQLGetData cannot be used on // columns that are already bound. @@ -5752,9 +6005,11 @@ SQLRETURN FetchOne_wrap(SqlHandlePtr StatementHandle, py::list& row, // Assume hStmt is already allocated and a query has been executed { - // Release the GIL during the blocking ODBC fetch + // Release the GIL during the blocking ODBC fetch. In async mode + // ExecuteWithPolling re-invokes SQLFetch while it returns + // SQL_STILL_EXECUTING, with backoff sleeps between polls. py::gil_scoped_release release; - ret = SQLFetch_ptr(hStmt); + ret = ExecuteWithPolling([&]() { return SQLFetch_ptr(hStmt); }, asyncMode, pollCfg); } if (SQL_SUCCEEDED(ret)) { // Retrieve column count @@ -5771,6 +6026,28 @@ SQLRETURN FetchOne_wrap(SqlHandlePtr StatementHandle, py::list& row, return ret; } +// Sync wrapper (exposed as DDBCSQLFetchOne). Existing pybind binding target; +// keeps the original signature and default arguments unchanged. +SQLRETURN FetchOne_wrap(SqlHandlePtr StatementHandle, py::list& row, + const std::string& charEncoding = "utf-16le", + const std::string& wcharEncoding = "utf-16le", + int charCtype = SQL_C_WCHAR) { + return FetchOne_impl(StatementHandle, row, charEncoding, wcharEncoding, charCtype, + /*asyncMode=*/false, PollingConfig{}); +} + +// Async wrapper (exposed as DDBCSQLFetchOneAsync). Called from Cursor.fetch_async +// via loop.run_in_executor. +SQLRETURN FetchOneAsync_wrap(SqlHandlePtr StatementHandle, py::list& row, + const std::string& charEncoding, const std::string& wcharEncoding, + int charCtype, double poll_initial_ms, double poll_max_ms) { + PollingConfig cfg; + cfg.initial_ms = poll_initial_ms; + cfg.max_ms = poll_max_ms; + return FetchOne_impl(StatementHandle, row, charEncoding, wcharEncoding, charCtype, + /*asyncMode=*/true, cfg); +} + // Wrap SQLMoreResults SQLRETURN SQLMoreResults_wrap(SqlHandlePtr StatementHandle) { LOG("SQLMoreResults_wrap: Check for more results"); @@ -5913,7 +6190,11 @@ PYBIND11_MODULE(ddbc_bindings, m) { .def("set_attr", &ConnectionHandle::setAttr, py::arg("attribute"), py::arg("value"), "Set connection attribute") .def("alloc_statement_handle", &ConnectionHandle::allocStatementHandle) - .def("get_info", &ConnectionHandle::getInfo, py::arg("info_type")); + .def("get_info", &ConnectionHandle::getInfo, py::arg("info_type")) + .def("is_async_capable", &ConnectionHandle::isAsyncCapable, + "Async POC: returns True iff the driver advertises statement-level " + "or connection-level async support via SQLGetInfo(SQL_ASYNC_MODE). " + "Result is cached per-connection."); m.def("enable_pooling", &enable_pooling, "Enable global connection pooling"); m.def("close_pooling", []() { ConnectionPoolManager::getInstance().closePools(); }); m.def("DDBCSQLExecDirect", &SQLExecDirect_wrap, "Execute a SQL query directly"); @@ -5942,6 +6223,46 @@ PYBIND11_MODULE(ddbc_bindings, m) { m.def("DDBCSQLFetchAll", &FetchAll_wrap, "Fetch all rows from the result set", py::arg("StatementHandle"), py::arg("rows"), py::arg("charEncoding") = "utf-16le", py::arg("wcharEncoding") = "utf-16le", py::arg("charCtype") = SQL_C_WCHAR); + + // ------------------------------------------------------------------------ + // Async POC: statement-level async execute + fetch via polling loop. + // Called from Python's Cursor.execute_async / Cursor.fetch_async through + // loop.run_in_executor, so the polling loop runs on a background thread + // with the GIL released and the asyncio event loop stays responsive. + // + // Each *Async binding takes two extra float args: + // poll_initial_ms: first sleep interval between SQL_STILL_EXECUTING polls + // poll_max_ms: capped exponential backoff ceiling + // The multiplier (1.5x) is fixed at the C++ default for this POC. + // + // Note: DDBCSQLExecuteAsync rejects DAE (Data-At-Execution) parameters + // up-front - use sync execute for large VARBINARY(MAX) / NVARCHAR(MAX). + // ------------------------------------------------------------------------ + m.def("DDBCSQLExecuteAsync", &SQLExecuteAsync_wrap, + "Prepare and execute a T-SQL statement with statement-level async polling", + py::arg("statementHandle"), py::arg("query"), py::arg("params"), + py::arg("paramInfos"), py::arg("isStmtPrepared"), py::arg("usePrepare"), + py::arg("encodingSettings"), py::arg("poll_initial_ms") = 0.5, + py::arg("poll_max_ms") = 20.0); + m.def("DDBCSQLFetchOneAsync", &FetchOneAsync_wrap, + "Fetch one row from the result set with statement-level async polling", + py::arg("StatementHandle"), py::arg("row"), + py::arg("charEncoding") = "utf-16le", py::arg("wcharEncoding") = "utf-16le", + py::arg("charCtype") = SQL_C_WCHAR, py::arg("poll_initial_ms") = 0.5, + py::arg("poll_max_ms") = 20.0); + m.def("DDBCSQLFetchManyAsync", &FetchManyAsync_wrap, + "Fetch many rows from the result set with statement-level async polling", + py::arg("StatementHandle"), py::arg("rows"), py::arg("fetchSize"), + py::arg("charEncoding") = "utf-16le", py::arg("wcharEncoding") = "utf-16le", + py::arg("charCtype") = SQL_C_WCHAR, py::arg("poll_initial_ms") = 0.5, + py::arg("poll_max_ms") = 20.0); + m.def("DDBCSQLFetchAllAsync", &FetchAllAsync_wrap, + "Fetch all rows from the result set with statement-level async polling", + py::arg("StatementHandle"), py::arg("rows"), + py::arg("charEncoding") = "utf-16le", py::arg("wcharEncoding") = "utf-16le", + py::arg("charCtype") = SQL_C_WCHAR, py::arg("poll_initial_ms") = 0.5, + py::arg("poll_max_ms") = 20.0); + m.def("DDBCSQLFetchArrowBatch", &FetchArrowBatch_wrap, "Fetch an arrow batch of given length from the result set"); m.def("DDBCSQLFreeHandle", &SQLFreeHandle_wrap, "Free a handle"); diff --git a/tests/test_030_async_execute_fetch.py b/tests/test_030_async_execute_fetch.py new file mode 100644 index 00000000..fa5b7144 --- /dev/null +++ b/tests/test_030_async_execute_fetch.py @@ -0,0 +1,948 @@ +""" +Copyright (c) Microsoft Corporation. +Licensed under the MIT license. + +Basic functional tests for the async POC on ``Cursor`` — ``execute_async`` +and ``fetch_async`` (see ``docs/async_query_poc_spec.md``). + +The suite is intentionally minimal for the POC. It covers: + 1. The Step-1 capability probe (``Connection.is_async_capable``) reports True + on hosts where the Microsoft ODBC driver supports statement-level async. + 2. A single ``execute_async`` / ``fetch_async`` round-trip returns the + expected row (end-to-end smoke). + 3. **The requested "100 concurrent async statements" workload** — fires 100 + ``execute_async`` + ``fetch_async`` pairs through ``asyncio.gather`` with + a bounded semaphore, and asserts all 100 complete with the correct row. + Uses one dedicated connection per task to avoid SQL Server's + Multiple-Active-Result-Sets (MARS) restriction on a single connection. + +The tests skip cleanly when: + * ``DB_CONNECTION_STRING`` is not set in the environment, OR + * the driver reports ``SQL_ASYNC_MODE == SQL_AM_NONE`` (async unavailable). + +No dependency on ``pytest-asyncio`` — each test uses ``asyncio.run(...)`` +directly, so the existing pytest install is sufficient. +""" + +import asyncio +import os +import time + +import pytest + +from mssql_python import connect +from mssql_python.exceptions import ProgrammingError + + +# ============================================================================ +# Fixtures +# ============================================================================ + + +@pytest.fixture(scope="module") +def conn_str(): + """Skip the whole module if ``DB_CONNECTION_STRING`` is unset.""" + cs = os.getenv("DB_CONNECTION_STRING") + if not cs: + pytest.skip("DB_CONNECTION_STRING environment variable not set") + return cs + + +@pytest.fixture(scope="module") +def _async_capable(conn_str): + """Skip all tests in the module if the driver doesn't advertise async support. + + Opens a short-lived connection just to run the SQLGetInfo(SQL_ASYNC_MODE) + probe added in async POC step 1. Result is cached inside the C++ Connection + object, so this doesn't cost anything for later tests that reconnect. + """ + conn = connect(conn_str) + try: + if not conn._conn.is_async_capable(): + pytest.skip( + "ODBC driver reports SQL_AM_NONE — async POC is not usable on this driver" + ) + finally: + conn.close() + + +# ============================================================================ +# Sanity tests +# ============================================================================ + + +def test_capability_probe_returns_true(conn_str, _async_capable): + """``Connection.is_async_capable()`` reports True on this driver. + + Redundant with the ``_async_capable`` fixture, but keeps the probe visible + in the test list so a regression on the C++ SQLGetInfo path shows up as a + named failure rather than a silent module skip. + """ + conn = connect(conn_str) + try: + assert conn._conn.is_async_capable() is True + finally: + conn.close() + + +def test_single_execute_async_and_fetch_async(conn_str, _async_capable): + """One ``execute_async`` followed by ``fetch_async()`` returns the expected row.""" + + async def _run(): + conn = connect(conn_str) + try: + cur = conn.cursor() + try: + await cur.execute_async( + "SELECT ? AS n, CAST(? AS NVARCHAR(16)) AS msg", 42, "hello" + ) + row = await cur.fetch_async() + assert row is not None + assert row[0] == 42 + assert row[1] == "hello" + + # After the single row, fetch_async() should return None. + assert await cur.fetch_async() is None + finally: + cur.close() + finally: + conn.close() + + asyncio.run(_run()) + + +# ============================================================================ +# The 100-concurrent async workload +# ============================================================================ +# Fires ASYNC_TEST_CONCURRENCY (default 100) SELECT statements concurrently +# through ``asyncio.gather``. Each task runs on its own connection so we +# don't hit SQL Server's default no-MARS restriction (a single connection +# only supports one active statement at a time). +# +# A bounded ``asyncio.Semaphore`` limits how many are simultaneously +# in flight. Both knobs are env-overridable so CI environments with +# tight connection limits (or slow bring-up) can tune without editing. +# ============================================================================ + +CONCURRENT_TASKS = int(os.getenv("ASYNC_TEST_CONCURRENCY", "100")) +CONCURRENT_LIMIT = int(os.getenv("ASYNC_TEST_MAX_INFLIGHT", "16")) +# Loose upper bound to catch obvious hangs; not a perf assertion. +WALL_CLOCK_BUDGET_SECONDS = float(os.getenv("ASYNC_TEST_WALL_BUDGET", "120")) + + +def test_100_concurrent_async_selects(conn_str, _async_capable): + """Fire 100 (or ``ASYNC_TEST_CONCURRENCY``) async SELECT statements concurrently. + + Each task: + 1. Opens its own connection (fresh HSTMT, avoids MARS restriction). + 2. Runs ``execute_async`` with two parameters (index + label). + 3. Runs ``fetch_async()`` and asserts the row matches its index. + 4. Closes cursor + connection. + + A ``Semaphore(CONCURRENT_LIMIT)`` throttles the number of simultaneously + in-flight tasks so we don't exhaust the client's default asyncio executor + (``ThreadPoolExecutor`` with ``max_workers = min(32, os.cpu_count()+4)``) + or SQL Server's login-handshake queue. + + Correctness assertion: every task must return its own index — proves that + responses don't get cross-wired between concurrent HSTMTs. + """ + + async def _one_query(idx: int, sem: asyncio.Semaphore): + async with sem: + conn = connect(conn_str) + try: + cur = conn.cursor() + try: + await cur.execute_async( + "SELECT ? AS idx, CAST(? AS NVARCHAR(16)) AS label", + idx, + f"task-{idx}", + ) + row = await cur.fetch_async() + assert row is not None, f"task {idx}: fetch_async returned None" + assert row[0] == idx, f"task {idx}: got idx={row[0]}, expected {idx}" + assert row[1] == f"task-{idx}", ( + f"task {idx}: got label={row[1]!r}, expected 'task-{idx}'" + ) + return idx + finally: + cur.close() + finally: + conn.close() + + async def _run(): + sem = asyncio.Semaphore(CONCURRENT_LIMIT) + start = time.perf_counter() + results = await asyncio.gather( + *[_one_query(i, sem) for i in range(CONCURRENT_TASKS)] + ) + elapsed = time.perf_counter() - start + return results, elapsed + + results, elapsed = asyncio.run(_run()) + + # Every task returned its own index — proves no cross-wiring between HSTMTs. + assert sorted(results) == list(range(CONCURRENT_TASKS)), ( + f"missing / duplicate task indices in results (got {len(results)} results, " + f"unique={len(set(results))})" + ) + + # Loose sanity bound to catch runaway hangs; NOT a perf assertion. + assert elapsed < WALL_CLOCK_BUDGET_SECONDS, ( + f"{CONCURRENT_TASKS} concurrent async queries took {elapsed:.1f}s " + f"(>{WALL_CLOCK_BUDGET_SECONDS}s budget) — likely a hang or serialization bug" + ) + + print( + f"\n[async POC] {CONCURRENT_TASKS} concurrent queries " + f"(semaphore={CONCURRENT_LIMIT}) completed in {elapsed:.2f}s" + ) + + +# ============================================================================ +# MARS variant: same workload on ONE MARS-enabled connection with N cursors +# ============================================================================ +# Complements the one-connection-per-task test above. Here we open ONE +# connection with MARS enabled and create N cursors from it, then fire the +# same 100 concurrent execute_async + fetch_async pairs. +# +# Important distinctions from the connection-per-task variant: +# * All requests share a single DBC handle and a single underlying TCP +# socket. The ODBC driver multiplexes concurrent statements via MARS. +# This is materially different from true network-level parallelism — +# execution is interleaved rather than simultaneous. +# * Exercises the per-HSTMT SQL_ATTR_ASYNC_ENABLE / AsyncEnableGuard path +# when many HSTMTs share one DBC, which is the "N cursors on 1 conn" +# pattern documented in async POC spec §4.5. +# +# History: an earlier revision of this test used ``MultipleActiveResultSets=Yes`` +# in the connection string, which the bundled msodbcsql18 driver silently +# ignores (MARS was never actually enabled). Under that regime, concurrent +# async on the shared DBC crashed with SIGSEGV because the driver returned +# SQL_ERROR ("connection busy") for every second-in-flight statement and our +# concurrent DDBCSQLCheckError calls raced on shared DBC diagnostic state. +# +# Root-cause fix (2026-07): switched to the ODBC-standard alias +# ``MARS_Connection=Yes`` (the only MARS keyword msodbcsql18 actually honors) +# and added it to _ALLOWED_CONNECTION_STRING_PARAMS. With MARS genuinely on, +# concurrent async on shared DBC works cleanly (100 tasks, semaphore=16, +# ~0.13 s wall-clock on local Docker). +# ============================================================================ + + +def _with_mars(conn_str: str) -> str: + """Return ``conn_str`` with ``MARS_Connection=Yes`` appended, or the + original string if the caller already set a MARS keyword. + + Uses the ODBC-standard spelling ``MARS_Connection`` rather than the + SQL Server-native ``MultipleActiveResultSets`` because the bundled + msodbcsql18 driver silently ignores the latter (empirically verified; + see ``mssql_python/constants.py`` for the allowlist rationale). + + Returns an empty string as a sentinel when the caller explicitly *disabled* + MARS — the caller should then skip the test rather than override the + user's choice. + """ + lower = conn_str.lower() + if "multipleactiveresultsets=yes" in lower or "mars_connection=yes" in lower: + return conn_str + if "multipleactiveresultsets=no" in lower or "mars_connection=no" in lower: + return "" # sentinel — caller skips + sep = "" if conn_str.rstrip().endswith(";") else ";" + return f"{conn_str}{sep}MARS_Connection=Yes" + + +def test_100_concurrent_async_selects_on_single_mars_connection(conn_str, _async_capable): + """N async SELECTs concurrently on ONE MARS-enabled connection (N cursors). + + Complements ``test_100_concurrent_async_selects``. That test opens one + connection per task; this one opens ONE connection with + ``MARS_Connection=Yes`` and creates N cursors on it, then fires the same + workload through ``asyncio.gather``. + + Correctness assertions match the connection-per-task variant: every task + must return its own index, proving that MARS + per-HSTMT + ``SQL_ATTR_ASYNC_ENABLE`` toggling doesn't cross-wire results between + cursors sharing a single DBC. + + Note: MARS multiplexes over one TCP socket, so this is interleaved + (not truly parallel) execution — expect a different wall-clock profile + than the connection-per-task variant. + + Kept forward-compatible: if a caller explicitly disables MARS via + ``MARS_Connection=No`` in DB_CONNECTION_STRING, the test skips rather + than override that choice. The keyword-rejection fallback below + remains in place in case a future mssql-python version removes MARS + from the allowlist (unlikely). + """ + mars_conn_str = _with_mars(conn_str) + if not mars_conn_str: + pytest.skip( + "DB_CONNECTION_STRING explicitly disables MARS — cannot run " + "the same-connection concurrency test" + ) + + # Defensive: skip cleanly if a future mssql-python version rejects the + # MARS keyword in its allowlist. Not expected on the current codebase, + # where mars_connection is in _ALLOWED_CONNECTION_STRING_PARAMS. + try: + _probe_conn = connect(mars_conn_str) + _probe_conn.close() + except Exception as e: + msg = str(e).lower() + if "multipleactiveresultsets" in msg or "mars_connection" in msg or "unknown keyword" in msg: + pytest.skip( + f"mssql-python does not currently accept MARS in the connection " + f"string allowlist — skipping same-connection concurrency test. " + f"Underlying error: {e}" + ) + raise + + async def _one_query(idx: int, cur, sem: asyncio.Semaphore): + async with sem: + await cur.execute_async( + "SELECT ? AS idx, CAST(? AS NVARCHAR(16)) AS label", + idx, + f"task-{idx}", + ) + row = await cur.fetch_async() + assert row is not None, f"task {idx}: fetch_async returned None" + assert row[0] == idx, f"task {idx}: got idx={row[0]}, expected {idx}" + assert row[1] == f"task-{idx}", ( + f"task {idx}: got label={row[1]!r}, expected 'task-{idx}'" + ) + return idx + + async def _run(): + conn = connect(mars_conn_str) + try: + cursors = [conn.cursor() for _ in range(CONCURRENT_TASKS)] + try: + sem = asyncio.Semaphore(CONCURRENT_LIMIT) + start = time.perf_counter() + results = await asyncio.gather( + *[_one_query(i, cursors[i], sem) for i in range(CONCURRENT_TASKS)] + ) + elapsed = time.perf_counter() - start + return results, elapsed + finally: + for c in cursors: + c.close() + finally: + conn.close() + + results, elapsed = asyncio.run(_run()) + + # Every task returned its own index — proves no cross-wiring between + # HSTMTs that share a single DBC via MARS. + assert sorted(results) == list(range(CONCURRENT_TASKS)), ( + f"missing / duplicate task indices in MARS results " + f"(got {len(results)} results, unique={len(set(results))})" + ) + + # Loose sanity bound; NOT a perf assertion (MARS interleaves over one + # socket, so this is expected to be slower than the connection-per-task + # variant, but nowhere near WALL_CLOCK_BUDGET_SECONDS). + assert elapsed < WALL_CLOCK_BUDGET_SECONDS, ( + f"{CONCURRENT_TASKS} async queries on 1 MARS conn took {elapsed:.1f}s " + f"(>{WALL_CLOCK_BUDGET_SECONDS}s budget) — likely a hang or serialization bug" + ) + + print( + f"\n[async POC MARS] {CONCURRENT_TASKS} async queries on 1 MARS conn, " + f"{CONCURRENT_TASKS} cursors (semaphore={CONCURRENT_LIMIT}) " + f"completed in {elapsed:.2f}s" + ) + + +# ============================================================================ +# Additional stability tests (cases 1, 2, 4, 6, 7, 8, 9, 10, 11, 12) +# ============================================================================ +# These tests exercise the async surface across several axes: long queries, +# large result sets, event-loop non-blocking behavior, multi-batch fetches, +# and sequential vs concurrent invocation patterns. All safe patterns use +# one connection per concurrent task (see block comment on +# test_100_concurrent_async_selects for why). +# +# Tunables (env-overridable): +# ASYNC_TEST_LARGE_ROWS rows for the large-result-set test (default 5000) +# ASYNC_TEST_STRESS_CONCURRENCY N for the 1000-task stress test (default 1000) +# ASYNC_TEST_WAITFOR_SECONDS server-side delay for long-query tests (default 2) +# ============================================================================ + +LARGE_ROWS = int(os.getenv("ASYNC_TEST_LARGE_ROWS", "5000")) +# The fetch-heartbeat test needs a materially slower fetch (multi-hundred ms) +# for the 10 ms heartbeat to have time to tick a meaningful number of times. +# Default to ~10× LARGE_ROWS so even a fast local SQL Server produces enough +# TDS traffic to keep the fetch worker thread busy for a while. +HB_FETCH_ROWS = int(os.getenv("ASYNC_TEST_HB_FETCH_ROWS", "50000")) +STRESS_CONCURRENCY = int(os.getenv("ASYNC_TEST_STRESS_CONCURRENCY", "1000")) +WAITFOR_SECONDS = int(os.getenv("ASYNC_TEST_WAITFOR_SECONDS", "2")) +WAITFOR_SQL = f"WAITFOR DELAY '00:00:0{WAITFOR_SECONDS}'" + + +# ---------- Case 1: Execute long-running query asynchronously -------------- + + +def test_execute_async_long_running_query(conn_str, _async_capable): + """Run a ~2-second server-side WAITFOR through execute_async and verify + the row afterward comes back correctly. + + Establishes that the polling loop survives realistic query durations + (many polling iterations of SQL_STILL_EXECUTING) and doesn't lose the + result set. + """ + + async def _run(): + conn = connect(conn_str) + try: + cur = conn.cursor() + try: + start = time.perf_counter() + await cur.execute_async(f"{WAITFOR_SQL}; SELECT 42 AS n") + elapsed = time.perf_counter() - start + row = await cur.fetch_async() + return elapsed, row + finally: + cur.close() + finally: + conn.close() + + elapsed, row = asyncio.run(_run()) + assert row is not None and row[0] == 42 + assert elapsed >= WAITFOR_SECONDS - 0.2, ( + f"execute_async returned too quickly ({elapsed:.2f}s < {WAITFOR_SECONDS}s) — " + f"WAITFOR was probably not honored" + ) + + +# ---------- Case 2: Event loop continues while query executes -------------- + + +def test_event_loop_progresses_during_execute_async(conn_str, _async_capable): + """A background heartbeat coroutine must keep ticking while a long query + is being polled by the C++ executor thread. + + Proves that ``execute_async`` does NOT starve the asyncio event loop: + the polling loop runs in the run_in_executor worker thread with the + GIL released, so the main thread's event loop stays free to schedule + other coroutines. + """ + + async def _run(): + conn = connect(conn_str) + heartbeats = 0 + + async def heartbeat(): + nonlocal heartbeats + while True: + await asyncio.sleep(0.01) # 10ms tick + heartbeats += 1 + + hb_task = asyncio.create_task(heartbeat()) + try: + cur = conn.cursor() + try: + await cur.execute_async(f"{WAITFOR_SQL}; SELECT 1") + _ = await cur.fetch_async() + finally: + cur.close() + finally: + hb_task.cancel() + try: + await hb_task + except asyncio.CancelledError: + pass + conn.close() + return heartbeats + + heartbeats = asyncio.run(_run()) + # Expect ~100 ticks/second × WAITFOR_SECONDS with plenty of margin for + # scheduling jitter. A value < ~30 (i.e. 15% of ideal) indicates the + # event loop was starved. + min_expected = max(30, WAITFOR_SECONDS * 30) + assert heartbeats >= min_expected, ( + f"event loop appears to have been blocked: only {heartbeats} heartbeats " + f"in {WAITFOR_SECONDS}s (expected >= {min_expected}) — " + f"execute_async is probably not releasing the event loop" + ) + print( + f"\n[async POC] heartbeat during {WAITFOR_SECONDS}s execute_async: " + f"{heartbeats} ticks (>= {min_expected} required)" + ) + + +# ---------- Case 4: 1000 async executes on different connections ----------- + + +@pytest.mark.stress +def test_1000_async_executes_on_different_connections(conn_str, _async_capable): + """Scale the connection-per-task workload up to 1000 tasks. + + Marked ``stress`` (excluded from the default pytest run per pytest.ini) + because opening 1000 connections stresses SQL Server's login-handshake + queue and the local machine's ephemeral port pool. Bounded by + ``ASYNC_TEST_MAX_INFLIGHT`` (default 16) so at most that many are being + established at any instant. + """ + + async def _one_query(idx: int, sem: asyncio.Semaphore): + async with sem: + conn = connect(conn_str) + try: + cur = conn.cursor() + try: + await cur.execute_async("SELECT ? AS idx", idx) + row = await cur.fetch_async() + assert row is not None and row[0] == idx + return idx + finally: + cur.close() + finally: + conn.close() + + async def _run(): + sem = asyncio.Semaphore(CONCURRENT_LIMIT) + start = time.perf_counter() + results = await asyncio.gather( + *[_one_query(i, sem) for i in range(STRESS_CONCURRENCY)] + ) + return results, time.perf_counter() - start + + results, elapsed = asyncio.run(_run()) + assert sorted(results) == list(range(STRESS_CONCURRENCY)) + print( + f"\n[async POC stress] {STRESS_CONCURRENCY} async queries " + f"(semaphore={CONCURRENT_LIMIT}) completed in {elapsed:.2f}s" + ) + + +# ---------- Case 6: Fetch multiple rows asynchronously --------------------- + + +def test_fetch_async_returns_batch_of_rows(conn_str, _async_capable): + """``fetch_async(size=N)`` should return a ``List[Row]`` with the first N + rows of the result set.""" + + async def _run(): + conn = connect(conn_str) + try: + cur = conn.cursor() + try: + # Simple VALUES clause yields 5 rows deterministically. + await cur.execute_async( + "SELECT * FROM (VALUES (1),(2),(3),(4),(5)) AS t(n)" + ) + rows = await cur.fetch_async(5) + return rows + finally: + cur.close() + finally: + conn.close() + + rows = asyncio.run(_run()) + assert isinstance(rows, list) and len(rows) == 5 + assert [r[0] for r in rows] == [1, 2, 3, 4, 5] + + +# ---------- Case 7: Fetch large result set --------------------------------- + + +def test_fetch_async_large_result_set(conn_str, _async_capable): + """``fetch_async(-1)`` should return all rows for a large result set. + + Uses a CROSS JOIN of ``sys.all_objects`` to generate ``LARGE_ROWS`` + rows, which produces a multi-KB result set spanning multiple TDS + packets. Verifies the async fetch path handles multi-batch responses + correctly. + """ + + async def _run(): + conn = connect(conn_str) + try: + cur = conn.cursor() + try: + await cur.execute_async( + f"SELECT TOP {LARGE_ROWS} " + f"ROW_NUMBER() OVER (ORDER BY a.object_id) AS n, " + f"CAST(a.name AS NVARCHAR(128)) AS obj_name " + f"FROM sys.all_objects a CROSS JOIN sys.all_objects b" + ) + rows = await cur.fetch_async(-1) + return rows + finally: + cur.close() + finally: + conn.close() + + rows = asyncio.run(_run()) + assert len(rows) == LARGE_ROWS, f"expected {LARGE_ROWS} rows, got {len(rows)}" + # Row numbering is 1..LARGE_ROWS and sequential. + assert rows[0][0] == 1 + assert rows[-1][0] == LARGE_ROWS + # Spot-check that name column decoded correctly (non-empty string). + assert isinstance(rows[0][1], str) and len(rows[0][1]) > 0 + + +# ---------- Case 8: Fetch does not block event loop ------------------------ + + +def test_event_loop_progresses_during_fetch_async(conn_str, _async_capable): + """A background heartbeat must keep ticking during a large ``fetch_async``. + + Symmetric to test_event_loop_progresses_during_execute_async but for + the fetch path. The fetch runs in an executor thread with the GIL + released, so the event loop should stay live. + + Uses ``HB_FETCH_ROWS`` (default 50000) rather than ``LARGE_ROWS`` so the + fetch is long enough (multi-hundred ms typical) for the 10 ms heartbeat + to tick a meaningful number of times. + """ + + async def _run(): + conn = connect(conn_str) + heartbeats = 0 + + async def heartbeat(): + nonlocal heartbeats + while True: + await asyncio.sleep(0.01) + heartbeats += 1 + + try: + cur = conn.cursor() + try: + # Prep the result set synchronously (fast). The interesting + # timing is on fetch, not execute. + await cur.execute_async( + f"SELECT TOP {HB_FETCH_ROWS} " + f"ROW_NUMBER() OVER (ORDER BY a.object_id) AS n, " + f"CAST(a.name AS NVARCHAR(128)) AS obj_name " + f"FROM sys.all_objects a CROSS JOIN sys.all_objects b" + ) + hb_task = asyncio.create_task(heartbeat()) + fetch_start = time.perf_counter() + rows = await cur.fetch_async(-1) + fetch_elapsed = time.perf_counter() - fetch_start + hb_task.cancel() + try: + await hb_task + except asyncio.CancelledError: + pass + return rows, heartbeats, fetch_elapsed + finally: + cur.close() + finally: + conn.close() + + rows, heartbeats, fetch_elapsed = asyncio.run(_run()) + assert len(rows) == HB_FETCH_ROWS + # If the fetch completes in less than ~50 ms the test is inconclusive + # (the heartbeat's 10 ms timer may not have fired even once even in an + # ideal system). Skip the tick-count assertion in that regime — a + # sub-50 ms fetch on a local SQL Server means the event loop wasn't + # blocked long enough to matter either way. + if fetch_elapsed < 0.05: + pytest.skip( + f"fetch_async completed too fast ({fetch_elapsed*1000:.0f}ms) to " + f"meaningfully measure event-loop responsiveness — increase " + f"ASYNC_TEST_HB_FETCH_ROWS on faster hardware" + ) + # Require at least ~1 tick per 30 ms of fetch (very loose to tolerate CI + # jitter). If the event loop were fully blocked we'd expect 0 ticks. + min_expected = max(1, int(fetch_elapsed * 30)) + assert heartbeats >= min_expected, ( + f"event loop blocked during fetch_async: {heartbeats} ticks in " + f"{fetch_elapsed*1000:.0f}ms (>= {min_expected} required)" + ) + print( + f"\n[async POC] heartbeat during {HB_FETCH_ROWS}-row fetch_async " + f"({fetch_elapsed*1000:.0f}ms): {heartbeats} ticks" + ) + + +# ---------- Case 9: Multiple concurrent execute_async operations ---------- + + +def test_multiple_concurrent_execute_async_small_batch(conn_str, _async_capable): + """A small variant of the connection-per-task pattern (N=10). + + Complementary to the 100-task test — kept small so it's included in + quick smoke runs. Each task uses its own connection; concurrent + execute on cursors sharing a DBC is intentionally NOT tested here + because it hits the ODBC no-MARS restriction (see block comment on + test_100_concurrent_async_selects_on_single_mars_connection). + """ + N = 10 + + async def _one(idx): + conn = connect(conn_str) + try: + cur = conn.cursor() + try: + await cur.execute_async("SELECT ? AS v", idx * 10) + row = await cur.fetch_async() + return row[0] + finally: + cur.close() + finally: + conn.close() + + async def _run(): + return await asyncio.gather(*[_one(i) for i in range(N)]) + + results = asyncio.run(_run()) + assert sorted(results) == [i * 10 for i in range(N)] + + +# ---------- Case 10: Execute async then fetch async ----------------------- + + +def test_execute_async_followed_by_fetch_async(conn_str, _async_capable): + """Verify the natural ``execute_async`` → ``fetch_async`` sequence for + each of the three ``fetch_async`` modes on the SAME cursor. + + Distinct from ``test_single_execute_async_and_fetch_async`` (which + exercises only the single-row mode) by covering all of size=None, + size=positive, and size=-1 on the SAME cursor after independent + executes. + """ + + async def _run(): + conn = connect(conn_str) + try: + cur = conn.cursor() + try: + # Mode 1: size=None → single Row + await cur.execute_async("SELECT 100 AS v") + r = await cur.fetch_async() + assert r is not None and r[0] == 100 + + # Mode 2: size=positive → List[Row] + await cur.execute_async( + "SELECT * FROM (VALUES (1),(2),(3)) AS t(n)" + ) + rows = await cur.fetch_async(3) + assert [x[0] for x in rows] == [1, 2, 3] + + # Mode 3: size=-1 → List[Row] (all) + await cur.execute_async( + "SELECT * FROM (VALUES ('a'),('b'),('c'),('d')) AS t(s)" + ) + rows = await cur.fetch_async(-1) + assert [x[0] for x in rows] == ["a", "b", "c", "d"] + finally: + cur.close() + finally: + conn.close() + + asyncio.run(_run()) + + +# ---------- Case 11: Multiple concurrent fetch_async ---------------------- + + +def test_multiple_concurrent_fetch_async_across_connections(conn_str, _async_capable): + """N connections, each with a pre-executed statement, then + ``fetch_async`` fired concurrently across all of them. + + Complements the execute-side concurrency tests: this exercises the + fetch code path (DDBCSQLFetchOneAsync) under concurrent gather. Uses + one connection per task so each fetch operates on its own DBC (safe + per the root-cause investigation). + """ + N = 20 + + async def _worker(idx): + conn = connect(conn_str) + cur = None + try: + cur = conn.cursor() + # Execute first (sequentially per task), then fetch in the + # concurrent phase below. + await cur.execute_async("SELECT ? AS v", idx * 7) + return cur, conn # keep alive for the fetch phase + except Exception: + if cur is not None: + cur.close() + conn.close() + raise + + async def _run(): + # Phase 1: prepare N cursors with statements ready to fetch. + prep = await asyncio.gather(*[_worker(i) for i in range(N)]) + # prep is [(cur, conn), ...] + try: + # Phase 2: concurrent fetch_async across all N. + rows = await asyncio.gather(*[c.fetch_async() for c, _ in prep]) + return [r[0] for r in rows] + finally: + for c, conn in prep: + c.close() + conn.close() + + results = asyncio.run(_run()) + assert sorted(results) == [i * 7 for i in range(N)] + + +# ---------- Case 12: Sequential execute_async calls ----------------------- + + +def test_sequential_execute_async_on_same_cursor(conn_str, _async_capable): + """Run multiple ``execute_async`` calls back-to-back on the SAME cursor. + + Verifies that the ``AsyncEnableGuard`` correctly toggles + ``SQL_ATTR_ASYNC_ENABLE`` OFF at the end of each call, so subsequent + async executes on the same HSTMT are not affected by leftover state. + Also verifies that ``_finalize_execute_async`` correctly resets cursor + metadata (description, rowcount, column cache) between calls. + """ + + async def _run(): + conn = connect(conn_str) + try: + cur = conn.cursor() + try: + for i in range(5): + await cur.execute_async("SELECT ? AS v", i) + row = await cur.fetch_async() + assert row is not None + assert row[0] == i, f"call {i}: got {row[0]}" + # Second fetch_async on the exhausted result set + # returns None (consistent with sync fetchone semantics). + assert await cur.fetch_async() is None + finally: + cur.close() + finally: + conn.close() + + asyncio.run(_run()) + + +# ============================================================================ +# MARS stability test (multiple iterations, high concurrency) +# ============================================================================ +# Complements test_100_concurrent_async_selects_on_single_mars_connection +# (a single-run correctness test) with a stability test designed to catch +# intermittent races, leaks, or state corruption that a single-run test +# might miss. +# +# Differences from the correctness test: +# * NO semaphore — all N statements truly in flight at once, not bounded +# to 16. +# * ITERATIONS repetitions of the workload reuse the SAME MARS connection +# across rounds, catching leaks / state corruption that only manifest +# after N runs (e.g. an HSTMT counter that overflows, an internal MARS +# session table that isn't reclaimed, or a slow diagnostic-record leak). +# +# Note on MARS + parallelism: MARS multiplexes multiple result sets over +# ONE TCP socket, but SQL Server assigns a single server-side worker +# thread per session (SPID). So 100 statements on one MARS connection do +# NOT execute in true parallel server-side — they queue on the shared +# session's worker. Real network-level parallelism requires N connections +# (see test_100_concurrent_async_selects). This test therefore stresses +# CLIENT-side concurrency (100 in-flight coroutines, executor threads +# calling into the MARS driver simultaneously) rather than server-side +# throughput. +# +# Marked @pytest.mark.stress (excluded from the default pytest run per +# pytest.ini). +# ============================================================================ + +MARS_STABILITY_N = int(os.getenv("ASYNC_TEST_MARS_STABILITY_N", "100")) +MARS_STABILITY_ITERS = int(os.getenv("ASYNC_TEST_MARS_STABILITY_ITERS", "10")) + + +@pytest.mark.stress +def test_mars_stability_100_cursors_high_concurrency_multi_iteration( + conn_str, _async_capable +): + """Stability: N cursors on ONE MARS connection, all truly concurrent, + repeated across ITERATIONS rounds on the SAME connection. + + Each iteration: + 1. Opens ``N`` cursors on the (single, long-lived) MARS connection. + 2. Fires all ``N`` execute_async + fetch_async pairs through a single + ``asyncio.gather`` with NO semaphore (all N in flight at once). + 3. Verifies every task returned its own (idx, iteration) pair — proves + no cross-wiring across the ``ITERATIONS × N`` combined ops, and no + leftover state from prior iterations. + 4. Closes all N cursors before the next iteration starts. + + Correctness AND stability are both asserted: the connection must remain + usable across all ITERATIONS rounds (a per-iteration leak or state + corruption would show up as a failure in later rounds). + """ + mars_conn_str = _with_mars(conn_str) + if not mars_conn_str: + pytest.skip("DB_CONNECTION_STRING explicitly disables MARS") + + async def _one_iteration(iteration: int, conn): + cursors = [conn.cursor() for _ in range(MARS_STABILITY_N)] + try: + async def one(i): + await cursors[i].execute_async( + "SELECT ? AS idx, ? AS iter", i, iteration + ) + row = await cursors[i].fetch_async() + return (row[0], row[1]) + + # NO semaphore — all N truly in flight simultaneously. The + # asyncio default ThreadPoolExecutor caps the actual number of + # concurrent OS threads calling into the ODBC driver, so this + # is bounded in practice regardless of N. + return await asyncio.gather(*[one(i) for i in range(MARS_STABILITY_N)]) + finally: + for c in cursors: + c.close() + + async def _run(): + conn = connect(mars_conn_str) + try: + start = time.perf_counter() + all_results = [] + for it in range(MARS_STABILITY_ITERS): + iter_start = time.perf_counter() + iter_results = await _one_iteration(it, conn) + all_results.append((it, iter_results, time.perf_counter() - iter_start)) + return all_results, time.perf_counter() - start + finally: + conn.close() + + all_results, total_elapsed = asyncio.run(_run()) + + # Correctness: every iteration returned N rows, each carrying its own + # (idx, iteration) tuple. Cross-wiring or corruption would break this. + for it, iter_results, _iter_elapsed in all_results: + assert len(iter_results) == MARS_STABILITY_N, ( + f"iter {it}: got {len(iter_results)} results, " + f"expected {MARS_STABILITY_N}" + ) + got_indices = sorted(r[0] for r in iter_results) + assert got_indices == list(range(MARS_STABILITY_N)), ( + f"iter {it}: missing / duplicate indices — cross-wiring bug? " + f"got unique={len(set(got_indices))}" + ) + for idx, iter_val in iter_results: + assert iter_val == it, ( + f"iter {it}, idx {idx}: got iter_val={iter_val} — " + f"cross-iteration cross-wiring" + ) + + # Loose sanity bound to catch runaway hangs. MARS serializes statements + # server-side per session, so we expect the total to scale with + # ITERATIONS × N × per-query overhead. On local Docker this is a few + # ms per query; on production networks add round-trip latency. Very + # generous upper bound: 5 minutes total. + assert total_elapsed < 300, ( + f"MARS stability test took {total_elapsed:.1f}s (>300s) — " + f"probable hang or catastrophic serialization" + ) + + per_iter = [f"{elapsed*1000:.0f}ms" for _, _, elapsed in all_results] + print( + f"\n[async POC MARS stability] {MARS_STABILITY_ITERS} iters × " + f"{MARS_STABILITY_N} concurrent cursors on 1 MARS conn " + f"completed in {total_elapsed:.2f}s. Per-iter: {per_iter}" + )