Skip to content

Commit 21187de

Browse files
jahnvi480Copilot
andauthored
FIX: Bound pooled connection checkout on unresponsive sockets (#859)
### Work Item / Issue Reference <!-- IMPORTANT: Please follow the PR template guidelines below. For mssql-python maintainers: Insert your ADO Work Item ID below For external contributors: Insert Github Issue number below Only one reference is required - either GitHub issue OR ADO Work Item. --> <!-- mssql-python maintainers: ADO Work Item --> <!-- External contributors: GitHub Issue --> > GitHub Issue: #854 ------------------------------------------------------------------- ### Summary <!-- Insert your summary of changes below. Minimum 10 characters required. --> Bound pooled connection reset and autocommit configuration using a temporary native connection timeout. Use the positive login timeout (explicit attrs_before takes precedence), or a 15-second validation default, and restore the previous connection timeout on success. After validation fails, open a replacement rather than probing every stale idle socket sequentially. Preserve isolation/transaction cleanup and document that cancellation, cleanup, and replacement login can add time: this is not an end-to-end connect deadline. Add stalled-socket regressions covering both autocommit modes, one/three stale sockets, timeout precedence, the finite default, and timeout restoration. Extend isolation reset coverage to SQL SET statements and adjust close-error mocks to inject failures after construction. Validation (Linux x64, Python 3.13, SQL Server 2022): - `bash build.sh` and subsequent `cmake --build /work/mssql_python/pybind/build --config Release -j2`: passed; rebuilt native extension used for testing. - `python -m pytest -v --tb=short`: 5,537 passed, 167 skipped, 42 deselected. Linux test copy used LF shell scripts to avoid Windows checkout CRLF failures. - `python -m pre_commit run black-check --all-files --hook-stage pre-push`: passed. - Actual TCP DROP reproduction: approximately 3 seconds with timeout=1 for one/three stale connections in both autocommit modes; replacement queries succeeded. - Windows/macOS validation not run. An initial combined packet-drop run exited with code 139; isolated and repeated runs subsequently passed, but the initial crash was not conclusively explained. <!-- ### PR Title Guide > For feature requests FEAT: (short-description) > For non-feature requests like test case updates, config updates , dependency updates etc CHORE: (short-description) > For Fix requests FIX: (short-description) > For doc update requests DOC: (short-description) > For Formatting, indentation, or styling update STYLE: (short-description) > For Refactor, without any feature changes REFACTOR: (short-description) > For performance improvements PERF: (short-description) > For AI tooling, agents, skills, prompts, or AI-assisted development workflows AI: (short-description) > For release related changes RELEASE: #<RELEASE_VERSION> (short-description) ### Contribution Guidelines External contributors: - Create a GitHub issue first: https://github.com/microsoft/mssql-python/issues/new - Link the GitHub issue in the "GitHub Issue" section above - Follow the PR title format and provide a meaningful summary mssql-python maintainers: - Create an ADO Work Item following internal processes - Link the ADO Work Item in the "ADO Work Item" section above - Follow the PR title format and provide a meaningful summary --> Fixes: #854 --------- Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
1 parent 185bc65 commit 21187de

8 files changed

Lines changed: 546 additions & 29 deletions

File tree

‎README.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,10 @@ EntraID authentication is now fully supported on MacOS and Linux but with certai
129129

130130
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).
131131

132+
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.
133+
134+
This favors bounded validation latency over trying every idle connection during a single checkout. If the replacement login fails, that checkout raises a connection exception even if a healthy idle connection remains; the healthy connection stays pooled and can be reused by a subsequent checkout.
135+
132136
> **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.
133137
134138
> **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.

‎mssql_python/connection.py‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -417,7 +417,11 @@ def __init__(
417417
query timeout. An explicit ``attrs_before[SQL_ATTR_LOGIN_TIMEOUT]``
418418
takes precedence over this value. (pyodbc migration note: pyodbc
419419
lets a positive ``timeout=`` override ``attrs_before``; here the
420-
explicit ``attrs_before`` entry wins.)
420+
explicit ``attrs_before`` entry wins.) Pooled candidate reset
421+
also uses this positive timeout, or 15 seconds when zero/default,
422+
temporarily as SQL_ATTR_CONNECTION_TIMEOUT. Driver cancellation
423+
and a replacement login can add time; this is not an overall
424+
connect deadline.
421425
native_uuid (bool, optional): Controls whether UNIQUEIDENTIFIER columns return
422426
uuid.UUID objects (True) or str (False) for cursors created from this connection.
423427
None (default) defers to the module-level ``mssql_python.native_uuid`` setting (True).
@@ -858,17 +862,22 @@ def _token_factory():
858862
_provider = ProviderManager.ensure_available()
859863
ddbc_bindings._set_odbc_provider(_provider)
860864

865+
# Configure a reused connection's mode inside native checkout, while
866+
# its temporary connection timeout still bounds attribute-setting I/O.
867+
native_attrs = self._attrs_before.copy()
868+
native_attrs[ConstantsDDBC.SQL_ATTR_AUTOCOMMIT.value] = int(bool(autocommit))
861869
try:
862870
self._conn = ddbc_bindings.Connection(
863871
self.connection_str,
864872
self._pooling,
865-
self._attrs_before,
873+
native_attrs,
866874
self._pool_key,
867875
self._token_factory,
868876
)
877+
if self._conn.get_autocommit() != autocommit:
878+
self.setautocommit(autocommit)
869879
except RuntimeError as e:
870880
_raise_connection_error(e)
871-
self.setautocommit(autocommit)
872881

873882
# Register this connection for cleanup before Python shutdown
874883
# This ensures ODBC handles are freed in correct order, preventing leaks

‎mssql_python/pybind/connection/connection.cpp‎

Lines changed: 50 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -600,11 +600,31 @@ bool Connection::isAlive() const {
600600
return SQL_SUCCEEDED(ret) && status == SQL_CD_FALSE;
601601
}
602602

603-
bool Connection::reset() {
603+
bool Connection::reset(unsigned int timeoutSeconds, bool autocommit) {
604604
if (!_dbcHandle) {
605605
ThrowStdException("Connection handle not allocated");
606606
}
607607
clearResultMetadata();
608+
// Login and statement timeouts do not bound connection-attribute calls.
609+
// Keep a finite connection timeout through reset and lease configuration;
610+
// retain it on failure so discarding the candidate is bounded too (#854).
611+
SQLULEN originalTimeout = 0;
612+
SQLRETURN ret = SQLGetConnectAttr_ptr(_dbcHandle->get(), SQL_ATTR_CONNECTION_TIMEOUT,
613+
&originalTimeout, 0, nullptr);
614+
if (!SQL_SUCCEEDED(ret)) {
615+
LOG("Failed to read pooled connection timeout (ret=%d). Marking as dead.", ret);
616+
return false;
617+
}
618+
{
619+
py::gil_scoped_release release;
620+
ret = SQLSetConnectAttr_ptr(
621+
_dbcHandle->get(), SQL_ATTR_CONNECTION_TIMEOUT,
622+
reinterpret_cast<SQLPOINTER>(static_cast<SQLULEN>(timeoutSeconds)), SQL_IS_UINTEGER);
623+
}
624+
if (!SQL_SUCCEEDED(ret)) {
625+
LOG("Failed to set pooled connection timeout (ret=%d). Marking as dead.", ret);
626+
return false;
627+
}
608628
LOG("Resetting connection via SQL_ATTR_RESET_CONNECTION");
609629
// NOTE: SQL_ATTR_RESET_CONNECTION is a pool-checkin reset: it asks the
610630
// driver to wipe per-session state (temp tables, open cursors, SET
@@ -618,7 +638,6 @@ bool Connection::reset() {
618638
// Connection (_attrStringBuffers / _attrBytesBuffers) are intentionally
619639
// retained here. Clearing them would reintroduce issue #594 in a new
620640
// form (UAF during silent reconnect).
621-
SQLRETURN ret;
622641
{
623642
// Release the GIL around the ODBC call for consistency with the
624643
// other connection-attribute paths; some attributes can block.
@@ -631,17 +650,27 @@ bool Connection::reset() {
631650
return false;
632651
}
633652

634-
// SQL_ATTR_RESET_CONNECTION does NOT reset the transaction isolation level.
635-
// Explicitly reset it to the default (SQL_TXN_READ_COMMITTED) to prevent
636-
// isolation level settings from leaking between pooled connection usages.
653+
// SQL_ATTR_RESET_CONNECTION does not reset transaction isolation.
637654
LOG("Resetting transaction isolation level to READ COMMITTED");
638655
{
639656
py::gil_scoped_release release;
640657
ret = SQLSetConnectAttr_ptr(_dbcHandle->get(), SQL_ATTR_TXN_ISOLATION,
641658
(SQLPOINTER)SQL_TXN_READ_COMMITTED, SQL_IS_INTEGER);
642659
}
643660
if (!SQL_SUCCEEDED(ret)) {
644-
LOG("Failed to reset transaction isolation level (ret=%d). Marking as dead.", ret);
661+
ErrorInfo error = SQLReadError(SQL_HANDLE_DBC, _dbcHandle->get(), ret);
662+
LOG("Pooled connection validation failed: %s", error.ddbcErrorMsg.c_str());
663+
return false;
664+
}
665+
setAutocommit(autocommit);
666+
{
667+
py::gil_scoped_release release;
668+
ret = SQLSetConnectAttr_ptr(
669+
_dbcHandle->get(), SQL_ATTR_CONNECTION_TIMEOUT,
670+
reinterpret_cast<SQLPOINTER>(originalTimeout), SQL_IS_UINTEGER);
671+
}
672+
if (!SQL_SUCCEEDED(ret)) {
673+
LOG("Failed to restore pooled connection timeout (ret=%d). Marking as dead.", ret);
645674
return false;
646675
}
647676

@@ -905,6 +934,21 @@ py::object Connection::getInfo(SQLUSMALLINT infoType) const {
905934
return result;
906935
}
907936

937+
SQLULEN ConnectionHandle::getConnectionTimeout() const {
938+
if (!_conn) {
939+
ThrowStdException("Connection object is not initialized");
940+
}
941+
SQLULEN timeout = 0;
942+
const auto& handle = _conn->getDbcHandle();
943+
SQLRETURN ret = SQLGetConnectAttr_ptr(handle->get(), SQL_ATTR_CONNECTION_TIMEOUT,
944+
&timeout, 0, nullptr);
945+
if (!SQL_SUCCEEDED(ret)) {
946+
ErrorInfo error = SQLCheckError_Wrap(SQL_HANDLE_DBC, handle, ret);
947+
ThrowStdException("Failed to read connection timeout: " + error.ddbcErrorMsg);
948+
}
949+
return timeout;
950+
}
951+
908952
py::object ConnectionHandle::getInfo(SQLUSMALLINT infoType) const {
909953
PERF_TIMER("ConnectionHandle::getInfo");
910954
if (!_conn) {

‎mssql_python/pybind/connection/connection.h‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ class Connection {
5858
// Check whether autocommit is enabled.
5959
bool getAutocommit() const;
6060
bool isAlive() const;
61-
bool reset();
61+
bool reset(unsigned int timeoutSeconds, bool autocommit = true);
6262
void prepareForPool(bool transactionAlreadyRolledBack = false);
6363
void updateLastUsed();
6464
std::chrono::steady_clock::time_point lastUsed() const;
@@ -158,6 +158,7 @@ class ConnectionHandle {
158158
void setAutocommit(bool enabled);
159159
bool getAutocommit() const;
160160
SqlHandlePtr allocStatementHandle();
161+
SQLULEN getConnectionTimeout() const;
161162
void setAttr(int attribute, py::object value);
162163

163164
// Get information about the driver and data source

‎mssql_python/pybind/connection/connection_pool.cpp‎

Lines changed: 36 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
#include <chrono>
77
#include <exception>
88
#include <memory>
9+
#include <limits>
910
#include <vector>
1011

1112
// Logging uses LOG() macro for all diagnostic output
@@ -60,6 +61,22 @@ std::shared_ptr<Connection> ConnectionPool::acquire(const std::u16string& connSt
6061
const py::dict& attrs_before,
6162
const py::object& token_factory) {
6263
PERF_TIMER("ConnectionPool::acquire");
64+
// A login timeout does not cover SQLSetConnectAttr on an existing socket.
65+
// Use it for bounded candidate reset as well; zero/unspecified
66+
// retains a finite validation default rather than an OS TCP retry budget.
67+
unsigned int validation_timeout = 15;
68+
const py::int_ login_timeout_attr(SQL_ATTR_LOGIN_TIMEOUT);
69+
if (attrs_before.contains(login_timeout_attr)) {
70+
const auto requested = attrs_before[login_timeout_attr].cast<long long>();
71+
if (requested > 0) {
72+
validation_timeout = static_cast<unsigned int>(std::min<long long>(
73+
requested, std::numeric_limits<unsigned int>::max()));
74+
}
75+
}
76+
const py::int_ autocommit_attr(SQL_ATTR_AUTOCOMMIT);
77+
const bool autocommit = attrs_before.contains(autocommit_attr)
78+
? attrs_before[autocommit_attr].cast<int>() != SQL_AUTOCOMMIT_OFF
79+
: true;
6380
std::vector<std::shared_ptr<Connection>> to_disconnect;
6481
std::shared_ptr<Connection> valid_conn = nullptr;
6582
bool needs_connect = false;
@@ -131,6 +148,7 @@ std::shared_ptr<Connection> ConnectionPool::acquire(const std::u16string& connSt
131148

132149
// Validate the candidate outside the mutex.
133150
bool reuse_candidate = false;
151+
bool validation_attempted = false;
134152
try {
135153
if (token_factory && !token_factory.is_none() &&
136154
candidate->isTokenNearExpiry(TOKEN_EXPIRY_THRESHOLD_SECS)) {
@@ -169,7 +187,10 @@ std::shared_ptr<Connection> ConnectionPool::acquire(const std::u16string& connSt
169187
// is taken in practice. We favor never handing out a token
170188
// that may expire mid-query over avoiding the churn.
171189
candidate->setTokenExpiry(fresh_expiry);
172-
reuse_candidate = candidate->isAlive() && candidate->reset();
190+
if (candidate->isAlive()) {
191+
validation_attempted = true;
192+
reuse_candidate = candidate->reset(validation_timeout, autocommit);
193+
}
173194
if (!reuse_candidate) {
174195
// The token is still valid but the socket is dead
175196
// (isAlive()/reset() failed). We already minted the
@@ -213,7 +234,10 @@ std::shared_ptr<Connection> ConnectionPool::acquire(const std::u16string& connSt
213234
}
214235
}
215236
} else {
216-
reuse_candidate = candidate->isAlive() && candidate->reset();
237+
if (candidate->isAlive()) {
238+
validation_attempted = true;
239+
reuse_candidate = candidate->reset(validation_timeout, autocommit);
240+
}
217241
}
218242
} catch (const std::exception& ex) {
219243
LOG("Candidate connection validation failed: %s", ex.what());
@@ -232,10 +256,10 @@ std::shared_ptr<Connection> ConnectionPool::acquire(const std::u16string& connSt
232256
if (_current_size > 0) --_current_size;
233257
}
234258

235-
// If a rotated token was captured, reserve a slot and reopen with it
236-
// immediately instead of churning through the remaining candidates
237-
// (which hold the same stale token and would all be discarded anyway).
238-
if (have_pending_token) {
259+
// After a reset failure, open a replacement instead of spending another
260+
// timeout on each idle socket. Token rotation likewise opens directly
261+
// using the captured token.
262+
if (have_pending_token || validation_attempted) {
239263
std::lock_guard<std::mutex> lock(_mutex);
240264
if (_current_size < _max_size) {
241265
// Reserve the slot here but construct the Connection outside
@@ -246,11 +270,12 @@ std::shared_ptr<Connection> ConnectionPool::acquire(const std::u16string& connSt
246270
needs_connect = true;
247271
break;
248272
}
249-
// Pool momentarily full; fall through and retry the loop. On the
250-
// retry another near-expiry candidate may re-invoke the factory and
251-
// overwrite pending_attrs/pending_expiry with a newer token. That
252-
// needs a full pool AND a simultaneous rotation, is rare, and is
253-
// harmless: we simply reopen with the most recently minted token.
273+
if (validation_attempted) {
274+
throw std::runtime_error(
275+
"ConnectionPool::acquire: pool size limit reached");
276+
}
277+
// A concurrent acquire reserved the slot freed by token rotation.
278+
// Retry token comparison with another idle candidate.
254279
}
255280
}
256281

‎mssql_python/pybind/ddbc_bindings.cpp‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6476,6 +6476,8 @@ PYBIND11_MODULE(ddbc_bindings, m) {
64766476
.def("rollback", &ConnectionHandle::rollback, "Rollback the current transaction")
64776477
.def("set_autocommit", &ConnectionHandle::setAutocommit)
64786478
.def("get_autocommit", &ConnectionHandle::getAutocommit)
6479+
.def("_get_connection_timeout", &ConnectionHandle::getConnectionTimeout,
6480+
"Internal: read the native SQL_ATTR_CONNECTION_TIMEOUT value")
64796481
.def("set_attr", &ConnectionHandle::setAttr, py::arg("attribute"), py::arg("value"),
64806482
"Set connection attribute")
64816483
.def("alloc_statement_handle", &ConnectionHandle::allocStatementHandle)

‎tests/test_006_exceptions.py‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -301,11 +301,12 @@ def test_close_cleans_up_after_autocommit_read_failure():
301301
from unittest.mock import MagicMock, patch
302302

303303
mock_conn = MagicMock()
304-
mock_conn.get_autocommit.side_effect = RuntimeError("SQLSTATE:08S01:Communication link failure")
304+
mock_conn.get_autocommit.return_value = False
305305

306306
with patch("mssql_python.connection.ddbc_bindings.Connection", return_value=mock_conn):
307307
conn = connect("Server=testserver;Database=mydb;Trusted_Connection=yes;")
308308

309+
mock_conn.get_autocommit.side_effect = RuntimeError("SQLSTATE:08S01:Communication link failure")
309310
with pytest.raises(OperationalError, match="Communication link failure"):
310311
conn.close()
311312

@@ -356,16 +357,16 @@ def test_native_close_error_takes_precedence_over_preclose_failure(preclose_fail
356357
from unittest.mock import MagicMock, patch
357358

358359
mock_conn = MagicMock()
360+
mock_conn.get_autocommit.return_value = False
361+
362+
with patch("mssql_python.connection.ddbc_bindings.Connection", return_value=mock_conn):
363+
conn = connect("Server=testserver;Database=mydb;Trusted_Connection=yes;")
364+
359365
if preclose_failure == "autocommit":
360366
mock_conn.get_autocommit.side_effect = RuntimeError("SQLSTATE:08S01:Autocommit read failed")
361367
else:
362-
mock_conn.get_autocommit.return_value = False
363368
mock_conn.rollback.side_effect = RuntimeError("SQLSTATE:08S01:Rollback failed")
364369
mock_conn.close.side_effect = RuntimeError("SQLSTATE:08003:Native close failed")
365-
366-
with patch("mssql_python.connection.ddbc_bindings.Connection", return_value=mock_conn):
367-
conn = connect("Server=testserver;Database=mydb;Trusted_Connection=yes;")
368-
369370
with pytest.raises(OperationalError, match="Native close failed"):
370371
conn.close()
371372

0 commit comments

Comments
 (0)