FEAT: token-expiry capture and identity-aware pool key

README.md

@@ -74,6 +74,12 @@ 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).
 
+> **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 only applies to tokens it acquired itself.
+
+> **`DefaultAzureCredential` is not recommended for multi-user pooling:** `ActiveDirectoryDefault` resolves to whatever ambient identity `DefaultAzureCredential` discovers (environment, managed identity, developer sign-in, etc.), which is a single process-wide identity. It is well suited to single-identity services but does **not** distinguish between end users, so pooled connections opened with it are not isolated per user. The driver emits a one-time warning in this case. For genuinely multi-user workloads, use a per-user token provider instead. (You will also see this noted in the emitted log warning.)
+
 ### DBAPI v2.0 Compliance
 
 The Microsoft **mssql-python** module is designed to be fully compliant with the DB API 2.0 specification. This ensures that the driver adheres to a standardized interface for database access in Python, providing consistency and reliability across different database systems. Key aspects of DBAPI v2.0 compliance include:

mssql_python/connection.py

@@ -44,7 +44,8 @@
     extract_auth_type,
     process_auth_parameters,
     remove_sensitive_params,
-    get_auth_token,
+    get_auth_token_info,
+    compute_identity_key,
 )
 from mssql_python.constants import ConstantsDDBC, GetInfoConstants
 from mssql_python.connection_string_parser import _ConnectionStringParser
@@ -340,6 +341,24 @@ def __init__(
         # them because UID is already gone.
         self._credential_kwargs: Optional[Dict[str, str]] = None
 
+        # Composite, identity-aware pool key. Empty means "key the native pool
+        # on the connection string" (legacy behavior, used for non-token auth).
+        # For Entra access-token auth it is set to connStr + identity so that
+        # distinct identities never share a pooled connection.
+        self._pool_key: str = ""
+
+        # Optional deferred-attrs callback handed to the native layer. When set,
+        # native invokes it *only* when it actually opens a physical connection
+        # (a pool miss or a non-pooled connect), so a same-identity pool hit
+        # skips token acquisition entirely. None means "no deferred
+        # token" — the token (if any) is already in self._attrs_before.
+        #
+        # NOTE: This is an internal, private callback and is intentionally NOT
+        # the public ``token_provider=`` credential parameter. It returns the
+        # full connect-attrs dict lazily; hence the distinct name
+        # ``_token_factory``.
+        self._token_factory = None
+
         # Handle Entra ID authentication if specified.
         # The parsed dict is used directly — no re-parsing of the connection string.
         if _KEY_AUTHENTICATION in parsed_params:
@@ -357,11 +376,82 @@ def __init__(
                 # Strip sensitive params and rebuild the connection string.
                 sanitized = remove_sensitive_params(parsed_params)
                 self.connection_str = _ConnectionStringBuilder(sanitized).build()
-                token = get_auth_token(auth_type, credential_kwargs)
-                if token:
-                    self._attrs_before[ConstantsDDBC.SQL_COPT_SS_ACCESS_TOKEN.value] = token
                 self._credential_kwargs = credential_kwargs
 
+                # Make the pool key identity-aware so two callers using the
+                # same server but different Entra identities never reuse each
+                # other's authenticated connection. auth_type here
+                # is the process_auth_parameters result, which is truthy only
+                # for the token-bearing types (default/devicecode/msi/
+                # interactive-non-Windows); ServicePrincipal and Windows
+                # Interactive return None and keep their identity in the
+                # connection string, so they stay on the legacy connStr key.
+                #
+                # First try to derive the identity *without* a token. For MSI
+                # the client/object id comes straight from the params, so the
+                # pool key is known before any token exists — which lets us
+                # defer token acquisition to a provider callback that native
+                # invokes only on a pool miss.
+                #
+                # The factory returns ``(attrs, expires_on)``: the connect-attrs
+                # dict plus the token's POSIX-epoch expiry (or None). Native
+                # stores the expiry so it can refresh/discard a pooled
+                # connection whose token is near expiry on checkout.
+                token_attr = ConstantsDDBC.SQL_COPT_SS_ACCESS_TOKEN.value
+                base_attrs = self._attrs_before
+
+                def _make_token_factory():
+                    def _token_factory():
+                        attrs = dict(base_attrs)
+                        info = get_auth_token_info(auth_type, credential_kwargs)
+                        if info and info.token_struct:
+                            attrs[token_attr] = info.token_struct
+                            return attrs, info.expires_on
+                        return attrs, None
+
+                    return _token_factory
+
+                identity = compute_identity_key(auth_type, credential_kwargs)
+                if identity:
+                    # A real connection string can
+                    # never contain \0, so the composite key can never collide
+                    # with a bare connStr pool key. std::u16string map keys hold
+                    # embedded NULs fine and the key is never used as a C string.
+                    self._pool_key = self.connection_str + "\x00" + identity
+
+                    # Lazy token acquisition: native invokes this only
+                    # when it opens a physical connection, so same-identity pool
+                    # hits never pay for a token.
+                    self._token_factory = _make_token_factory()
+                else:
+                    # Token/account-dependent identity: acquire once to derive
+                    # the key. Interactive / Device-code yield a stable
+                    # home_account_id (key ``acct:``) so subsequent acquisitions
+                    # can be deferred to the factory (silent refresh reuses the
+                    # pool). DefaultAzureCredential / raw token key on the token
+                    # hash (``tok:``); the pooled connection is bound to that
+                    # exact token, so it is kept in attrs_before with no factory.
+                    info = get_auth_token_info(auth_type, credential_kwargs)
+                    token = info.token_struct if info else None
+                    home_account_id = info.home_account_id if info else None
+                    identity = compute_identity_key(
+                        auth_type,
+                        credential_kwargs,
+                        token_struct=token,
+                        home_account_id=home_account_id,
+                    )
+                    if identity:
+                        self._pool_key = self.connection_str + "\x00" + identity
+                    if identity and identity.startswith("acct:"):
+                        # Account-stable key: safe to defer to the factory so a
+                        # same-account pool hit skips token acquisition and a
+                        # near-expiry checkout can refresh silently.
+                        self._token_factory = _make_token_factory()
+                    elif token:
+                        # Token-hash key: bind the pooled connection to this
+                        # exact token so its hash always matches the pool key.
+                        self._attrs_before[token_attr] = token
+
             # Store auth type so bulkcopy() can acquire a fresh token later.
             # On Windows Interactive, process_auth_parameters returns None
             # (DDBC handles auth natively), so fall back to extract_auth_type.
@@ -403,7 +493,11 @@ def __init__(
         self._pooling = PoolingManager.is_enabled()
         try:
             self._conn = ddbc_bindings.Connection(
-                self.connection_str, self._pooling, self._attrs_before
+                self.connection_str,
+                self._pooling,
+                self._attrs_before,
+                self._pool_key,
+                self._token_factory,
             )
         except RuntimeError as e:
             _raise_connection_error(e)

mssql_python/pybind/connection/connection.cpp

@@ -12,7 +12,6 @@
 #include <utility>
 #include <vector>
 
-#define SQL_COPT_SS_ACCESS_TOKEN 1256  // Custom attribute ID for access token
 #define SQL_MAX_SMALL_INT 32767        // Maximum value for SQLSMALLINT
 
 // Logging uses LOG() macro for all diagnostic output
@@ -513,14 +512,66 @@ std::chrono::steady_clock::time_point Connection::lastUsed() const {
     return _lastUsed;
 }
 
+py::dict Connection::invokeTokenFactory(const py::object& tokenFactory,
+                                        long long& outExpiryEpoch) {
+    outExpiryEpoch = 0;
+    py::object result = tokenFactory();
+    // New contract: factory returns (attrs, expires_on). Remain
+    // backward compatible with the legacy contract where it returned a
+    // bare attrs dict.
+    if (py::isinstance<py::tuple>(result)) {
+        py::tuple parts = result.cast<py::tuple>();
+        py::dict attrs = parts[0].cast<py::dict>();
+        if (parts.size() > 1 && !parts[1].is_none()) {
+            outExpiryEpoch = parts[1].cast<long long>();
+        }
+        return attrs;
+    }
+    return result.cast<py::dict>();
+}
+
+void Connection::setTokenExpiry(long long epochSeconds) {
+    _tokenExpiryEpoch = epochSeconds;
+}
+
+bool Connection::isTokenNearExpiry(int thresholdSecs) const {
+    if (_tokenExpiryEpoch == 0) {
+        return false;  // Unknown / non-token auth: never treated as expiring.
+    }
+    const long long now = static_cast<long long>(
+        std::chrono::duration_cast<std::chrono::seconds>(
+            std::chrono::system_clock::now().time_since_epoch())
+            .count());
+    return (now + static_cast<long long>(thresholdSecs)) >= _tokenExpiryEpoch;
+}
+
+std::string Connection::currentAccessToken() const {
+    auto it = _attrBytesBuffers.find(SQL_COPT_SS_ACCESS_TOKEN);
+    return it != _attrBytesBuffers.end() ? it->second : std::string();
+}
+
 ConnectionHandle::ConnectionHandle(const std::u16string& connStr, bool usePool,
-                                   const py::dict& attrsBefore)
-    : _usePool(usePool), _connStr(connStr) {
+                                   const py::dict& attrsBefore, const std::u16string& poolKey,
+                                   const py::object& tokenFactory)
+    : _usePool(usePool), _connStr(connStr), _poolKey(poolKey.empty() ? connStr : poolKey) {
     if (_usePool) {
-        _conn = ConnectionPoolManager::getInstance().acquireConnection(_connStr, attrsBefore);
+        _conn = ConnectionPoolManager::getInstance().acquireConnection(_connStr, attrsBefore,
+                                                                       _poolKey, tokenFactory);
     } else {
         _conn = std::make_shared<Connection>(_connStr, false);
-        _conn->connect(attrsBefore);
+        // Non-pooled connect still honors the lazy token factory: a
+        // token is materialized only when a physical connection is opened. The
+        // factory may also carry the token expiry, recorded on the
+        // connection for completeness (a non-pooled connection is never reused,
+        // so expiry-aware checkout does not apply here).
+        if (tokenFactory && !tokenFactory.is_none()) {
+            long long expiry = 0;
+            py::dict connect_attrs = Connection::invokeTokenFactory(tokenFactory, expiry);
+            _conn->connect(connect_attrs);
+            _conn->setTokenExpiry(expiry);
+        } else {
+            _conn->connect(attrsBefore);
+        }
     }
 }
 
@@ -535,7 +586,7 @@ void ConnectionHandle::close() {
         ThrowStdException("Connection object is not initialized");
     }
     if (_usePool) {
-        ConnectionPoolManager::getInstance().returnConnection(_connStr, _conn);
+        ConnectionPoolManager::getInstance().returnConnection(_poolKey, _conn);
     } else {
         _conn->disconnect();
     }

mssql_python/pybind/connection/connection.h

@@ -8,6 +8,14 @@
 #include <mutex>
 #include <unordered_map>
 
+// Custom msodbcsql (SQL Server ODBC) connection-attribute id carrying the Entra
+// access-token struct. It is consumed once at login. Shared single source of
+// truth: connection.cpp applies it, and expiry-aware pooled checkout in
+// connection_pool.cpp uses it to compare a freshly minted token against the one
+// a pooled connection already holds. constexpr at namespace scope has internal
+// linkage, so each translation unit gets its own copy (no ODR issue).
+constexpr long SQL_COPT_SS_ACCESS_TOKEN = 1256;
+
 // Represents a single ODBC database connection.
 // Manages connection handles.
 // Note: This class does NOT implement pooling logic directly.
@@ -48,6 +56,29 @@ class Connection {
     void updateLastUsed();
     std::chrono::steady_clock::time_point lastUsed() const;
 
+    // Materialize connect-attrs from a Python token-factory callback.
+    // The factory may return either a bare attrs dict (legacy) or a
+    // ``(attrs, expires_on)`` tuple (expiry-aware pooling),
+    // where expires_on is the token's POSIX-epoch expiry or None. Returns the
+    // attrs dict; outExpiryEpoch is set to the expiry, or 0 when unknown.
+    // GIL must be held by the caller (it invokes Python).
+    static py::dict invokeTokenFactory(const py::object& tokenFactory,
+                                       long long& outExpiryEpoch);
+
+    // Record / inspect the pooled access token's expiry so a near-expiry
+    // connection is refreshed on checkout instead of being handed out with a
+    // token about to expire. Epoch seconds; 0 = unknown
+    // (non-token auth), for which isTokenNearExpiry() always returns false.
+    void setTokenExpiry(long long epochSeconds);
+    bool isTokenNearExpiry(int thresholdSecs) const;
+
+    // Returns the raw SQL_COPT_SS_ACCESS_TOKEN bytes this connection last
+    // authenticated with, or an empty string if it never used a token. Used by
+    // expiry-aware checkout to decide whether a freshly minted token differs
+    // from the one the pooled connection already holds: if unchanged, the
+    // healthy connection is reused; only a rotated token forces reopen.
+    std::string currentAccessToken() const;
+
     // Allocate a new statement handle on this connection.
     SqlHandlePtr allocStatementHandle();
 
@@ -69,6 +100,10 @@ class Connection {
     bool _autocommit = true;
     SqlHandlePtr _dbcHandle;
     std::chrono::steady_clock::time_point _lastUsed;
+    // POSIX-epoch expiry (seconds) of the access token this connection last
+    // authenticated with. 0 means unknown / non-token auth, in which case the
+    // expiry-aware checkout logic never treats the connection as near-expiry.
+    long long _tokenExpiryEpoch = 0;
     // Per-attribute owned buffers for connect attributes whose pointer the
     // driver may dereference *after* SQLSetConnectAttr returns (deferred
     // attributes, e.g. SQL_COPT_SS_ACCESS_TOKEN). Keyed by attribute ID so
@@ -97,7 +132,9 @@ class Connection {
 class ConnectionHandle {
   public:
     ConnectionHandle(const std::u16string& connStr, bool usePool,
-                     const py::dict& attrsBefore = py::dict());
+                     const py::dict& attrsBefore = py::dict(),
+                     const std::u16string& poolKey = std::u16string(),
+                     const py::object& tokenFactory = py::object());
     ~ConnectionHandle();
 
     void close();
@@ -115,4 +152,9 @@ class ConnectionHandle {
     std::shared_ptr<Connection> _conn;
     bool _usePool;
     std::u16string _connStr;
+    // Key under which this connection's pool is stored. Defaults to _connStr
+    // (legacy behavior) but is set to an identity-aware composite key for
+    // Entra access-token auth so distinct identities never share a pool.
+    // Empty is never stored; the ctor falls back to _connStr.
+    std::u16string _poolKey;
 };