From b4749afe51ad788baf8398c7aca89ad9244d7bd4 Mon Sep 17 00:00:00 2001 From: Jahnvi Thakkar Date: Thu, 22 Oct 2026 12:03:37 +0530 Subject: [PATCH 1/2] FIX: Bound pooled connection checkout on unresponsive sockets Apply a temporary native connection timeout during pooled reset and autocommit configuration, then restore it on success. Replace a failed candidate instead of probing all idle sockets sequentially. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- README.md | 2 + mssql_python/connection.py | 15 +- mssql_python/pybind/connection/connection.cpp | 41 +++- mssql_python/pybind/connection/connection.h | 2 +- .../pybind/connection/connection_pool.cpp | 47 +++- tests/test_006_exceptions.py | 13 +- tests/test_009_pooling.py | 205 +++++++++++++++++- 7 files changed, 296 insertions(+), 29 deletions(-) diff --git a/README.md b/README.md index 40fc409e7..b87090693 100644 --- a/README.md +++ b/README.md @@ -129,6 +129,8 @@ EntraID authentication is now fully supported on MacOS and Linux but with certai The Microsoft mssql_python driver provides built-in support for connection pooling, which helps improve performance and scalability by reusing active database connections instead of creating a new connection for every request. This feature is enabled by default. For more information, refer [Connection Pooling Wiki](https://github.com/microsoft/mssql-python/wiki/Connection#connection-pooling). +Pooled checkout validates and resets a reused connection with a temporary native connection timeout. It uses the positive login timeout from `connect(timeout=N)` or, with higher precedence, `attrs_before[SQL_ATTR_LOGIN_TIMEOUT]`; if neither specifies a positive value, validation uses 15 seconds. A failed validation discards that candidate and attempts a new physical connection instead of probing every idle socket sequentially. The previous connection timeout is restored after successful checkout, and query timeouts remain controlled separately by `conn.timeout`. These are driver operation timeouts, not a strict end-to-end `connect()` deadline: cancellation/cleanup and a replacement login can add time. Pool return still requires network I/O and is not bounded by the checkout validation timeout. + > **Interactive / Device-code authentication in multi-user processes:** For `ActiveDirectoryInteractive` and `ActiveDirectoryDeviceCode`, the driver caches one credential instance per authentication type for the lifetime of the process so that pooled reconnects refresh silently instead of re-prompting. As a result, every interactive connection opened in the same process shares that one signed-in account — the first user to authenticate. If your application serves multiple end users from a single process, do **not** rely on interactive/device-code auth to isolate them; instead supply your own per-user token via a token provider so each user's connection is keyed to their own identity. > **Bring-your-own token (`token_provider` / raw access token):** When you supply your own access token (via a token provider or a raw token in `attrs_before`), managing its lifetime is **your application's responsibility**. The driver stores a pooled connection under the identity that opened it, but it cannot refresh a token it did not mint. If a pooled connection is reused after its token has expired, the reconnect (ODBC's implicit connection resiliency) will fail. Ensure your token provider returns a currently-valid token on every call, and account for token expiry in your own logic. The driver's built-in expiry-aware checkout (refreshing a pooled connection whose token is near expiry) only applies to Managed Identity and interactive/device-code pools, whose tokens the driver can silently re-mint — not to bring-your-own tokens. diff --git a/mssql_python/connection.py b/mssql_python/connection.py index 4d0b4cd72..f65c26d39 100644 --- a/mssql_python/connection.py +++ b/mssql_python/connection.py @@ -417,7 +417,11 @@ def __init__( query timeout. An explicit ``attrs_before[SQL_ATTR_LOGIN_TIMEOUT]`` takes precedence over this value. (pyodbc migration note: pyodbc lets a positive ``timeout=`` override ``attrs_before``; here the - explicit ``attrs_before`` entry wins.) + explicit ``attrs_before`` entry wins.) Pooled candidate reset + also uses this positive timeout, or 15 seconds when zero/default, + temporarily as SQL_ATTR_CONNECTION_TIMEOUT. Driver cancellation + and a replacement login can add time; this is not an overall + connect deadline. native_uuid (bool, optional): Controls whether UNIQUEIDENTIFIER columns return uuid.UUID objects (True) or str (False) for cursors created from this connection. None (default) defers to the module-level ``mssql_python.native_uuid`` setting (True). @@ -858,17 +862,22 @@ def _token_factory(): _provider = ProviderManager.ensure_available() ddbc_bindings._set_odbc_provider(_provider) + # Configure a reused connection's mode inside native checkout, while + # its temporary connection timeout still bounds attribute-setting I/O. + native_attrs = self._attrs_before.copy() + native_attrs[ConstantsDDBC.SQL_ATTR_AUTOCOMMIT.value] = int(bool(autocommit)) try: self._conn = ddbc_bindings.Connection( self.connection_str, self._pooling, - self._attrs_before, + native_attrs, self._pool_key, self._token_factory, ) + if self._conn.get_autocommit() != autocommit: + self.setautocommit(autocommit) except RuntimeError as e: _raise_connection_error(e) - self.setautocommit(autocommit) # Register this connection for cleanup before Python shutdown # This ensures ODBC handles are freed in correct order, preventing leaks diff --git a/mssql_python/pybind/connection/connection.cpp b/mssql_python/pybind/connection/connection.cpp index 9c4304d2f..bb56ef284 100644 --- a/mssql_python/pybind/connection/connection.cpp +++ b/mssql_python/pybind/connection/connection.cpp @@ -600,11 +600,31 @@ bool Connection::isAlive() const { return SQL_SUCCEEDED(ret) && status == SQL_CD_FALSE; } -bool Connection::reset() { +bool Connection::reset(unsigned int timeoutSeconds, bool autocommit) { if (!_dbcHandle) { ThrowStdException("Connection handle not allocated"); } clearResultMetadata(); + // Login and statement timeouts do not bound connection-attribute calls. + // Keep a finite connection timeout through reset and lease configuration; + // retain it on failure so discarding the candidate is bounded too (#854). + SQLULEN originalTimeout = 0; + SQLRETURN ret = SQLGetConnectAttr_ptr(_dbcHandle->get(), SQL_ATTR_CONNECTION_TIMEOUT, + &originalTimeout, 0, nullptr); + if (!SQL_SUCCEEDED(ret)) { + LOG("Failed to read pooled connection timeout (ret=%d). Marking as dead.", ret); + return false; + } + { + py::gil_scoped_release release; + ret = SQLSetConnectAttr_ptr( + _dbcHandle->get(), SQL_ATTR_CONNECTION_TIMEOUT, + reinterpret_cast(static_cast(timeoutSeconds)), SQL_IS_UINTEGER); + } + if (!SQL_SUCCEEDED(ret)) { + LOG("Failed to set pooled connection timeout (ret=%d). Marking as dead.", ret); + return false; + } 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 @@ -618,7 +638,6 @@ bool Connection::reset() { // Connection (_attrStringBuffers / _attrBytesBuffers) are intentionally // retained here. Clearing them would reintroduce issue #594 in a new // form (UAF during silent reconnect). - SQLRETURN ret; { // Release the GIL around the ODBC call for consistency with the // other connection-attribute paths; some attributes can block. @@ -631,9 +650,7 @@ bool Connection::reset() { return false; } - // SQL_ATTR_RESET_CONNECTION does NOT reset the transaction isolation level. - // Explicitly reset it to the default (SQL_TXN_READ_COMMITTED) to prevent - // isolation level settings from leaking between pooled connection usages. + // SQL_ATTR_RESET_CONNECTION does not reset transaction isolation. LOG("Resetting transaction isolation level to READ COMMITTED"); { py::gil_scoped_release release; @@ -641,7 +658,19 @@ bool Connection::reset() { (SQLPOINTER)SQL_TXN_READ_COMMITTED, SQL_IS_INTEGER); } if (!SQL_SUCCEEDED(ret)) { - LOG("Failed to reset transaction isolation level (ret=%d). Marking as dead.", ret); + ErrorInfo error = SQLReadError(SQL_HANDLE_DBC, _dbcHandle->get(), ret); + LOG("Pooled connection validation failed: %s", error.ddbcErrorMsg.c_str()); + return false; + } + setAutocommit(autocommit); + { + py::gil_scoped_release release; + ret = SQLSetConnectAttr_ptr( + _dbcHandle->get(), SQL_ATTR_CONNECTION_TIMEOUT, + reinterpret_cast(originalTimeout), SQL_IS_UINTEGER); + } + if (!SQL_SUCCEEDED(ret)) { + LOG("Failed to restore pooled connection timeout (ret=%d). Marking as dead.", ret); return false; } diff --git a/mssql_python/pybind/connection/connection.h b/mssql_python/pybind/connection/connection.h index f43613842..7eaa4dcce 100644 --- a/mssql_python/pybind/connection/connection.h +++ b/mssql_python/pybind/connection/connection.h @@ -58,7 +58,7 @@ class Connection { // Check whether autocommit is enabled. bool getAutocommit() const; bool isAlive() const; - bool reset(); + bool reset(unsigned int timeoutSeconds, bool autocommit = true); void prepareForPool(bool transactionAlreadyRolledBack = false); void updateLastUsed(); std::chrono::steady_clock::time_point lastUsed() const; diff --git a/mssql_python/pybind/connection/connection_pool.cpp b/mssql_python/pybind/connection/connection_pool.cpp index 3ea3ccc5f..1c1089ead 100644 --- a/mssql_python/pybind/connection/connection_pool.cpp +++ b/mssql_python/pybind/connection/connection_pool.cpp @@ -6,6 +6,7 @@ #include #include #include +#include #include // Logging uses LOG() macro for all diagnostic output @@ -60,6 +61,22 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt const py::dict& attrs_before, const py::object& token_factory) { PERF_TIMER("ConnectionPool::acquire"); + // A login timeout does not cover SQLSetConnectAttr on an existing socket. + // Use it for bounded candidate reset as well; zero/unspecified + // retains a finite validation default rather than an OS TCP retry budget. + unsigned int validation_timeout = 15; + const py::int_ login_timeout_attr(SQL_ATTR_LOGIN_TIMEOUT); + if (attrs_before.contains(login_timeout_attr)) { + const auto requested = attrs_before[login_timeout_attr].cast(); + if (requested > 0) { + validation_timeout = static_cast(std::min( + requested, std::numeric_limits::max())); + } + } + const py::int_ autocommit_attr(SQL_ATTR_AUTOCOMMIT); + const bool autocommit = attrs_before.contains(autocommit_attr) + ? attrs_before[autocommit_attr].cast() != SQL_AUTOCOMMIT_OFF + : true; std::vector> to_disconnect; std::shared_ptr valid_conn = nullptr; bool needs_connect = false; @@ -131,6 +148,7 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt // Validate the candidate outside the mutex. bool reuse_candidate = false; + bool validation_attempted = false; try { if (token_factory && !token_factory.is_none() && candidate->isTokenNearExpiry(TOKEN_EXPIRY_THRESHOLD_SECS)) { @@ -169,7 +187,10 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt // is taken in practice. We favor never handing out a token // that may expire mid-query over avoiding the churn. candidate->setTokenExpiry(fresh_expiry); - reuse_candidate = candidate->isAlive() && candidate->reset(); + if (candidate->isAlive()) { + validation_attempted = true; + reuse_candidate = candidate->reset(validation_timeout, autocommit); + } if (!reuse_candidate) { // The token is still valid but the socket is dead // (isAlive()/reset() failed). We already minted the @@ -213,7 +234,10 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt } } } else { - reuse_candidate = candidate->isAlive() && candidate->reset(); + if (candidate->isAlive()) { + validation_attempted = true; + reuse_candidate = candidate->reset(validation_timeout, autocommit); + } } } catch (const std::exception& ex) { LOG("Candidate connection validation failed: %s", ex.what()); @@ -232,10 +256,10 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt if (_current_size > 0) --_current_size; } - // If a rotated token was captured, reserve a slot and reopen with it - // immediately instead of churning through the remaining candidates - // (which hold the same stale token and would all be discarded anyway). - if (have_pending_token) { + // After a reset failure, open a replacement instead of spending another + // timeout on each idle socket. Token rotation likewise opens directly + // using the captured token. + if (have_pending_token || validation_attempted) { std::lock_guard lock(_mutex); if (_current_size < _max_size) { // Reserve the slot here but construct the Connection outside @@ -246,11 +270,12 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt needs_connect = true; break; } - // Pool momentarily full; fall through and retry the loop. On the - // retry another near-expiry candidate may re-invoke the factory and - // overwrite pending_attrs/pending_expiry with a newer token. That - // needs a full pool AND a simultaneous rotation, is rare, and is - // harmless: we simply reopen with the most recently minted token. + if (validation_attempted) { + throw std::runtime_error( + "ConnectionPool::acquire: pool size limit reached"); + } + // A concurrent acquire reserved the slot freed by token rotation. + // Retry token comparison with another idle candidate. } } diff --git a/tests/test_006_exceptions.py b/tests/test_006_exceptions.py index 371b9662f..88b36fbf1 100644 --- a/tests/test_006_exceptions.py +++ b/tests/test_006_exceptions.py @@ -301,11 +301,12 @@ def test_close_cleans_up_after_autocommit_read_failure(): from unittest.mock import MagicMock, patch mock_conn = MagicMock() - mock_conn.get_autocommit.side_effect = RuntimeError("SQLSTATE:08S01:Communication link failure") + mock_conn.get_autocommit.return_value = False with patch("mssql_python.connection.ddbc_bindings.Connection", return_value=mock_conn): conn = connect("Server=testserver;Database=mydb;Trusted_Connection=yes;") + mock_conn.get_autocommit.side_effect = RuntimeError("SQLSTATE:08S01:Communication link failure") with pytest.raises(OperationalError, match="Communication link failure"): conn.close() @@ -356,16 +357,16 @@ def test_native_close_error_takes_precedence_over_preclose_failure(preclose_fail from unittest.mock import MagicMock, patch mock_conn = MagicMock() + mock_conn.get_autocommit.return_value = False + + with patch("mssql_python.connection.ddbc_bindings.Connection", return_value=mock_conn): + conn = connect("Server=testserver;Database=mydb;Trusted_Connection=yes;") + if preclose_failure == "autocommit": mock_conn.get_autocommit.side_effect = RuntimeError("SQLSTATE:08S01:Autocommit read failed") else: - mock_conn.get_autocommit.return_value = False mock_conn.rollback.side_effect = RuntimeError("SQLSTATE:08S01:Rollback failed") mock_conn.close.side_effect = RuntimeError("SQLSTATE:08003:Native close failed") - - with patch("mssql_python.connection.ddbc_bindings.Connection", return_value=mock_conn): - conn = connect("Server=testserver;Database=mydb;Trusted_Connection=yes;") - with pytest.raises(OperationalError, match="Native close failed"): conn.close() diff --git a/tests/test_009_pooling.py b/tests/test_009_pooling.py index 28b7465bc..e14526168 100644 --- a/tests/test_009_pooling.py +++ b/tests/test_009_pooling.py @@ -308,7 +308,8 @@ def test_autocommit_explicit_transaction_is_rolled_back_on_pool_checkin(conn_str ) -def test_connection_pooling_isolation_level_reset(conn_str): +@pytest.mark.parametrize("use_sql", [False, True], ids=["attribute", "sql"]) +def test_connection_pooling_isolation_level_reset(conn_str, use_sql): """Test that pooling correctly resets session state for isolation level. This test verifies that when a connection is returned to the pool and then @@ -325,7 +326,11 @@ def test_connection_pooling_isolation_level_reset(conn_str): conn1 = connect(conn_str) # Set isolation level to SERIALIZABLE (non-default) - conn1.set_attr(mssql_python.SQL_ATTR_TXN_ISOLATION, mssql_python.SQL_TXN_SERIALIZABLE) + if use_sql: + with conn1.cursor() as cur: + cur.execute("SET TRANSACTION ISOLATION LEVEL SERIALIZABLE") + else: + conn1.set_attr(mssql_python.SQL_ATTR_TXN_ISOLATION, mssql_python.SQL_TXN_SERIALIZABLE) # Verify the isolation level was set cursor1 = conn1.cursor() @@ -851,6 +856,202 @@ def phys_id(conn): # ============================================================================= +@pytest.mark.parametrize( + "pool_size,autocommit,options", + [ + pytest.param(pool_size, autocommit, options, id=f"{name}-{pool_size}-{autocommit}") + for name, options in ( + ("constructor", {"timeout": 1}), + ( + "attrs_before", + {"timeout": 30, "attrs_before": {mssql_python.SQL_ATTR_LOGIN_TIMEOUT: 1}}, + ), + ) + for pool_size in (1, 3) + for autocommit in (False, True) + ] + + [pytest.param(3, False, {"timeout": 0}, id="finite-default")], +) +def test_pool_checkout_validation_timeout(conn_str, autocommit, pool_size, options): + """GH #854: stalled checkout validation must time out and open a fresh socket.""" + _run_in_subprocess( + """ + import os + import socket + import sys + import threading + import time + from mssql_python import connect, pooling + from mssql_python.connection_string_parser import _ConnectionStringParser + from mssql_python.connection_string_builder import _ConnectionStringBuilder + + params = _ConnectionStringParser._normalize_params( + _ConnectionStringParser()._parse(os.environ["DB_CONNECTION_STRING"]) + ) + server = params["Server"] + if server.lower().startswith("tcp:"): + server = server[4:] + host, _, port = server.partition(",") + if "\\\\" in host: + print("Socket proxy requires an explicit TCP endpoint", file=sys.stderr) + sys.exit(77) + port = int(port) if port else 1433 + pool_size = POOL_SIZE + autocommit = AUTOCOMMIT + options = TIMEOUT_OPTIONS + use_default = options.get("timeout") == 0 + gate = threading.Event() + gate.set() + probed = set() + stopping = threading.Event() + sockets = [] + threads = [] + errors = [] + listener = socket.socket() + loopback = "127.0.0.1" # DevSkim: ignore DS162092 + listener.bind((loopback, 0)) + listener.listen() + listener.settimeout(0.2) + + def relay(source, destination, allowed, index): + try: + while not stopping.is_set(): + data = source.recv(65536) + if not data: + return + if not allowed.is_set(): + probed.add(index) + allowed.wait() + if not stopping.is_set(): + destination.sendall(data) + except (ConnectionResetError, BrokenPipeError): + return + except OSError as exc: + if not stopping.is_set(): + errors.append(str(exc)) + + def accept(): + try: + while not stopping.is_set(): + try: + client, _ = listener.accept() + except socket.timeout: + continue + sockets.append(client) + upstream = socket.create_connection((host, port), timeout=10) + upstream.settimeout(None) + sockets.append(upstream) + index = len(sockets) // 2 + allowed = gate if index <= pool_size else threading.Event() + if allowed is not gate: + allowed.set() + for source, destination in ((client, upstream), (upstream, client)): + thread = threading.Thread( + target=relay, args=(source, destination, allowed, index), daemon=True + ) + threads.append(thread) + thread.start() + except OSError as exc: + if not stopping.is_set(): + errors.append(str(exc)) + + accept_thread = threading.Thread(target=accept, daemon=True) + accept_thread.start() + params["Server"] = f"{loopback},{listener.getsockname()[1]}" + params.setdefault("HostNameInCertificate", host) + params["ConnectRetryCount"] = "0" + proxied = _ConnectionStringBuilder(params).build() + leases = [] + resume = None + try: + pooling(max_size=pool_size, idle_timeout=600) + identities = set() + for _ in range(pool_size): + conn = connect(proxied, autocommit=autocommit, **options) + leases.append(conn) + with conn.cursor() as cur: + cur.execute("SELECT @@SPID") + identities.add(cur.fetchone()[0]) + assert len(identities) == pool_size + for conn in leases: + conn.close() + leases.clear() + + # Keep both TCP peers open but withhold all TDS responses. The + # timer bounds the old implementation's hang without killing it. + gate.clear() + resume = threading.Timer(60 if use_default else 20, gate.set) + resume.start() + start = time.monotonic() + leases.append(connect(proxied, autocommit=autocommit, **options)) + elapsed = time.monotonic() - start + assert len(probed) == 1, f"Checkout probed {len(probed)} unresponsive sockets" + lower, upper = (10, 55) if use_default else (0.5, 12) + assert lower <= elapsed < upper, f"Validation timeout not honored: {elapsed:.2f}s" + assert not gate.is_set(), "Safety timer fired before checkout finished" + for conn in leases: + with conn.cursor() as cur: + cur.execute("SELECT @@SPID, @@TRANCOUNT") + spid, transaction_count = cur.fetchone() + assert spid not in identities, "Unresponsive session was reused" + assert transaction_count == (0 if autocommit else 1) + assert len(sockets) // 2 == pool_size + 1 + assert not errors, errors + finally: + gate.set() + if resume is not None: + resume.cancel() + resume.join() + try: + for conn in leases: + conn.close() + pooling(enabled=False) + finally: + stopping.set() + for sock in sockets: + try: + sock.shutdown(socket.SHUT_RDWR) + except OSError: + pass + sock.close() + accept_thread.join(timeout=12) + listener.close() + for thread in threads: + thread.join(timeout=2) + assert not accept_thread.is_alive() + assert all(not thread.is_alive() for thread in threads) + """.replace("POOL_SIZE", str(pool_size)) + .replace("AUTOCOMMIT", repr(autocommit)) + .replace("TIMEOUT_OPTIONS", repr(options)), + conn_str, + ) + + +def test_pool_checkout_restores_connection_timeout(conn_str): + """Temporary checkout timeout must not become the next lease's query timeout.""" + _run_in_subprocess( + """ + import os + import mssql_python as ms + + ms.pooling(max_size=1) + conn_str = os.environ["DB_CONNECTION_STRING"] + first = ms.connect(conn_str, autocommit=True, timeout=1) + with first.cursor() as cur: + cur.execute("SELECT @@SPID") + identity = cur.fetchone()[0] + first.close() + with ms.connect(conn_str, autocommit=True, timeout=1) as reused: + assert reused.timeout == 0 + with reused.cursor() as cur: + cur.execute("WAITFOR DELAY '00:00:02'; SELECT @@SPID") + assert cur.fetchone()[0] == identity + ms.pooling(enabled=False) + """, + conn_str, + ) + + def test_pool_removes_invalid_connections(conn_str): """Pool must replace a pooled connection whose server-side session has died. From bd0d540169267dc23ef0eaa6a5cf82eaa640839d Mon Sep 17 00:00:00 2001 From: Jahnvi Thakkar Date: Thu, 22 Oct 2026 12:24:00 +0530 Subject: [PATCH 2/2] TEST: Assert native connection timeout restoration on pool reuse Read SQL_ATTR_CONNECTION_TIMEOUT through a private native accessor and check zero and nonzero values on the same reused SQL Server session. Mutation tests reject omitted restoration and unconditional zero. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- mssql_python/pybind/connection/connection.cpp | 15 +++++++++++++++ mssql_python/pybind/connection/connection.h | 1 + mssql_python/pybind/ddbc_bindings.cpp | 2 ++ tests/test_009_pooling.py | 11 ++++++++--- 4 files changed, 26 insertions(+), 3 deletions(-) diff --git a/mssql_python/pybind/connection/connection.cpp b/mssql_python/pybind/connection/connection.cpp index bb56ef284..39e217bbf 100644 --- a/mssql_python/pybind/connection/connection.cpp +++ b/mssql_python/pybind/connection/connection.cpp @@ -934,6 +934,21 @@ py::object Connection::getInfo(SQLUSMALLINT infoType) const { return result; } +SQLULEN ConnectionHandle::getConnectionTimeout() const { + if (!_conn) { + ThrowStdException("Connection object is not initialized"); + } + SQLULEN timeout = 0; + const auto& handle = _conn->getDbcHandle(); + SQLRETURN ret = SQLGetConnectAttr_ptr(handle->get(), SQL_ATTR_CONNECTION_TIMEOUT, + &timeout, 0, nullptr); + if (!SQL_SUCCEEDED(ret)) { + ErrorInfo error = SQLCheckError_Wrap(SQL_HANDLE_DBC, handle, ret); + ThrowStdException("Failed to read connection timeout: " + error.ddbcErrorMsg); + } + return timeout; +} + py::object ConnectionHandle::getInfo(SQLUSMALLINT infoType) const { PERF_TIMER("ConnectionHandle::getInfo"); if (!_conn) { diff --git a/mssql_python/pybind/connection/connection.h b/mssql_python/pybind/connection/connection.h index 7eaa4dcce..0ae225688 100644 --- a/mssql_python/pybind/connection/connection.h +++ b/mssql_python/pybind/connection/connection.h @@ -158,6 +158,7 @@ class ConnectionHandle { void setAutocommit(bool enabled); bool getAutocommit() const; SqlHandlePtr allocStatementHandle(); + SQLULEN getConnectionTimeout() const; void setAttr(int attribute, py::object value); // Get information about the driver and data source diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index 4e4523955..d0ddd0010 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -6476,6 +6476,8 @@ PYBIND11_MODULE(ddbc_bindings, m) { .def("rollback", &ConnectionHandle::rollback, "Rollback the current transaction") .def("set_autocommit", &ConnectionHandle::setAutocommit) .def("get_autocommit", &ConnectionHandle::getAutocommit) + .def("_get_connection_timeout", &ConnectionHandle::getConnectionTimeout, + "Internal: read the native SQL_ATTR_CONNECTION_TIMEOUT value") .def("set_attr", &ConnectionHandle::setAttr, py::arg("attribute"), py::arg("value"), "Set connection attribute") .def("alloc_statement_handle", &ConnectionHandle::allocStatementHandle) diff --git a/tests/test_009_pooling.py b/tests/test_009_pooling.py index e14526168..b55b00e06 100644 --- a/tests/test_009_pooling.py +++ b/tests/test_009_pooling.py @@ -1027,8 +1027,9 @@ def accept(): ) -def test_pool_checkout_restores_connection_timeout(conn_str): - """Temporary checkout timeout must not become the next lease's query timeout.""" +@pytest.mark.parametrize("original_timeout", [0, 7]) +def test_pool_checkout_restores_connection_timeout(conn_str, original_timeout): + """Checkout must restore the native timeout, not just the Python query field.""" _run_in_subprocess( """ import os @@ -1037,17 +1038,21 @@ def test_pool_checkout_restores_connection_timeout(conn_str): ms.pooling(max_size=1) conn_str = os.environ["DB_CONNECTION_STRING"] first = ms.connect(conn_str, autocommit=True, timeout=1) + original_timeout = ORIGINAL_TIMEOUT + first.set_attr(ms.SQL_ATTR_CONNECTION_TIMEOUT, original_timeout) + assert first._conn._get_connection_timeout() == original_timeout with first.cursor() as cur: cur.execute("SELECT @@SPID") identity = cur.fetchone()[0] first.close() with ms.connect(conn_str, autocommit=True, timeout=1) as reused: + assert reused._conn._get_connection_timeout() == original_timeout assert reused.timeout == 0 with reused.cursor() as cur: cur.execute("WAITFOR DELAY '00:00:02'; SELECT @@SPID") assert cur.fetchone()[0] == identity ms.pooling(enabled=False) - """, + """.replace("ORIGINAL_TIMEOUT", str(original_timeout)), conn_str, )