FEAT: token-expiry capture and identity-aware pool key#660
Open
jahnvi480 wants to merge 17 commits into
Open
Conversation
) Foundation for identity-aware connection pooling. No behavior change; all additions are dormant until the connection/native layers consume them. - Capture token expiry: _acquire_token returns (struct, raw, expires_on); add TokenInfo, AADAuth.get_token_info, module-level get_auth_token_info. Public get_token/get_raw_token/get_token_struct signatures unchanged. - compute_identity_key(): per-identity pool-key discriminator. None for non-token auth (bare-connStr key, backward compatible); msi:<client_id>/msi:system derived without a token (enables #659 skip-on-hit); tok:<sha256> fail-safe hash otherwise. - is_token_near_expiry(): 5-min single-threshold check; None expiry treated as near (fail safe). - Tests: TestTokenExpiryCapture, TestComputeIdentityKey, TestIsTokenNearExpiry (110 auth tests pass).
📊 Code Coverage Report
Diff CoverageDiff: main...HEAD, staged and unstaged changes
Summary
mssql_python/pybind/connection/connection.cppLines 303-313 303 // cast in the string branch below would mangle that struct; worse, the
304 // Python identity-aware pool-key logic only hashes bytes/bytearray tokens,
305 // so a str token slips through with the bare connection-string pool key and
306 // two callers passing different str tokens against the same server could
! 307 // share a pooled, authenticated connection. Reject any non-binary token at
! 308 // this native boundary so the cross-identity invariant ("a token is present
! 309 // => the pool key is never the bare connStr") holds regardless of how the
310 // Connection was constructed.
311 if (attribute == SQL_COPT_SS_ACCESS_TOKEN && !py::isinstance<py::bytes>(value) &&
312 !py::isinstance<py::bytearray>(value)) {
313 LOG("Rejecting non-binary SQL_COPT_SS_ACCESS_TOKEN (attribute=%d): access token "Lines 595-604 595 long long expiry = 0;
596 py::dict connect_attrs = Connection::invokeTokenFactory(tokenFactory, expiry);
597 _conn->connect(connect_attrs);
598 _conn->setTokenExpiry(expiry);
! 599 } else {
! 600 _conn->connect(attrsBefore);
601 }
602 }
603 }mssql_python/pybind/connection/connection_pool.cppLines 20-37 20 // beyond now + thresholdSecs. A zero/negative (unknown/missing) expiry is
21 // treated as NOT safe so the caller fails closed rather than reusing a token it
22 // cannot prove is still valid. Caller need not hold any lock; reads the wall
23 // clock only.
! 24 static bool tokenExpirySafelyBeyond(long long expiryEpoch, int thresholdSecs) {
! 25 if (expiryEpoch <= 0) {
! 26 return false;
! 27 }
! 28 const long long now = static_cast<long long>(
! 29 std::chrono::duration_cast<std::chrono::seconds>(
! 30 std::chrono::system_clock::now().time_since_epoch())
! 31 .count());
! 32 return (now + static_cast<long long>(thresholdSecs)) < expiryEpoch;
! 33 }
34
35 // Pull the raw access-token bytes out of a connect-attrs dict returned by the
36 // token factory, or an empty string if the dict carries no token. Caller must
37 // hold the GIL. Keys in the dict are attribute ids (ints). SQL_COPT_SS_ACCESS_TOKENLines 42-51 42 item.first.cast<long>() == SQL_COPT_SS_ACCESS_TOKEN) {
43 try {
44 return item.second.cast<std::string>();
45 } catch (const py::cast_error&) {
! 46 return std::string();
! 47 }
48 }
49 }
50 return std::string();
51 }Lines 123-131 123 _pool.pop_front();
124 }
125
126 // Validate the candidate outside the mutex.
! 127 bool reuse_candidate = false;
128 try {
129 if (token_factory && !token_factory.is_none() &&
130 candidate->isTokenNearExpiry(TOKEN_EXPIRY_THRESHOLD_SECS)) {
131 // Expiry-aware checkout with token compare: the pooled token isLines 161-170 161 // proactively (returning a token with a fresh, far-out
162 // expiry) before the threshold, so the safe-reuse path above
163 // is taken in practice. We favor never handing out a token
164 // that may expire mid-query over avoiding the churn.
! 165 candidate->setTokenExpiry(fresh_expiry);
! 166 reuse_candidate = candidate->isAlive() && candidate->reset();
167 } else {
168 // Token rotated, or the "fresh" token is still at/near
169 // expiry (or has an unknown expiry): discard and reopen with
170 // the fresh attrs. Remember the fresh token to reopen with,Lines 179-200 179 pending_expiry = fresh_expiry;
180 have_pending_token = true;
181 const std::string stale_token = candidate->currentAccessToken();
182 if (!stale_token.empty()) {
! 183 std::lock_guard<std::mutex> lock(_mutex);
! 184 _pool.erase(
! 185 std::remove_if(
! 186 _pool.begin(), _pool.end(),
! 187 [&](const std::shared_ptr<Connection>& sibling) {
! 188 if (sibling->currentAccessToken() == stale_token) {
! 189 to_disconnect.push_back(sibling);
! 190 if (_current_size > 0) --_current_size;
! 191 return true;
! 192 }
! 193 return false;
! 194 }),
! 195 _pool.end());
! 196 }
197 }
198 } else {
199 reuse_candidate = candidate->isAlive() && candidate->reset();
200 }Lines 314-323 314 return false;
315 }
316 // Nothing checked out and the pool is empty: safe to drop immediately.
317 if (_pool.empty()) {
! 318 return true;
! 319 }
320 // Nothing checked out but idle connections remain. Evict the whole pool
321 // only once EVERY pooled connection has been idle longer than the idle
322 // timeout. This is what reclaims pools for rotating / single-use identities
323 // (e.g. per-request Entra users keyed by token hash): such a pool is neverLines 328-337 328 for (const auto& conn : _pool) {
329 auto idle_time =
330 std::chrono::duration_cast<std::chrono::seconds>(now - conn->lastUsed()).count();
331 if (idle_time <= _idle_timeout_secs) {
! 332 return false;
! 333 }
334 }
335 return true;
336 }Lines 429-438 429 for (auto& evicted_pool : evicted) {
430 try {
431 evicted_pool->close();
432 } catch (const std::exception& ex) {
! 433 LOG("ConnectionPoolManager: closing evicted pool failed: %s", ex.what());
! 434 }
435 }
436 // Call acquire() outside _manager_mutex. acquire() may release the GIL
437 // during the ODBC connect call; holding _manager_mutex across that would
438 // create a mutex/GIL lock-ordering deadlock. connStr (not key) is used toLines 463-472 463 if (conn) {
464 try {
465 conn->disconnect();
466 } catch (const std::exception& ex) {
! 467 LOG("ConnectionPoolManager::returnConnection: disconnect of orphaned "
! 468 "connection failed: %s",
469 ex.what());
470 }
471 }
472 }Lines 504-512 504 // Close each pool outside _manager_mutex.
505 for (auto& pool : to_close) {
506 try {
507 pool->close();
! 508 } catch (const std::exception& ex) {
509 LOG("ConnectionPoolManager::closePools: closing pool failed: %s", ex.what());
510 }
511 }
512 }📋 Files Needing Attention📉 Files with overall lowest coverage (click to expand)mssql_python.pybind.logger_bridge.cpp: 59.2%
mssql_python.pybind.ddbc_bindings.h: 59.9%
mssql_python.pybind.logger_bridge.hpp: 70.8%
mssql_python.pybind.ddbc_bindings.cpp: 76.4%
mssql_python.__init__.py: 77.3%
mssql_python.row.py: 77.6%
mssql_python.ddbc_bindings.py: 79.6%
mssql_python.pybind.connection.connection_pool.cpp: 82.9%
mssql_python.pybind.connection.connection.cpp: 83.6%
mssql_python.logging.py: 85.5%🔗 Quick Links
|
Pooled Entra connections were keyed only on the sanitized connection string, so two callers hitting the same server with different Entra identities could reuse each other's authenticated connection (cross-identity reuse, #651). connection.py now computes an identity-aware pool key (connStr + NUL + security-context) whenever an access token is present: system/user-assigned MSI keyed on client_id, other token auth keyed on a SHA-256 of the token struct. Non-token auth (SQL, trusted, native ServicePrincipal, Windows Interactive) keeps the legacy bare-connStr key, so behavior is unchanged. Native pool threads an optional pool_key through ConnectionPoolManager::acquireConnection/returnConnection and the pybind Connection init (defaulted, backward compatible); the physical connect still uses the connection string. Adds 6 TestConnectionPoolKey unit tests.
…#659) For MSI auth the identity (client id) is known without a token, so the pool key can be computed up front. Acquisition of the access token is deferred to an internal token_factory callback that the native layer invokes only when it opens a physical connection. A same-identity pool hit therefore reuses the pooled connection and never pays for a token. Token-dependent identities (DefaultAzureCredential / raw token / interactive / device-code) still acquire eagerly because the pool key is the token hash. The internal token_factory is intentionally distinct from the public token_provider= credential API proposed in PR #603 (issue #577).
github-actions
Bot
added
pr-size: large
Add unit tests for auth.get_auth_token_info Windows/interactive and error paths, and is_token_near_expiry wall-clock default. Add native integration tests driving ddbc_bindings.Connection with a token_factory to exercise the C++ lazy-token branches (pool miss vs reuse, non-pooled with/without factory).
… isolation test - Correct SP claims: ServicePrincipal is ODBC-native (creds retained in connStr), so it is already isolated and needs no dedicated key; fix design doc s3.1/s4/s12. - Add an 'Implementation status (PR #660)' section clarifying wired vs deferred (home_account_id, expiry-aware checkout, returnConnection-on-missing, eviction). - Clarify get_token_info docstring: expiry capture is foundation, not yet consumed on checkout. - Add native integration test proving distinct pool keys do not share connections (and that the \x00 key separator survives str->u16string).
…re pooling Remove unreachable acquire-timeout/condition_variable path from the native connection pool, drop the unused is_token_near_expiry helper and orphaned get_auth_token, and delete their now-redundant tests. Clean up code comments.
…pools auth: try get_token() silently first for interactive/device-code credentials and only call authenticate() when the SDK raises AuthenticationRequiredError; cache the resolved home_account_id so subsequent acquisitions stay silent. pool: compare a freshly minted token against a near-expiry connection's current token so a matching token reuses the connection; make canEvict() evaluate the idle timeout so pools holding only idle connections are reclaimed, and close evicted pools outside the manager mutex. tests: add silent-first regression/branch tests (auth) and a real-DB idle identity-pool eviction test (pooling).
pool: throttle the lazy-eviction sweep to at most once per idle-timeout window instead of running on every acquireConnection, cutting global-lock contention under many-identity workloads; eagerly drain sibling idle connections holding a rotated (stale) token in one pass; consolidate the SQL_COPT_SS_ACCESS_TOKEN attribute id into connection.h as a single source of truth. docs: document in README that bring-your-own token lifetime is the application's responsibility, that DefaultAzureCredential is not recommended for multi-user pooling, and that interactive/device-code auth shares one signed-in account per process.
Clear _last_sweep in configure() and closePools() so the first acquireConnection() after a pooling reconfigure or shutdown performs a fresh sweep instead of honoring a stale throttle interval.
Contributor
There was a problem hiding this comment.
Pull request overview
This PR hardens and improves connection pooling for Entra ID (Azure AD) authentication by making the native pool key identity-aware (preventing cross-identity reuse of authenticated connections), deferring token acquisition until a physical connection is actually opened, and adding token-expiry awareness plus lazy eviction of idle identity pools.
Changes:
- Add
TokenInfo(token struct + expiry + optional account id),get_auth_token_info(), andcompute_identity_key()to support identity-aware pooling and expiry-aware behavior. - Extend the Python
Connectionlayer to compute an identity-awarepool_keyand optionally pass a lazytoken_factorycallback down to the native layer. - Update native pooling to key by
pool_key, invoketoken_factoryonly on pool misses, refresh/reopen near-expiry token connections, and lazily evict idle identity pools; add integration/diagnostic tests and document limitations.
Reviewed changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/test_009_pooling.py | Adds integration/regression tests for idle identity-pool eviction and native lazy token-factory behavior. |
| tests/test_008_auth.py | Adds unit coverage for silent-first interactive/device-code flow, token-expiry capture, identity key computation, and pool-key wiring. |
| README.md | Documents multi-user caveats for interactive/device-code, BYO token lifetime responsibility, and DefaultAzureCredential pooling guidance. |
| mssql_python/pybind/ddbc_bindings.cpp | Extends the pybind Connection constructor to accept pool_key and token_factory, and exposes _pool_count() for tests. |
| mssql_python/pybind/connection/connection.h | Introduces token expiry tracking APIs and a shared SQL_COPT_SS_ACCESS_TOKEN constant used by pooling/connection code. |
| mssql_python/pybind/connection/connection.cpp | Implements token-factory invocation, token-expiry recording, and routes pooling acquire/return via identity-aware keys. |
| mssql_python/pybind/connection/connection_pool.h | Extends pool acquire/manager APIs to take pool_key and token_factory; adds canEvict()/poolCount(). |
| mssql_python/pybind/connection/connection_pool.cpp | Implements lazy token acquisition, expiry-aware checkout (token compare + reopen), and throttled idle identity-pool eviction. |
| mssql_python/connection.py | Builds identity-aware pool keys and defers token acquisition via _token_factory where possible; forwards both to native. |
| mssql_python/auth.py | Adds TokenInfo, silent-first interactive/device-code acquisition with cached home_account_id, DefaultAzureCredential warning, and identity key computation. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
…tion - auth.get_auth_token_info: propagate acquisition errors instead of swallowing them and returning None (avoids opening a connection with Authentication= stripped and no token) - connection: token factory and eager auth branch raise InterfaceError when no usable token can be produced - pool acquire: only reuse a same-token near-expiry candidate when the refreshed expiry is safely beyond the threshold - pool eviction sweep: only evict pools with use_count == 1 so an in-flight acquirer's pool is never closed from under it - pooling.disable now calls disable_pooling() to disarm the native manager - add tests and coverage tooling fix
enable() calls the native setAccepting(true) but did not reset the Python _pools_closed flag, so an enable -> disable -> enable -> disable sequence skipped the second disable_pooling() and desynced from the native manager. Reset the flag in enable() and add a regression test.
…hurn edge - connection: wrap get_auth_token_info ValueError/RuntimeError from the eager auth branch and the token factory as DB-API InterfaceError so every auth failure from connect() is a consistent, catchable driver exception - update credential-failure test to expect InterfaceError - connection_pool: document the narrow same-token near-expiry churn edge (a misbehaving provider returning the same near-expiry token repeatedly causes discard+reopen churn; well-behaved SDKs refresh proactively)
- setAttribute (native): reject SQL_COPT_SS_ACCESS_TOKEN (1256) that is not bytes/bytearray with SQL_ERROR, so a str/other token can never be applied - connection (Python): reject a non-bytes attr 1256 with InterfaceError before any physical connect, so a str token can never fall through with the bare connStr pool key and be shared across identities; enforced at both boundaries - tests: non-binary token (str/int/memoryview/list) fails closed before connect; factory token with unknown/zero expiry is reused (documents epoch-0 semantics); factory raising mid near-expiry checkout recovers the reserved pool slot
Contributor
Author
Reviewer's guide: identity-aware connection poolingThe 2-minute backgroundThe bug: the label was built only from the connection string. That's fine for username/password logins (the password is in the string). But with Entra ID (Azure AD), you authenticate with a short-lived access token, not a password — so two different people hitting the same server produce the same connection string, get the same label, and land in the same pool. User B could be handed a connection authenticated as User A. That identity leak is what this PR fixes. The fix in one line: when a token is involved, add the caller's identity to the pool key. Terms you'll see:
The one rule ("invariant") to keep in mind while reviewing:
The big picturePython figures out who is connecting and builds the pool key; the native C++ layer owns the actual pool. Note for C++ reviewers: the pool has two locks (manager, then pool), and any network call (connect / disconnect / validate) must run with no lock held — that's the main thing to watch for in Passes 4–6. flowchart TB
subgraph PY["Python layer"]
C["connection.py<br/>builds pool key + factory"]
A["auth.py<br/>picks identity, gets tokens"]
P["pooling.py<br/>PoolingManager on/off"]
end
subgraph NATIVE["Native C++ — pool lives here"]
M["Pool manager<br/>key to pool + switch"]
Pool["One pool<br/>shelf of connections"]
Conn["Connection<br/>ODBC handle + expiry"]
end
C --> A
C --> M
P --> M
M --> Pool --> Conn
Pass 1 — How the identity key is chosen🔗 What it does: given the auth type and its parameters, returns a short string for "who is connecting" — or Why it matters: this function is the source of truth for the whole fix. Wrong answer here = identities can collide. What to check:
flowchart TD
S["connect()"] --> AT{"auth type?"}
AT -->|"password-style"| BARE["bare connStr key<br/>unchanged"]
AT -->|"Managed Identity"| MSI["msi:client_id"]
AT -->|"Interactive"| ACCT["acct:account_id"]
AT -->|"DAC / raw token"| HASH["tok:sha256"]
BARE --> CMP["pool key = connStr + identity"]
MSI --> CMP
ACCT --> CMP
HASH --> CMP
Pass 2 — Where the pool key gets attached🔗 token-info helper · token-factory helper · key-building block · raw-token safety net + fail-closed guard · same guard in native · test: What it does: when a Why it matters: this is the enforcement point for the invariant — every token path must set a key; nothing token-related may fall through to the bare-string path. What to check:
flowchart TD
S["Connection.__init__"] --> K["compute_identity_key"]
K --> Q{"identity without a token?<br/>Managed Identity"}
Q -->|yes| F1["set key<br/>lazy factory"]
Q -->|no| E["get token now<br/>fail: InterfaceError"]
E --> Q2{"identity kind?"}
Q2 -->|"account id"| F2["set key<br/>lazy factory"]
Q2 -->|"raw token"| T["key from hash<br/>pass token through"]
Q2 -->|"none"| FC["InterfaceError"]
F1 --> RAW{"raw token in attrs,<br/>no key yet?"}
F2 --> RAW
T --> RAW
RAW -->|"bytes"| RN["key = connStr + tok:hash"]
RAW -->|"str / other"| REJ["InterfaceError<br/>fail closed"]
RAW -->|"none"| NAT["hand to native Connection"]
RN --> NAT
Pass 3 — Interactive tokens without nagging the user🔗 What it does: for interactive/device-code logins, try the cached session first; only show a login prompt if that genuinely fails. Why it matters: pooling reconnects often — without this, every reconnect could pop a login prompt. We also cache the account ID so later reconnects stay silent. What to check:
sequenceDiagram
participant App as connect()
participant Auth as acquire_token
participant Cred as Credential
participant Cache as account cache
App->>Auth: need interactive token
Auth->>Cache: known account id?
Auth->>Cred: get_token silently
alt cached session works
Cred-->>Auth: token, no prompt
else login required
Auth->>Cred: authenticate, prompt once
Cred-->>Auth: account id
Auth->>Cache: remember it
Auth->>Cred: get_token
Cred-->>Auth: token
end
Auth-->>App: token + expiry + account id
Pass 4 — Native pool: keying + minting tokens only on a miss🔗 What it does: the C++ pool files connections under the pool key (not the raw connection string), and calls the token factory only when it must open a new physical connection (a "pool miss"). Why it matters: (1) keying by pool key is what physically separates identities. (2) "mint on miss only" means a pool hit reuses a connection with zero token work — the performance win. What to check (C++ reviewers):
flowchart TD
A["acquire()"] --> P1["prune dead/idle"]
P1 --> LOOP{"connection on shelf?"}
LOOP -->|"yes"| POP["validate outside lock"]
LOOP -->|"no, room"| RES["reserve slot"]
LOOP -->|"no, full"| FULL["pool full error"]
POP --> NE{"factory set +<br/>near expiry?"}
NE -->|"no"| V{"alive + resettable?"}
NE -->|"yes"| EXP["expiry refresh<br/>Pass 5"]
V -->|"yes"| REUSE["reuse<br/>no token minted"]
V -->|"no"| DISC["discard"]
EXP --> DISC
DISC --> LOOP
RES --> OPEN["mint token + connect<br/>outside lock"]
REUSE --> DONE["hand to caller"]
OPEN --> DONE
Pass 5 — Don't hand out a connection whose token is about to expire🔗 near-expiry refresh block · What it does: if a pooled connection's token is within ~5 minutes of expiring, mint a fresh token and compare. Same token → reuse the connection. Different (rotated) token → discard the stale connections and open fresh ones with the new token. Why it matters: without this you could hand out a connection that dies mid-query because its token just expired. What to check:
flowchart TD
NE["token near expiry"] --> MINT["mint fresh token"]
MINT --> CMP{"same token +<br/>far from expiry?"}
CMP -->|"yes"| SAFE["update expiry<br/>reuse"]
CMP -->|"no, rotated"| ROT["drain stale siblings<br/>discard"]
ROT --> REOPEN["reopen with new token"]
Pass 6 — Cleaning up unused pools, and the on/off switch🔗 What it does: two housekeeping concerns. (1) Eviction: every distinct identity now creates its own pool, so a one-off identity would leave an idle pool forever — on later calls we sweep out pools idle too long. (2) Lifecycle: Why it matters: without eviction, per-identity pools slowly leak memory/handles. And disable→re-enable must actually work (there was a latent bug that kept it "closed"). What to check:
stateDiagram-v2
[*] --> Off
Off --> On: enable()
On --> Off: disable()
Off --> On: enable() again
On --> On: connect, pooled
Off --> Off: connect, non-pooled
Two practical notes for reviewers
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Work Item / Issue Reference
Summary
FEAT: Identity-aware connection pooling
The problem
The connection pool keyed only on the connection string. When two callers hit the same server with different Entra (Azure AD) identities, they collided onto one pool — so a connection authenticated as user A could be handed to user B. Two secondary issues: a token was acquired on every connect (even pool hits), and a pooled connection could be reused with a near-expired token.
What this PR does
Keys the pool by security context, not just the connection string, upholding one invariant: a token-backed connection is never keyed on the bare connection string. It also defers token acquisition to pool-misses and refreshes pooled connections whose token is near expiry.
Design decisions (aligned with the cross-driver discussion)
connStr + "\x00" + <identity>.compute_identity_key): MSI →msi:<client_id>/msi:system(derived from params, no token needed); Interactive/DeviceCode →acct:<home_account_id>; DefaultAzureCredential / raw token →tok:<sha256(token)>.Implementation highlights
\x00separator means the composite key can never collide with a bare connStr key.disable_automatic_authentication=True; interactive prompt only onAuthenticationRequiredError;home_account_idcached so reconnects stay silent.use_count()==1guard, closed outside the manager mutex to respect lock ordering).InterfaceErrorrather than connecting token-less; a non-bytesaccess-token attribute is rejected at both the Python and native boundaries so it can't slip through with a bare pool key.Behavior matrix
msi:<client_id>/msi:systemacct:<home_account_id>tok:<sha256>Testing
tests/test_008_auth.py): identity-key mapping, silent-first interactive auth, fail-closed acquisition, pool-key wiring, and adversarial non-binary-token cases.tests/test_009_pooling.py, require a live SQL Server): cross-identity isolation, lazy factory invoked-on-miss / skipped-on-reuse, near-expiry refresh & token-bytes extraction, idle-pool eviction, checked-out-pool-not-evicted, enable/disable re-arm, concurrent-disable, and slot recovery on factory failure.Documented limitations (README)
DefaultAzureCredentialis not recommended for multi-user pooling (one-time warning emitted; keyed on token hash, so not expiry-aware).Out of scope / follow-ups