diff --git a/README.md b/README.md index d59dd5b9..564f8460 100644 --- a/README.md +++ b/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 (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. + +> **`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.) Because a `DefaultAzureCredential` pool is keyed on the token hash rather than a stable identity, it is also **not** expiry-aware: the driver-acquired token is reused until the connection dies, so it is best suited to short-lived or single-identity use. + ### 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: diff --git a/generate_codecov.sh b/generate_codecov.sh index 66b3604b..99edc511 100644 --- a/generate_codecov.sh +++ b/generate_codecov.sh @@ -38,7 +38,19 @@ echo "===================================" # Cleanup old coverage rm -f .coverage coverage.xml python-coverage.info cpp-coverage.info total.info -rm -rf htmlcov unified-coverage +rm -f default.profraw default.profdata +rm -rf htmlcov unified-coverage profraw + +# Capture one raw profile *per process* so subprocess-based tests count too. +# Several pooling tests (idle eviction, pool-full, orphan return) must run in a +# fresh interpreter because the C++ pool config is locked in via std::call_once; +# they spawn `python -c` workers. With a fixed LLVM_PROFILE_FILE every process +# writes the same default.profraw and the last writer (the main pytest process) +# clobbers the workers, dropping their C++ coverage entirely. Using %p (PID) and +# %m (binary signature) gives each instrumented process its own file, which we +# merge below. Subprocess workers inherit this env var (os.environ.copy()). +mkdir -p "$(pwd)/profraw" +export LLVM_PROFILE_FILE="$(pwd)/profraw/default-%p-%m.profraw" # Run pytest with Python coverage (XML + HTML output) python -m pytest -v \ @@ -57,13 +69,22 @@ echo "===================================" echo "[STEP 3] Processing C++ coverage (Clang/LLVM)" echo "===================================" -# Merge raw profile data from pybind runs -if [ ! -f default.profraw ]; then - echo "[ERROR] default.profraw not found. Did you build with -fprofile-instr-generate?" +# Merge raw profile data from every instrumented process (main pytest run plus +# any `python -c` subprocess workers). Each wrote its own profraw/*.profraw via +# the LLVM_PROFILE_FILE pattern set in STEP 2. +shopt -s nullglob +PROFRAW_FILES=(profraw/*.profraw) +# Fallback: pick up a legacy single default.profraw if one exists in CWD. +if [ -f default.profraw ]; then + PROFRAW_FILES+=(default.profraw) +fi +if [ ${#PROFRAW_FILES[@]} -eq 0 ]; then + echo "[ERROR] No .profraw files found. Did you build with -fprofile-instr-generate?" exit 1 fi -llvm-profdata merge -sparse default.profraw -o default.profdata +echo "[INFO] Merging ${#PROFRAW_FILES[@]} raw profile file(s)" +llvm-profdata merge -sparse "${PROFRAW_FILES[@]}" -o default.profdata # Find the pybind .so file (Linux build) PYBIND_SO=$(find mssql_python -name "*.so" | head -n 1) diff --git a/mssql_python/auth.py b/mssql_python/auth.py index 4abd2fb9..9cdc50d5 100644 --- a/mssql_python/auth.py +++ b/mssql_python/auth.py @@ -8,7 +8,7 @@ import platform import struct import threading -from typing import Tuple, Dict, Optional +from typing import Tuple, Dict, NamedTuple, Optional from mssql_python.logging import logger from mssql_python.constants import ( @@ -28,12 +28,82 @@ # # Cache is keyed on (auth_type, sorted credential_kwargs), which is # bounded by the distinct credentials a single process ever uses. +# +# CAVEAT (interactive / device-code): the key for keyless interactive auth is +# just the auth type (e.g. "interactive"), so ALL interactive connections in a +# process share one credential instance and therefore one signed-in account — +# the first user to authenticate. This is what keeps pooled reconnects silent, +# but it means interactive/device-code auth cannot isolate multiple end users +# within a single process. Multi-user apps must bring their own per-user token +# provider instead of relying on interactive auth. (Documented for users in the +# Connection Pooling section of README.md.) _credential_cache: Dict[object, object] = {} _credential_cache_lock = threading.Lock() +# Stable home_account_id captured the first time an interactive/device-code +# credential runs authenticate(). Lets later *silent* get_token() acquisitions +# still key the pool on the account without re-authenticating. Keyed like +# _credential_cache and guarded by the same lock. +_account_id_cache: Dict[object, Optional[str]] = {} + # Canonical keys to strip when handing an Entra-token connection to ODBC. _SENSITIVE_KEYS = frozenset({_KEY_UID, _KEY_PWD, _KEY_TRUSTED_CONNECTION, _KEY_AUTHENTICATION}) + +class TokenInfo(NamedTuple): + """Result of an access-token acquisition. + + ``token_struct`` is the SQL Server / ODBC access-token blob (length-prefixed + UTF-16LE) placed in ``attrs_before``. ``expires_on`` is the POSIX epoch + second at which the underlying JWT expires, or ``None`` when the credential + did not surface an expiry (older SDKs / test doubles). ``home_account_id`` + is the stable per-account identifier surfaced by ``authenticate()`` for + Interactive / Device-code auth (used as the pool's security-context key so + a silent token refresh reuses the pool but a different account gets its own + pool); ``None`` for auth types that do not expose it. + """ + + token_struct: bytes + expires_on: Optional[int] + home_account_id: Optional[str] = None + + +# Scope requested for all SQL Server / Azure SQL access tokens. +_SQL_SCOPE = "https://database.windows.net/.default" + +# One-time guard so the DefaultAzureCredential multi-user pooling caveat is +# logged at most once per process rather than on every token acquisition. +_dac_pooling_warned = False +_dac_pooling_warned_lock = threading.Lock() + + +def _warn_default_credential_pooling() -> None: + """Emit a one-time warning that DefaultAzureCredential is not recommended + for multi-user connection pooling. + + DefaultAzureCredential walks a credential chain and can silently resolve to + a different underlying identity between calls, so we isolate it per token + (a new pool whenever the token changes) rather than reusing across + identities. Applications that need stable multi-user pooling should use a + concrete credential (Managed Identity, Service Principal) or supply their + own token provider. + """ + global _dac_pooling_warned # pylint: disable=global-statement + if _dac_pooling_warned: + return + with _dac_pooling_warned_lock: + if _dac_pooling_warned: + return + _dac_pooling_warned = True + logger.warning( + "DefaultAzureCredential is not recommended for multi-user connection " + "pooling: its credential chain can resolve to different identities " + "between calls, so connections are isolated per token (a new pool per " + "distinct token). Use a concrete credential (ManagedIdentity / " + "ServicePrincipal) or a token provider for stable pooled reuse." + ) + + # Map Authentication connection-string values to internal short names. _AUTH_TYPE_MAP: Dict[str, str] = { AuthType.INTERACTIVE.value: _AuthInternal.INTERACTIVE, @@ -77,9 +147,30 @@ def get_token_struct(token: str) -> bytes: @staticmethod def get_token(auth_type: str, credential_kwargs: Optional[Dict[str, str]] = None) -> bytes: """Get DDBC token struct for the specified authentication type.""" - token_struct, _ = AADAuth._acquire_token(auth_type, credential_kwargs) + token_struct, _, _, _ = AADAuth._acquire_token(auth_type, credential_kwargs) return token_struct + @staticmethod + def get_token_info( + auth_type: str, credential_kwargs: Optional[Dict[str, str]] = None + ) -> TokenInfo: + """Acquire a token and return both the ODBC struct and its expiry. + + Foundation for expiry-aware pool checkout: it surfaces ``expires_on`` so + a pooled connection's token can eventually be refreshed/discarded near + expiry. Interactive / Device-code additionally surface + ``home_account_id`` so the pool can key on a stable account identity + across silent refreshes. + """ + token_struct, _, expires_on, home_account_id = AADAuth._acquire_token( + auth_type, credential_kwargs + ) + return TokenInfo( + token_struct=token_struct, + expires_on=expires_on, + home_account_id=home_account_id, + ) + @staticmethod def get_raw_token(auth_type: str, credential_kwargs: Optional[Dict[str, str]] = None) -> str: """Acquire a raw JWT for the mssql-py-core connection (bulk copy). @@ -88,18 +179,54 @@ def get_raw_token(auth_type: str, credential_kwargs: Optional[Dict[str, str]] = built-in token cache can serve a valid token without a round-trip when the previous token has not yet expired. """ - _, raw_token = AADAuth._acquire_token(auth_type, credential_kwargs) + _, raw_token, _, _ = AADAuth._acquire_token(auth_type, credential_kwargs) return raw_token + @staticmethod + def _authenticate_interactive(credential) -> Optional[str]: + """Run the interactive ``authenticate()`` step and return the resulting + ``home_account_id``. + + This is the *fallback* half of the silent-first pattern: it is invoked + only when :meth:`_acquire_token` catches ``AuthenticationRequiredError`` + from a silent ``get_token()``, i.e. when interactive login is genuinely + required (first sign-in or a changed account). ``authenticate()`` + performs the interactive step and caches the record on the credential, + so subsequent silent ``get_token()`` calls for the same account succeed + without prompting; a different account yields a new ``home_account_id`` + (rotating the pool key). Returns ``None`` when the SDK / test double does + not expose ``authenticate()`` or ``home_account_id`` — the caller then + falls back to the token-hash pool key so the safety invariant still holds. + + Real authentication failures (e.g. a cancelled prompt raising + ``ClientAuthenticationError``) are allowed to propagate to + :meth:`_acquire_token`'s handler rather than being swallowed here. + """ + authenticate = getattr(credential, "authenticate", None) + if authenticate is None: + return None + try: + record = authenticate(scopes=[_SQL_SCOPE]) + except TypeError: + # Older/mocked signatures may not accept the scopes keyword. + record = authenticate() + return getattr(record, "home_account_id", None) + @staticmethod def _acquire_token( auth_type: str, credential_kwargs: Optional[Dict[str, str]] = None - ) -> Tuple[bytes, str]: - """Internal: acquire token and return (ddbc_struct, raw_jwt).""" + ) -> Tuple[bytes, str, Optional[int], Optional[str]]: + """Internal: acquire token and return + ``(ddbc_struct, raw_jwt, expires_on, home_account_id)``. + + ``home_account_id`` is populated only for Interactive / Device-code auth + (via ``authenticate()``); it is ``None`` for every other auth type. + """ # Import Azure libraries inside method to support test mocking # pylint: disable=import-outside-toplevel try: from azure.identity import ( + AuthenticationRequiredError, DefaultAzureCredential, DeviceCodeCredential, InteractiveBrowserCredential, @@ -131,6 +258,19 @@ def _acquire_token( credential_class.__name__, ) + # DefaultAzureCredential is not recommended for multi-user pooling; warn + # once so operators understand why these connections isolate per token. + if auth_type == _AuthInternal.DEFAULT: + _warn_default_credential_pooling() + + # Interactive / Device-code follow a silent-first pattern (see the + # get_token() call below): disable_automatic_authentication makes + # get_token() acquire silently from the credential's cached account and + # raise AuthenticationRequiredError (instead of prompting) when + # interactive login is genuinely required, at which point we call + # authenticate(). This keeps pooled reconnects prompt-free. + is_interactive = auth_type in (_AuthInternal.INTERACTIVE, _AuthInternal.DEVICE_CODE) + kwargs = credential_kwargs or {} cache_key = _credential_cache_key(auth_type, kwargs) try: @@ -140,20 +280,45 @@ def _acquire_token( "get_token: Creating new credential instance for auth_type=%s", auth_type, ) - _credential_cache[cache_key] = credential_class(**kwargs) + construct_kwargs: Dict[str, object] = dict(kwargs) + if is_interactive: + construct_kwargs["disable_automatic_authentication"] = True + _credential_cache[cache_key] = credential_class(**construct_kwargs) else: logger.debug( "get_token: Reusing cached credential instance for auth_type=%s", auth_type, ) credential = _credential_cache[cache_key] - raw_token = credential.get_token("https://database.windows.net/.default").token + + home_account_id: Optional[str] = None + if is_interactive: + # Silent-first: try the cached account, and only fall back to + # the interactive authenticate() step when the SDK signals it is + # actually required. A browser/device-code prompt on every + # pooled checkout would defeat pooling for these auth types. + with _credential_cache_lock: + home_account_id = _account_id_cache.get(cache_key) + try: + access_token = credential.get_token(_SQL_SCOPE) + except AuthenticationRequiredError: + home_account_id = AADAuth._authenticate_interactive(credential) + with _credential_cache_lock: + _account_id_cache[cache_key] = home_account_id + access_token = credential.get_token(_SQL_SCOPE) + else: + access_token = credential.get_token(_SQL_SCOPE) + + raw_token = access_token.token + # expires_on is a POSIX epoch second on azure-identity's AccessToken. + # getattr keeps older SDKs / test doubles (which may omit it) working. + expires_on = getattr(access_token, "expires_on", None) logger.info( "get_token: Azure AD token acquired successfully - token_length=%d chars", len(raw_token), ) token_struct = AADAuth.get_token_struct(raw_token) - return token_struct, raw_token + return token_struct, raw_token, expires_on, home_account_id except ClientAuthenticationError as e: logger.error( "get_token: Azure AD authentication failed - credential_class=%s, error=%s", @@ -402,30 +567,88 @@ def remove_sensitive_params(parsed_params: Dict[str, str]) -> Dict[str, str]: return {k: v for k, v in parsed_params.items() if k not in _SENSITIVE_KEYS} -def get_auth_token( +def get_auth_token_info( auth_type: str, credential_kwargs: Optional[Dict[str, str]] = None -) -> Optional[bytes]: - """Get DDBC authentication token struct based on auth type.""" - logger.debug("get_auth_token: Starting - auth_type=%s", auth_type) +) -> Optional[TokenInfo]: + """Acquire an access token and return both its ODBC struct and expiry. + + Returns ``None`` only when there is no auth type or when the driver handles + authentication natively (Windows Interactive), so callers can treat a + ``None`` result as "no Python-acquired token to place in attrs_before". + + For every other (token-backed) auth type this **fails closed**: an + acquisition error from :meth:`AADAuth.get_token_info` (``ValueError`` for an + unsupported auth type, ``RuntimeError`` for a credential/network/auth + failure) is allowed to propagate. Swallowing it and returning ``None`` would + let the caller fall through to an ODBC connect with ``Authentication=`` + stripped and no token attached, which either surfaces a confusing login + error or — worse — authenticates as an unintended identity. + """ + logger.debug("get_auth_token_info: Starting - auth_type=%s", auth_type) if not auth_type: - logger.debug("get_auth_token: No auth_type specified, returning None") return None - # Handle platform-specific logic for interactive auth if auth_type == _AuthInternal.INTERACTIVE and platform.system().lower() == "windows": - logger.debug("get_auth_token: Windows interactive auth - delegating to native handler") - return None # Let Windows handle AADInteractive natively + logger.debug("get_auth_token_info: Windows interactive auth - delegating to native handler") + return None - try: - token = AADAuth.get_token(auth_type, credential_kwargs) - logger.info("get_auth_token: Token acquired successfully - auth_type=%s", auth_type) - return token - except (ValueError, RuntimeError) as e: - logger.warning( - "get_auth_token: Token acquisition failed - auth_type=%s, error=%s", auth_type, str(e) - ) + info = AADAuth.get_token_info(auth_type, credential_kwargs) + logger.info("get_auth_token_info: Token acquired successfully - auth_type=%s", auth_type) + return info + + +def compute_identity_key( + auth_type: Optional[str], + credential_kwargs: Optional[Dict[str, str]] = None, + token_struct: Optional[bytes] = None, + home_account_id: Optional[str] = None, +) -> Optional[str]: + """Build the per-identity discriminator appended to the native pool key. + + This is the heart of identity-aware pooling: today the pool + keys only on the sanitized connection string, so two callers using + different Entra identities but the same server collide onto one pool. The + discriminator returned here is combined with the connection string to keep + distinct identities in distinct pools. + + Design invariant: when a token is present the pool key must NEVER be the + bare connection string; the fail-safe is the token hash. + + Return values: + * ``None`` — no token auth (SQL / trusted / native ServicePrincipal / + Windows Interactive). The connection string alone already isolates the + identity, so the pool key is unchanged (fully backward compatible). + * ``"msi:"`` / ``"msi:system"`` — Managed Identity, derived + from connection params *without* acquiring a token (enables skipping + token acquisition on a pool hit). + * ``"acct:"`` — Interactive / Device-code, keyed on the + stable account id from ``authenticate()`` so a silent token refresh + reuses the pool but a different signed-in account gets its own pool. + * ``"tok:"`` — fail-safe hash of the token struct for auth types + whose identity is not derivable from params (DefaultAzureCredential, + raw token, and interactive / device-code when ``home_account_id`` is + unavailable). Requires ``token_struct``; when it is not supplied the + function returns ``None`` to signal "acquire a token, then recompute". + """ + if not auth_type: return None + if auth_type == _AuthInternal.MSI: + client_id = (credential_kwargs or {}).get("client_id") + return f"msi:{client_id}" if client_id else "msi:system" + + # Interactive / Device-code: prefer the stable account id when available so + # silent refreshes reuse the pool; fall back to the token hash otherwise. + if auth_type in (_AuthInternal.INTERACTIVE, _AuthInternal.DEVICE_CODE) and home_account_id: + return f"acct:{home_account_id}" + + if token_struct is not None: + return "tok:" + hashlib.sha256(token_struct).hexdigest() + + # A token-backed auth type whose key needs the token, but none was + # supplied yet. Signal the caller to acquire first, then recompute. + return None + def extract_auth_type(parsed_params: Dict[str, str]) -> Optional[str]: """Map the Authentication connection-string value to an internal type name. diff --git a/mssql_python/connection.py b/mssql_python/connection.py index 94fb0924..eaed870b 100644 --- a/mssql_python/connection.py +++ b/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,132 @@ 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 _acquire_token_info(): + # DB-API boundary: get_auth_token_info fails closed by + # letting the underlying Azure error propagate as a + # ValueError (unsupported auth type) or RuntimeError + # (credential/network/auth failure). Re-wrap those as a + # DB-API 2.0 InterfaceError so every auth failure surfaced + # from connect() / the token factory is a consistent, + # catchable driver exception rather than a bare RuntimeError. + try: + return get_auth_token_info(auth_type, credential_kwargs) + except (RuntimeError, ValueError) as e: + raise InterfaceError( + driver_error=( + "Failed to acquire an Entra ID access token for " + f"authentication type '{auth_type}': {e}" + ), + ddbc_error=str(e), + ) from e + + def _make_token_factory(): + # Deferred connect-attrs provider handed to native. Native + # invokes it only when it actually opens a physical + # connection (a pool miss or non-pooled connect), so a + # same-identity pool hit never pays for a token. + attrs = dict(base_attrs) + info = _acquire_token_info() + if info and info.token_struct: + attrs[token_attr] = info.token_struct + return attrs, info.expires_on + # Fail closed: a token-backed pool must never open a + # physical connection without the token it is meant to + # carry. get_auth_token_info already raises when + # acquisition fails; this guards the empty-token edge so + # native never connects with Authentication= stripped. + raise InterfaceError( + driver_error=( + "Unable to acquire an Entra ID access token for " + f"authentication type '{auth_type}'." + ), + ddbc_error="Token factory produced no usable token.", + ) + + 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 = _acquire_token_info() + 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. + # No token factory is attached, so this pool is NOT + # expiry-aware — the near-expiry refresh in the native + # layer only runs for factory-backed (msi:/acct:) pools. + # A driver-acquired DAC/raw token is reused until the + # connection dies; see the pooling notes in the README. + self._attrs_before[token_attr] = token + else: + # Fail closed: a token-backed auth type reached here but + # we could derive neither a deferred factory nor an + # eager token, so connecting now would strip + # Authentication= and open with no token (potentially + # authenticating as an unintended identity). Surface a + # clear error instead of silently falling through. + raise InterfaceError( + driver_error=( + "Unable to acquire an Entra ID access token for " + f"authentication type '{auth_type}'." + ), + ddbc_error="Token acquisition returned no usable 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. @@ -397,13 +537,52 @@ def __init__( # Initialize search escape character self._searchescape = None + # Safety net for the raw access-token pattern: a caller may pass + # SQL_COPT_SS_ACCESS_TOKEN directly in attrs_before with no + # Authentication= keyword (the documented msodbcsql raw-token pattern), + # or via the public token_provider= API. That path never runs the + # identity-key logic above, so without this guard the pool key would + # stay empty and two callers with different raw tokens against the same + # server would share a pool — handing user B user A's authenticated + # connection on a pool hit. Bind the pool key to the token hash so + # distinct tokens never collide (upholds the "a token is present => + # key is never the bare connStr" invariant). + if not self._pool_key: + _raw_token = self._attrs_before.get(ConstantsDDBC.SQL_COPT_SS_ACCESS_TOKEN.value) + if _raw_token is not None and not isinstance(_raw_token, (bytes, bytearray)): + # Fail closed: an access token supplied as anything other than + # raw bytes (e.g. a str) cannot be hashed into an identity-aware + # pool key, so it would fall through with the bare connStr key + # and two callers passing different str tokens against the same + # server could share a pooled, authenticated connection. Reject + # it up front with a clear DB-API error instead of relying on + # ODBC to mangle/reject the byte-struct downstream. The native + # setAttribute() enforces the same rule, so the invariant holds + # at both boundaries. + raise InterfaceError( + driver_error=( + "SQL_COPT_SS_ACCESS_TOKEN must be supplied as bytes (the " + "raw [length][UTF-16LE token] struct), not " + f"{type(_raw_token).__name__}." + ), + ddbc_error="Non-binary access token attribute rejected.", + ) + if isinstance(_raw_token, (bytes, bytearray)): + _token_identity = compute_identity_key("default", token_struct=bytes(_raw_token)) + if _token_identity: + self._pool_key = self.connection_str + "\x00" + _token_identity + # Auto-enable pooling if user never called if not PoolingManager.is_initialized(): PoolingManager.enable() 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) diff --git a/mssql_python/pooling.py b/mssql_python/pooling.py index a2811d9f..6543b07e 100644 --- a/mssql_python/pooling.py +++ b/mssql_python/pooling.py @@ -66,6 +66,14 @@ def enable(cls, max_size: int = 100, idle_timeout: int = 600) -> None: cls._config["max_size"] = max_size cls._config["idle_timeout"] = idle_timeout cls._enabled = True + # Re-arm the disable guard: enable_pooling() calls setAccepting(true) + # on the native manager, so the Python state must mirror that a + # fresh disable() is once again required to disarm it. Without this + # reset, an enable -> disable -> enable -> disable sequence would + # leave _pools_closed True after the first disable, causing the + # second disable to skip disable_pooling() and desync from the + # native manager (which is still accepting). + cls._pools_closed = False cls._initialized = True logger.info("PoolingManager.enable: Connection pooling enabled successfully") @@ -83,7 +91,7 @@ def disable(cls) -> None: cls._enabled and not cls._pools_closed ): # Only cleanup if enabled and not already closed logger.info("PoolingManager.disable: Closing connection pools") - ddbc_bindings.close_pooling() + ddbc_bindings.disable_pooling() logger.info("PoolingManager.disable: Connection pools closed successfully") else: logger.debug("PoolingManager.disable: Pooling already disabled or closed") diff --git a/mssql_python/pybind/connection/connection.cpp b/mssql_python/pybind/connection/connection.cpp index 4b366575..beb23fb3 100644 --- a/mssql_python/pybind/connection/connection.cpp +++ b/mssql_python/pybind/connection/connection.cpp @@ -12,7 +12,6 @@ #include #include -#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 @@ -298,6 +297,25 @@ SQLRETURN Connection::setAttribute(SQLINTEGER attribute, py::object value) { // SQLPOINTER ptr = nullptr; // SQLINTEGER length = 0; + // Fail closed on a non-binary access token. SQL_COPT_SS_ACCESS_TOKEN (1256) + // MUST be the raw [DWORD byte-length][UTF-16LE token] struct passed as + // bytes/bytearray. If a caller supplies it as a py::str, the str->UTF-16 + // cast in the string branch below would mangle that struct; worse, the + // Python identity-aware pool-key logic only hashes bytes/bytearray tokens, + // so a str token slips through with the bare connection-string pool key and + // two callers passing different str tokens against the same server could + // share a pooled, authenticated connection. Reject any non-binary token at + // this native boundary so the cross-identity invariant ("a token is present + // => the pool key is never the bare connStr") holds regardless of how the + // Connection was constructed. + if (attribute == SQL_COPT_SS_ACCESS_TOKEN && !py::isinstance(value) && + !py::isinstance(value)) { + LOG("Rejecting non-binary SQL_COPT_SS_ACCESS_TOKEN (attribute=%d): access token " + "must be bytes/bytearray", + attribute); + return SQL_ERROR; + } + if (py::isinstance(value)) { // Get the integer value int64_t longValue = value.cast(); @@ -513,14 +531,90 @@ 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(result)) { + py::tuple parts = result.cast(); + // Defensive: a well-formed factory always returns at least (attrs,). + // Guard the index so a misbehaving/empty tuple falls through to the + // cast below (which raises a clear tuple->dict error) instead of an + // out-of-range access on parts[0]. + if (parts.size() >= 1) { + py::dict attrs = parts[0].cast(); + if (parts.size() > 1 && !parts[1].is_none()) { + outExpiryEpoch = parts[1].cast(); + } + return attrs; + } + } + return result.cast(); +} + +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( + std::chrono::duration_cast( + std::chrono::system_clock::now().time_since_epoch()) + .count()); + return (now + static_cast(thresholdSecs)) >= _tokenExpiryEpoch; +} + +std::string Connection::currentAccessToken() const { + auto it = _attrBytesBuffers.find(SQL_COPT_SS_ACCESS_TOKEN); + return it != _attrBytesBuffers.end() ? it->second : std::string(); +} + +bool Connection::accessTokenEquals(const std::string& token) const { + auto it = _attrBytesBuffers.find(SQL_COPT_SS_ACCESS_TOKEN); + if (it == _attrBytesBuffers.end()) { + // No token on this connection; matches only an empty comparand, + // mirroring currentAccessToken() == token semantics. + return token.empty(); + } + return it->second == token; +} + 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); - } else { + _conn = ConnectionPoolManager::getInstance().acquireConnection(_connStr, attrsBefore, + _poolKey, tokenFactory); + // acquireConnection returns nullptr when pooling was disabled out from + // under us (a disable_pooling() won the race). Fall back to a non-pooled + // connection and flip _usePool so close() disconnects it directly rather + // than trying to return it to a pool that no longer exists. + if (!_conn) { + _usePool = false; + } + } + if (!_usePool) { _conn = std::make_shared(_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 +629,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(); } diff --git a/mssql_python/pybind/connection/connection.h b/mssql_python/pybind/connection/connection.h index 981bbc0a..0ae81ecb 100644 --- a/mssql_python/pybind/connection/connection.h +++ b/mssql_python/pybind/connection/connection.h @@ -8,6 +8,14 @@ #include #include +// 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,36 @@ 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; + + // True if this connection's last access token equals `token`, compared in + // place against the internal buffer (no copy). Used by the sibling-drain + // sweep so rotating one token does not copy every pooled connection's token + // per comparison. Matches currentAccessToken() == token semantics (a + // connection with no token matches only an empty comparand). + bool accessTokenEquals(const std::string& token) const; + // Allocate a new statement handle on this connection. SqlHandlePtr allocStatementHandle(); @@ -69,6 +107,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 +139,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 +159,9 @@ class ConnectionHandle { std::shared_ptr _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; }; diff --git a/mssql_python/pybind/connection/connection_pool.cpp b/mssql_python/pybind/connection/connection_pool.cpp index db891d08..61435685 100644 --- a/mssql_python/pybind/connection/connection_pool.cpp +++ b/mssql_python/pybind/connection/connection_pool.cpp @@ -2,6 +2,8 @@ // Licensed under the MIT license. #include "connection/connection_pool.h" +#include +#include #include #include #include @@ -9,11 +11,53 @@ // Logging uses LOG() macro for all diagnostic output #include "logger_bridge.hpp" +// Refresh threshold for expiry-aware checkout: a pooled connection whose access +// token expires within this many seconds is discarded and reopened with a fresh +// token rather than handed out (<=5 min). +static constexpr int TOKEN_EXPIRY_THRESHOLD_SECS = 300; + +// True only when *expiryEpoch* is a known POSIX-second expiry that is safely +// beyond now + thresholdSecs. A zero/negative (unknown/missing) expiry is +// treated as NOT safe so the caller fails closed rather than reusing a token it +// cannot prove is still valid. Caller need not hold any lock; reads the wall +// clock only. +static bool tokenExpirySafelyBeyond(long long expiryEpoch, int thresholdSecs) { + if (expiryEpoch <= 0) { + return false; + } + const long long now = static_cast( + std::chrono::duration_cast( + std::chrono::system_clock::now().time_since_epoch()) + .count()); + return (now + static_cast(thresholdSecs)) < expiryEpoch; +} + +// Pull the raw access-token bytes out of a connect-attrs dict returned by the +// token factory, or an empty string if the dict carries no token. Caller must +// hold the GIL. Keys in the dict are attribute ids (ints). SQL_COPT_SS_ACCESS_TOKEN +// is defined once in connection.h (shared with connection.cpp). +static std::string extractAccessToken(const py::dict& attrs) { + for (auto item : attrs) { + if (py::isinstance(item.first) && + item.first.cast() == SQL_COPT_SS_ACCESS_TOKEN) { + try { + return item.second.cast(); + } catch (const py::cast_error&) { + return std::string(); + } + } + } + return std::string(); +} + ConnectionPool::ConnectionPool(size_t max_size, int idle_timeout_secs) - : _max_size(max_size), _idle_timeout_secs(idle_timeout_secs), _current_size(0) {} + : _max_size(max_size), + _idle_timeout_secs(idle_timeout_secs), + _current_size(0) {} std::shared_ptr ConnectionPool::acquire(const std::u16string& connStr, - const py::dict& attrs_before) { + const py::dict& attrs_before, + const py::object& token_factory) { std::vector> to_disconnect; std::shared_ptr valid_conn = nullptr; bool needs_connect = false; @@ -50,54 +94,181 @@ std::shared_ptr ConnectionPool::acquire(const std::u16string& connSt // mutex. isAlive() and reset() perform ODBC calls that release the // GIL; calling them while holding the mutex would create a mutex/GIL // lock-ordering deadlock when multiple threads acquire concurrently. + // + // Expiry-aware checkout may capture a freshly minted token here so a + // rotated-token pool can be reopened without invoking the factory twice. + py::dict pending_attrs; + long long pending_expiry = 0; + bool have_pending_token = false; while (true) { std::shared_ptr candidate; { - std::lock_guard lock(_mutex); + std::unique_lock lock(_mutex); if (_pool.empty()) { // No more candidates — try to reserve a slot for a new connection. if (_current_size < _max_size) { valid_conn = std::make_shared(connStr, true); ++_current_size; needs_connect = true; - } else { - // NOTE: Another thread may be validating a popped candidate - // outside the mutex right now. If that candidate fails, a - // slot will open up — but we can't wait for it here without - // adding a condition-variable retry loop. This is an - // acceptable trade-off: transient "pool full" errors under - // heavy contention are rare and callers can retry. - throw std::runtime_error("ConnectionPool::acquire: pool size limit reached"); + break; } - break; + // Pool is full — throw immediately. Another thread may be + // validating a popped candidate outside the mutex right now, so + // a transient "pool full" is an acceptable trade-off that + // callers can retry. + throw std::runtime_error( + "ConnectionPool::acquire: pool size limit reached"); } candidate = _pool.front(); _pool.pop_front(); } // Validate the candidate outside the mutex. + bool reuse_candidate = false; try { - if (candidate->isAlive() && candidate->reset()) { - valid_conn = candidate; - break; + if (token_factory && !token_factory.is_none() && + candidate->isTokenNearExpiry(TOKEN_EXPIRY_THRESHOLD_SECS)) { + // Expiry-aware checkout with token compare: the pooled token is + // at/near expiry, so mint a fresh one and compare. If the + // provider returns the SAME token (its cache is still valid), + // the connection is healthy — refresh the recorded expiry and + // reuse it rather than needlessly churning. Only a DIFFERENT + // (rotated) token forces discard-and-reopen, and we carry the + // fresh attrs forward so the reopen below does not invoke the + // factory a second time. + long long fresh_expiry = 0; + py::dict fresh_attrs = + Connection::invokeTokenFactory(token_factory, fresh_expiry); + std::string fresh_token = extractAccessToken(fresh_attrs); + if (!fresh_token.empty() && + fresh_token == candidate->currentAccessToken() && + tokenExpirySafelyBeyond(fresh_expiry, TOKEN_EXPIRY_THRESHOLD_SECS)) { + // Same token AND its refreshed expiry is safely beyond the + // threshold: the provider's cache is still valid and the + // connection is healthy, so refresh the recorded expiry and + // reuse. We deliberately do NOT reuse when the returned + // expiry is unknown (<=0) or still inside the threshold — + // extending the recorded expiry and handing the connection + // back would defeat the very refresh this checkout intended + // (the token could expire mid-query). Those cases fall + // through to discard-and-reopen below. + // + // Narrow edge: a MISBEHAVING provider that repeatedly hands + // back the same token still inside the threshold makes every + // checkout discard + reopen (and get the same near-expiry + // token) — pure churn, no benefit. This is acceptable: a + // well-behaved azure-identity credential refreshes + // proactively (returning a token with a fresh, far-out + // expiry) before the threshold, so the safe-reuse path above + // 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 (!reuse_candidate) { + // The token is still valid but the socket is dead + // (isAlive()/reset() failed). We already minted the + // fresh attrs, so carry them forward and let Phase 3 + // reopen with them instead of invoking the factory a + // second time. No sibling drain: siblings hold the same + // still-valid token and remain reusable. + pending_attrs = fresh_attrs; + pending_expiry = fresh_expiry; + have_pending_token = true; + } + } else { + // Token rotated, or the "fresh" token is still at/near + // expiry (or has an unknown expiry): discard and reopen with + // the fresh attrs. Remember the fresh token to reopen with, + // and eagerly drain the sibling idle connections that still + // hold the now-stale token. They were all minted from the + // same provider before the rotation, so they are equally + // stale; discarding them together here avoids rediscovering + // each one (and paying another factory compare) on later + // checkouts. No ODBC calls under the mutex — the actual + // disconnects happen in Phase 4, outside the lock. + pending_attrs = fresh_attrs; + pending_expiry = fresh_expiry; + have_pending_token = true; + const std::string stale_token = candidate->currentAccessToken(); + if (!stale_token.empty()) { + std::lock_guard lock(_mutex); + _pool.erase( + std::remove_if( + _pool.begin(), _pool.end(), + [&](const std::shared_ptr& sibling) { + if (sibling->accessTokenEquals(stale_token)) { + to_disconnect.push_back(sibling); + if (_current_size > 0) --_current_size; + return true; + } + return false; + }), + _pool.end()); + } + } + } else { + reuse_candidate = candidate->isAlive() && candidate->reset(); } } catch (const std::exception& ex) { LOG("Candidate connection validation failed: %s", ex.what()); } - // Candidate is dead or reset failed — mark for disconnect and - // decrement the pool size. + if (reuse_candidate) { + valid_conn = candidate; + break; + } + + // Candidate is dead, reset failed, or its token rotated — mark for + // disconnect and decrement the pool size. to_disconnect.push_back(candidate); { std::lock_guard lock(_mutex); 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) { + std::lock_guard lock(_mutex); + if (_current_size < _max_size) { + valid_conn = std::make_shared(connStr, true); + ++_current_size; + 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. + } } // Phase 3: Connect the new connection outside the mutex. if (needs_connect) { try { - valid_conn->connect(attrs_before); + if (have_pending_token) { + // Reopen with the fresh token captured during expiry-aware + // checkout (the previous connection's token had rotated). + valid_conn->connect(pending_attrs); + valid_conn->setTokenExpiry(pending_expiry); + } else if (token_factory && !token_factory.is_none()) { + // Lazy token acquisition: only now, when a physical + // connection is actually being opened, do we materialize the + // token. On a pool reuse this whole branch is skipped, so a + // same-identity hit never acquires a token. The GIL is held here + // (connect() releases it only around the ODBC call itself), so + // invoking the Python callback is safe. + long long expiry = 0; + py::dict connect_attrs = Connection::invokeTokenFactory(token_factory, expiry); + valid_conn->connect(connect_attrs); + // Record the token expiry so a later checkout can refresh this + // connection before the token lapses. + valid_conn->setTokenExpiry(expiry); + } else { + valid_conn->connect(attrs_before); + } } catch (...) { // Connect failed — release the reserved slot { @@ -144,6 +315,37 @@ void ConnectionPool::release(std::shared_ptr conn) { } } +bool ConnectionPool::canEvict() { + std::lock_guard lock(_mutex); + // Never evict while any connection is checked out or in-flight. Reserved + // capacity (_current_size) beyond what is sitting idle in _pool means a + // caller still holds one, so the pool must stay. + size_t checked_out = (_current_size > _pool.size()) ? (_current_size - _pool.size()) : 0; + if (checked_out > 0) { + return false; + } + // Nothing checked out and the pool is empty: safe to drop immediately. + if (_pool.empty()) { + return true; + } + // Nothing checked out but idle connections remain. Evict the whole pool + // only once EVERY pooled connection has been idle longer than the idle + // timeout. This is what reclaims pools for rotating / single-use identities + // (e.g. per-request Entra users keyed by token hash): such a pool is never + // acquired again, so its idle connection is never pruned by acquire() and + // _current_size would otherwise stay > 0 forever. Evaluating the idle + // timeout here lets the next acquireConnection() on any key sweep it away. + auto now = std::chrono::steady_clock::now(); + for (const auto& conn : _pool) { + auto idle_time = + std::chrono::duration_cast(now - conn->lastUsed()).count(); + if (idle_time <= _idle_timeout_secs) { + return false; + } + } + return true; +} + void ConnectionPool::close() { std::vector> to_close; { @@ -169,29 +371,92 @@ ConnectionPoolManager& ConnectionPoolManager::getInstance() { } std::shared_ptr ConnectionPoolManager::acquireConnection(const std::u16string& connStr, - const py::dict& attrs_before) { + const py::dict& attrs_before, + const std::u16string& pool_key, + const py::object& token_factory) { + // Key the pool by pool_key when provided (identity-aware), + // else fall back to the connection string (legacy behavior). + const std::u16string& key = pool_key.empty() ? connStr : pool_key; std::shared_ptr pool; + std::vector> evicted; { std::lock_guard lock(_manager_mutex); - auto& pool_ref = _pools[connStr]; + // Pooling disabled (a concurrent disable_pooling() disarmed us): decline + // to create or hand out a pool. Because this check and the pool creation + // below share _manager_mutex with the setAccepting(false) in + // disable_pooling(), the decision is atomic — a connect either creates + // its pool before the disable (and closePools() then reaps it) or sees + // _accepting == false here and never creates one. The caller + // (ConnectionHandle) falls back to a non-pooled connection. + if (!_accepting) { + return nullptr; + } + // Lazy eviction: drop pools whose connections are all idle past the + // idle timeout (and none checked out) so distinct short-lived + // identities (e.g. per-request Entra users keyed by token hash) do not + // accumulate pools forever. canEvict() only inspects state (no ODBC + // calls), so it is safe under _manager_mutex; the actual disconnects + // happen via close() below, outside the lock. The pool we are about to + // use is skipped so it is never evicted from under us. + // + // The sweep is O(pools × idle-conns) under the global mutex, so it is + // throttled: a pool can only become evictable after its connections + // sit idle past the idle timeout, so sweeping more often than that + // window is pure overhead. Between sweeps we skip straight to the pool + // lookup, keeping the hot path cheap under a many-identity connect load. + auto now = std::chrono::steady_clock::now(); + auto sweep_interval = std::chrono::seconds(std::max(1, _default_idle_secs)); + if (now - _last_sweep >= sweep_interval) { + _last_sweep = now; + for (auto it = _pools.begin(); it != _pools.end();) { + // Only evict a pool that no one else is holding: use_count == 1 + // means the map is the sole owner. An in-flight acquirer copies + // its pool shared_ptr while holding _manager_mutex (same section + // as this sweep) and keeps that copy across the unlocked + // acquire(); returnConnection() likewise takes a ref under the + // mutex before releasing. Either bumps use_count above 1 for the + // whole window, so this guard prevents evicting — and then + // closing (disconnecting) — a pool a peer thread has already + // selected but not yet finished using. + if (it->first != key && it->second && it->second.use_count() == 1 && + it->second->canEvict()) { + evicted.push_back(it->second); + it = _pools.erase(it); + } else { + ++it; + } + } + } + auto& pool_ref = _pools[key]; if (!pool_ref) { LOG("Creating new connection pool"); pool_ref = std::make_shared(_default_max_size, _default_idle_secs); } pool = pool_ref; } + // Close evicted pools outside _manager_mutex: close() disconnects ODBC + // handles (releasing the GIL), which must never run while holding + // _manager_mutex or we risk a mutex/GIL lock-ordering deadlock. + for (auto& evicted_pool : evicted) { + try { + evicted_pool->close(); + } catch (const std::exception& ex) { + LOG("ConnectionPoolManager: closing evicted pool failed: %s", ex.what()); + } + } // Call acquire() outside _manager_mutex. acquire() may release the GIL // during the ODBC connect call; holding _manager_mutex across that would - // create a mutex/GIL lock-ordering deadlock. - return pool->acquire(connStr, attrs_before); + // create a mutex/GIL lock-ordering deadlock. connStr (not key) is used to + // establish new physical connections. + return pool->acquire(connStr, attrs_before, token_factory); } -void ConnectionPoolManager::returnConnection(const std::u16string& conn_str, +void ConnectionPoolManager::returnConnection(const std::u16string& pool_key, const std::shared_ptr conn) { std::shared_ptr pool; { std::lock_guard lock(_manager_mutex); - auto it = _pools.find(conn_str); + auto it = _pools.find(pool_key); if (it != _pools.end()) { pool = it->second; } @@ -199,6 +464,22 @@ void ConnectionPoolManager::returnConnection(const std::u16string& conn_str, // Call release() outside _manager_mutex to avoid deadlock. if (pool) { pool->release(conn); + } else { + // No pool is registered under this key (e.g. the pool was lazily + // evicted while this connection was checked out, or the key changed). + // Disconnect the orphaned connection instead of leaking it: + // dropping the shared_ptr alone would keep the ODBC handle + // around until GC, and returnConnection is the deterministic close + // path. Done outside _manager_mutex (disconnect releases the GIL). + if (conn) { + try { + conn->disconnect(); + } catch (const std::exception& ex) { + LOG("ConnectionPoolManager::returnConnection: disconnect of orphaned " + "connection failed: %s", + ex.what()); + } + } } } @@ -206,17 +487,47 @@ void ConnectionPoolManager::configure(int max_size, int idle_timeout_secs) { std::lock_guard lock(_manager_mutex); _default_max_size = max_size; _default_idle_secs = idle_timeout_secs; + // Reset the sweep throttle so the new idle timeout takes effect on the very + // next acquireConnection() instead of waiting out a stale interval. + _last_sweep = std::chrono::steady_clock::time_point{}; } void ConnectionPoolManager::closePools() { - std::lock_guard lock(_manager_mutex); - // Keep _manager_mutex held for the full close operation so that - // acquireConnection()/returnConnection() cannot create or use pools - // while closePools() is in progress. - for (auto& [conn_str, pool] : _pools) { - if (pool) { + // Mirror the eviction-sweep pattern: under the mutex, move every pool into + // a local vector and clear the map, then release the mutex before closing. + // close() disconnects ODBC handles (releasing the GIL), which must never + // run while holding _manager_mutex or we risk a mutex/GIL lock-ordering + // deadlock with a concurrent acquireConnection()/returnConnection(). + std::vector> to_close; + { + std::lock_guard lock(_manager_mutex); + to_close.reserve(_pools.size()); + for (auto& [conn_str, pool] : _pools) { + if (pool) { + to_close.push_back(pool); + } + } + _pools.clear(); + // Nothing left to sweep; reset the throttle so a fresh pool set after + // this is swept on its next acquireConnection(). + _last_sweep = std::chrono::steady_clock::time_point{}; + } + // Close each pool outside _manager_mutex. + for (auto& pool : to_close) { + try { pool->close(); + } catch (const std::exception& ex) { + LOG("ConnectionPoolManager::closePools: closing pool failed: %s", ex.what()); } } - _pools.clear(); +} + +size_t ConnectionPoolManager::poolCount() { + std::lock_guard lock(_manager_mutex); + return _pools.size(); +} + +void ConnectionPoolManager::setAccepting(bool accepting) { + std::lock_guard lock(_manager_mutex); + _accepting = accepting; } diff --git a/mssql_python/pybind/connection/connection_pool.h b/mssql_python/pybind/connection/connection_pool.h index f4d9e0f4..6ee26994 100644 --- a/mssql_python/pybind/connection/connection_pool.h +++ b/mssql_python/pybind/connection/connection_pool.h @@ -19,9 +19,17 @@ class ConnectionPool { public: ConnectionPool(size_t max_size, int idle_timeout_secs); - // Acquires a connection from the pool or creates a new one if under limit + // Acquires a connection from the pool or creates a new one if under limit. + // token_factory, when set, is a Python callable returning the connect-attrs + // (including the access token) to connect with; it is invoked *only* when a + // new physical connection is opened, so a pool hit skips it entirely. + // (Internal callback — unrelated to the public token_provider= API.) + // When token_factory is set, a pooled candidate whose access token is within + // the expiry threshold is discarded and reopened so a caller never receives + // a connection with an about-to-expire token. std::shared_ptr acquire(const std::u16string& connStr, - const py::dict& attrs_before = py::dict()); + const py::dict& attrs_before = py::dict(), + const py::object& token_factory = py::object()); // Returns a connection to the pool for reuse void release(std::shared_ptr conn); @@ -29,6 +37,10 @@ class ConnectionPool { // Closes all connections in the pool, releasing resources void close(); + // True when the pool holds no live or in-flight connections and can be + // dropped by the manager to reclaim memory (lazy eviction). + bool canEvict(); + private: size_t _max_size; // Maximum number of connections allowed int _idle_timeout_secs; // Idle time before connections are stale @@ -45,16 +57,38 @@ class ConnectionPoolManager { void configure(int max_size, int idle_timeout); - // Gets a connection from the appropriate pool (creates one if none exists) - std::shared_ptr acquireConnection(const std::u16string& conn_str, - const py::dict& attrs_before = py::dict()); - - // Returns a connection to its original pool - void returnConnection(const std::u16string& conn_str, std::shared_ptr conn); + // Gets a connection from the appropriate pool (creates one if none exists). + // The pool is keyed by pool_key when supplied, else by conn_str. conn_str + // is always used to establish new physical connections. Keying separately + // keeps distinct Entra identities in distinct pools. + // token_factory is forwarded to ConnectionPool::acquire for lazy token + // acquisition on a pool miss. + // + // Returns nullptr when pooling is not currently accepting new pools (i.e. a + // concurrent disable won the race). The enabled-check and the pool creation + // share _manager_mutex, so this decision is atomic with respect to + // closePools(): the caller must fall back to a non-pooled connection. + std::shared_ptr acquireConnection( + const std::u16string& conn_str, const py::dict& attrs_before = py::dict(), + const std::u16string& pool_key = std::u16string(), + const py::object& token_factory = py::object()); + + // Arms (true) or disarms (false) new-pool creation. Disarming, done under + // _manager_mutex, guarantees that any acquireConnection() serialized after + // it declines to create a pool, closing the disable()-vs-connect() race. + void setAccepting(bool accepting); + + // Returns a connection to its original pool, identified by pool_key + // (the same key passed to acquireConnection). + void returnConnection(const std::u16string& pool_key, std::shared_ptr conn); // Closes all pools and their connections void closePools(); + // Diagnostic: number of pools currently tracked by the manager. Used by + // tests to observe lazy eviction of idle identity pools. + size_t poolCount(); + private: ConnectionPoolManager() = default; ~ConnectionPoolManager() = default; @@ -67,6 +101,22 @@ class ConnectionPoolManager { size_t _default_max_size = 10; int _default_idle_secs = 300; + // When false, acquireConnection() refuses to create/hand out pools so a + // connect racing a disable() cannot resurrect a pool after closePools() + // cleared the map. Read and written only under _manager_mutex. Defaults to + // true so direct native pool use (and auto-enable) works without an + // explicit enable_pooling() call; only disable_pooling() disarms it, and + // enable_pooling() re-arms it. + bool _accepting = true; + + // Throttle for the lazy-eviction sweep in acquireConnection(). The sweep + // iterates every pool (and every idle connection within each) under + // _manager_mutex, so running it on literally every connect is an O(pools × + // conns) contention hotspot for many-identity workloads. A pool cannot + // become evictable faster than the idle timeout, so sweeping more often + // than that buys nothing; we run it at most once per idle-timeout window. + std::chrono::steady_clock::time_point _last_sweep{}; + // Prevent copying ConnectionPoolManager(const ConnectionPoolManager&) = delete; ConnectionPoolManager& operator=(const ConnectionPoolManager&) = delete; diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp index 3cb00814..e83ae09c 100644 --- a/mssql_python/pybind/ddbc_bindings.cpp +++ b/mssql_python/pybind/ddbc_bindings.cpp @@ -5837,8 +5837,13 @@ SQLLEN SQLRowCount_wrap(SqlHandlePtr StatementHandle) { static std::once_flag pooling_init_flag; void enable_pooling(int maxSize, int idleTimeout) { + // configure() locks in the default max_size/idle_timeout for the process and + // must run exactly once, but re-arming must happen on every enable so a + // disable_pooling() -> enable_pooling() cycle resumes pooling (the + // std::call_once body would otherwise be skipped on the second enable). std::call_once(pooling_init_flag, [&]() { ConnectionPoolManager::getInstance().configure(maxSize, idleTimeout); }); + ConnectionPoolManager::getInstance().setAccepting(true); } // Thread-safe decimal separator setting @@ -5903,8 +5908,10 @@ PYBIND11_MODULE(ddbc_bindings, m) { "Internal: close the cursor without freeing the prepared statement"); py::class_(m, "Connection") - .def(py::init(), py::arg("conn_str"), - py::arg("use_pool"), py::arg("attrs_before") = py::dict()) + .def(py::init(), + py::arg("conn_str"), py::arg("use_pool"), py::arg("attrs_before") = py::dict(), + py::arg("pool_key") = std::u16string(), py::arg("token_factory") = py::none()) .def("close", &ConnectionHandle::close, "Close the connection") .def("commit", &ConnectionHandle::commit, "Commit the current transaction") .def("rollback", &ConnectionHandle::rollback, "Rollback the current transaction") @@ -5916,6 +5923,18 @@ PYBIND11_MODULE(ddbc_bindings, m) { .def("get_info", &ConnectionHandle::getInfo, py::arg("info_type")); m.def("enable_pooling", &enable_pooling, "Enable global connection pooling"); m.def("close_pooling", []() { ConnectionPoolManager::getInstance().closePools(); }); + m.def("disable_pooling", []() { + // Disarm new-pool creation *before* closing so a connect racing this + // disable cannot resurrect a pool after the map is cleared: any + // acquireConnection serialized after setAccepting(false) declines and + // falls back to a non-pooled connection. closePools() then reaps every + // pool created before the disarm. + auto& manager = ConnectionPoolManager::getInstance(); + manager.setAccepting(false); + manager.closePools(); + }, "Disable global connection pooling and close all pools"); + m.def("_pool_count", []() { return ConnectionPoolManager::getInstance().poolCount(); }, + "Diagnostic: number of connection pools currently tracked (used by tests)"); m.def("DDBCSQLExecDirect", &SQLExecDirect_wrap, "Execute a SQL query directly"); m.def("DDBCSQLExecute", &SQLExecute_wrap, "Prepare and execute T-SQL statements", py::arg("statementHandle"), py::arg("query"), py::arg("params"), py::arg("paramInfos"), diff --git a/tests/test_008_auth.py b/tests/test_008_auth.py index d82ecaea..3b1f439f 100644 --- a/tests/test_008_auth.py +++ b/tests/test_008_auth.py @@ -12,15 +12,19 @@ from mssql_python.auth import ( AADAuth, ServicePrincipalAuth, + TokenInfo, _parse_tenant_id, process_auth_parameters, remove_sensitive_params, - get_auth_token, + get_auth_token_info, + compute_identity_key, extract_auth_type, _credential_cache, _credential_cache_lock, + _account_id_cache, ) from mssql_python.constants import AuthType, ConstantsDDBC +from mssql_python.exceptions import InterfaceError import secrets SAMPLE_TOKEN = secrets.token_hex(44) @@ -37,11 +41,30 @@ class MockDefaultAzureCredential: def get_token(self, scope): return MockToken() + class MockAuthenticationRecord: + # Minimal stand-in for azure.identity.AuthenticationRecord: only the + # public home_account_id is consumed by the pooling key logic. + home_account_id = "mock-home-account-id" + class MockDeviceCodeCredential: + def __init__(self, **kwargs): + # Accept disable_automatic_authentication (and any future kwargs) + # the same way real azure-identity credentials do. + self._init_kwargs = kwargs + + def authenticate(self, **kwargs): + return MockAuthenticationRecord() + def get_token(self, scope): return MockToken() class MockInteractiveBrowserCredential: + def __init__(self, **kwargs): + self._init_kwargs = kwargs + + def authenticate(self, **kwargs): + return MockAuthenticationRecord() + def get_token(self, scope): return MockToken() @@ -82,12 +105,25 @@ def __init__(self, **kwargs): class MockClientAuthenticationError(Exception): pass + # Prefer the real AuthenticationRequiredError so production code exercises + # the exact exception type it will see in production; fall back to a local + # subclass only when azure-identity is not installed (the scenario this mock + # exists to cover). Either way the mocked module exposes the same class + # object that production imports, so silent-first's except-clause matches. + try: + from azure.identity import AuthenticationRequiredError as MockAuthenticationRequiredError + except ImportError: + + class MockAuthenticationRequiredError(MockClientAuthenticationError): + """Local stand-in when azure-identity is unavailable.""" + class MockIdentity: DefaultAzureCredential = MockDefaultAzureCredential DeviceCodeCredential = MockDeviceCodeCredential InteractiveBrowserCredential = MockInteractiveBrowserCredential ClientSecretCredential = MockClientSecretCredential ManagedIdentityCredential = MockManagedIdentityCredential + AuthenticationRequiredError = MockAuthenticationRequiredError class MockCore: class exceptions: @@ -124,10 +160,12 @@ class transport: @pytest.fixture(autouse=True) def clear_credential_cache(): - """Clear the module-level credential cache between tests.""" + """Clear the module-level credential and account-id caches between tests.""" _credential_cache.clear() + _account_id_cache.clear() yield _credential_cache.clear() + _account_id_cache.clear() class TestAuthType: @@ -240,7 +278,7 @@ def get_token(self, scope): # Test different exception types class MockCredentialWithTypeError: - def __init__(self): + def __init__(self, **kwargs): raise TypeError("Invalid argument type passed") azure_identity.DeviceCodeCredential = MockCredentialWithTypeError @@ -270,7 +308,7 @@ def test_get_token_general_exception_handling_token_error(self): # Create a credential that fails during get_token call class MockCredentialWithTokenError: - def __init__(self): + def __init__(self, **kwargs): pass # Successful initialization def get_token(self, scope): @@ -338,6 +376,224 @@ def __init__(self): azure_identity.DefaultAzureCredential = original_default +class TestSilentFirstInteractive: + """Regression tests for the silent-first pattern (interactive / device-code). + + Contract: ``_acquire_token`` must try ``get_token()`` silently first and + invoke ``authenticate()`` ONLY when the SDK raises + ``AuthenticationRequiredError`` — never on every acquisition. This guards + against reverting to the old "authenticate() on every checkout" behavior, + which would prompt the user (browser / device code) on each pooled reconnect. + """ + + def test_silent_success_never_calls_authenticate(self): + """When get_token() succeeds silently, authenticate() is never invoked.""" + az = sys.modules["azure.identity"] + original = az.InteractiveBrowserCredential + + class SilentCred: + def __init__(self, **kwargs): + self.authenticate_calls = 0 + + def authenticate(self, **kwargs): + self.authenticate_calls += 1 + + class R: + home_account_id = "acct" + + return R() + + def get_token(self, scope): + class T: + token = SAMPLE_TOKEN + + return T() + + try: + az.InteractiveBrowserCredential = SilentCred + AADAuth._acquire_token("interactive") + cred = _credential_cache["interactive"] + assert cred.authenticate_calls == 0, "authenticate() must not run on a silent success" + finally: + az.InteractiveBrowserCredential = original + + def test_authentication_required_triggers_authenticate_once(self): + """AuthenticationRequiredError from the silent get_token() falls back to + authenticate() exactly once, then retries get_token() successfully.""" + az = sys.modules["azure.identity"] + original = az.DeviceCodeCredential + auth_required = az.AuthenticationRequiredError + + class ChallengeCred: + def __init__(self, **kwargs): + self.authenticate_calls = 0 + self.get_token_calls = 0 + + def authenticate(self, **kwargs): + self.authenticate_calls += 1 + + class R: + home_account_id = "device-acct" + + return R() + + def get_token(self, scope): + self.get_token_calls += 1 + if self.get_token_calls == 1: + raise auth_required("interactive authentication required") + + class T: + token = SAMPLE_TOKEN + + return T() + + try: + az.DeviceCodeCredential = ChallengeCred + _, _, _, home_account_id = AADAuth._acquire_token("devicecode") + cred = _credential_cache["devicecode"] + # Exactly one interactive step, and get_token tried twice + # (silent attempt that raised, then the post-authenticate retry). + assert cred.authenticate_calls == 1 + assert cred.get_token_calls == 2 + assert home_account_id == "device-acct" + finally: + az.DeviceCodeCredential = original + + def test_subsequent_acquisitions_are_silent(self): + """After the first interactive authenticate(), later acquisitions for + the same cached credential acquire silently and never re-prompt.""" + az = sys.modules["azure.identity"] + original = az.InteractiveBrowserCredential + auth_required = az.AuthenticationRequiredError + + class OnceChallengeCred: + def __init__(self, **kwargs): + self.authenticate_calls = 0 + self.challenged = False + + def authenticate(self, **kwargs): + self.authenticate_calls += 1 + + class R: + home_account_id = "acct" + + return R() + + def get_token(self, scope): + if not self.challenged: + self.challenged = True + raise auth_required("need interactive") + + class T: + token = SAMPLE_TOKEN + + return T() + + try: + az.InteractiveBrowserCredential = OnceChallengeCred + AADAuth._acquire_token("interactive") # challenge -> authenticate() + AADAuth._acquire_token("interactive") # silent + AADAuth._acquire_token("interactive") # silent + cred = _credential_cache["interactive"] + assert cred.authenticate_calls == 1, "authenticate() must run only on the first prompt" + finally: + az.InteractiveBrowserCredential = original + + def test_home_account_id_is_cached_and_reused(self): + """The home_account_id resolved by the interactive step is cached in + _account_id_cache and surfaced on subsequent silent acquisitions.""" + az = sys.modules["azure.identity"] + original = az.InteractiveBrowserCredential + auth_required = az.AuthenticationRequiredError + + class OnceChallengeCred: + def __init__(self, **kwargs): + self.challenged = False + + def authenticate(self, **kwargs): + class R: + home_account_id = "stable-acct" + + return R() + + def get_token(self, scope): + if not self.challenged: + self.challenged = True + raise auth_required("need interactive") + + class T: + token = SAMPLE_TOKEN + + return T() + + try: + az.InteractiveBrowserCredential = OnceChallengeCred + _, _, _, hid1 = AADAuth._acquire_token("interactive") + # The cache now holds the account id for the interactive cache key. + assert "stable-acct" in _account_id_cache.values() + _, _, _, hid2 = AADAuth._acquire_token("interactive") + assert hid1 == "stable-acct" + assert hid2 == "stable-acct" + finally: + az.InteractiveBrowserCredential = original + + +class TestAuthenticateInteractive: + """Direct branch coverage for the _authenticate_interactive fallback helper.""" + + def test_returns_none_when_credential_has_no_authenticate(self): + """A credential/test double lacking authenticate() yields None so the + caller falls back to the token-hash pool key.""" + + class NoAuth: + pass + + assert AADAuth._authenticate_interactive(NoAuth()) is None + + def test_uses_scopes_keyword_when_supported(self): + captured = {} + + class Cred: + def authenticate(self, **kwargs): + captured.update(kwargs) + + class R: + home_account_id = "acct-1" + + return R() + + assert AADAuth._authenticate_interactive(Cred()) == "acct-1" + assert "scopes" in captured + + def test_falls_back_to_no_args_when_scopes_rejected(self): + """Older/mocked authenticate() signatures that reject the scopes + keyword are retried with no arguments.""" + calls = {"with_scopes": 0, "no_args": 0} + + class Cred: + def authenticate(self, *args, **kwargs): + if kwargs: + calls["with_scopes"] += 1 + raise TypeError("unexpected keyword argument 'scopes'") + calls["no_args"] += 1 + + class R: + home_account_id = "acct-2" + + return R() + + assert AADAuth._authenticate_interactive(Cred()) == "acct-2" + assert calls["with_scopes"] == 1 + assert calls["no_args"] == 1 + + def test_returns_none_when_record_lacks_home_account_id(self): + class Cred: + def authenticate(self, **kwargs): + return object() # no home_account_id attribute + + assert AADAuth._authenticate_interactive(Cred()) is None + + class TestProcessAuthParameters: def test_empty_parameters(self): auth_type = process_auth_parameters({}) @@ -686,7 +942,7 @@ class Token: azure_identity.DefaultAzureCredential = MockCredentialWithRefresh # First call — gets initial token - _, raw_token_1 = AADAuth._acquire_token("default") + _, raw_token_1, _, _ = AADAuth._acquire_token("default") assert raw_token_1 == "initial_token_abc123" assert call_count == 1 @@ -696,7 +952,7 @@ class Token: # Second call — same credential instance, but SDK returns refreshed token # (simulating post-expiry refresh) - _, raw_token_2 = AADAuth._acquire_token("default") + _, raw_token_2, _, _ = AADAuth._acquire_token("default") assert raw_token_2 == "refreshed_token_xyz789" assert call_count == 2 @@ -779,12 +1035,20 @@ def test_empty_authentication_value(self): class TestTokenFailureFallthrough: - """Verify that connect() succeeds without a token when credential creation fails.""" + """Token acquisition must fail *closed*: when a token-backed auth type is + detected but the credential cannot be created / the token cannot be + acquired, connect() must raise rather than silently opening a connection + with Authentication= stripped and no token attached.""" @patch("mssql_python.connection.ddbc_bindings.Connection") - def test_connect_proceeds_without_token_on_credential_failure(self, mock_ddbc_conn): - """When auth type is detected but token acquisition fails, - the connection should still be attempted (just without a token).""" + def test_connect_raises_on_credential_failure(self, mock_ddbc_conn): + """When auth type is detected but token acquisition fails, connect() + must surface a clear error instead of proceeding without a token. + + The underlying Azure failure (RuntimeError from AADAuth) is re-wrapped + at the DB-API boundary as InterfaceError so callers get a consistent, + catchable driver exception. + """ mock_ddbc_conn.return_value = MagicMock() import sys @@ -799,46 +1063,8 @@ def __init__(self): azure_identity.DefaultAzureCredential = CredentialThatAlwaysFails from mssql_python import connect - conn = connect("Server=test;Authentication=ActiveDirectoryDefault;Database=testdb") - assert conn._auth_type == "default" - # Token should not be in attrs_before since acquisition failed - assert ConstantsDDBC.SQL_COPT_SS_ACCESS_TOKEN.value not in conn._attrs_before - conn.close() - finally: - azure_identity.DefaultAzureCredential = original - - -class TestGetAuthTokenEdgeCases: - """Cover the Windows-interactive and token-failure branches.""" - - def test_no_auth_type_returns_none(self): - result = get_auth_token(None) - assert result is None - - def test_empty_auth_type_returns_none(self): - result = get_auth_token("") - assert result is None - - def test_windows_interactive_returns_none(self, monkeypatch): - monkeypatch.setattr(platform, "system", lambda: "Windows") - result = get_auth_token("interactive") - assert result is None - - def test_token_acquisition_failure_returns_none(self): - """When AADAuth.get_token raises, get_auth_token returns None.""" - import sys - - azure_identity = sys.modules["azure.identity"] - original = azure_identity.DefaultAzureCredential - - class FailingCredential: - def __init__(self): - raise RuntimeError("credential creation exploded") - - try: - azure_identity.DefaultAzureCredential = FailingCredential - result = get_auth_token("default") - assert result is None + with pytest.raises(InterfaceError): + connect("Server=test;Authentication=ActiveDirectoryDefault;Database=testdb") finally: azure_identity.DefaultAzureCredential = original @@ -1380,3 +1606,542 @@ def test_factory_cache_key_does_not_contain_raw_secret(self): ) for key in _credential_cache.keys(): assert secret_marker not in repr(key), f"Raw secret leaked into cache key: {key!r}" + + +class TestTokenExpiryCapture: + """Foundation for identity-aware pooling: the acquisition + path must surface the token expiry alongside the ODBC struct.""" + + def test_get_token_info_returns_tokeninfo(self): + info = AADAuth.get_token_info("default") + assert isinstance(info, TokenInfo) + assert isinstance(info.token_struct, bytes) + assert len(info.token_struct) > 4 + + def test_get_token_info_expiry_none_when_absent(self): + # The default test double's token object has no expires_on attribute, + # so capture must fall back to None rather than raising. + info = AADAuth.get_token_info("default") + assert info.expires_on is None + + def test_get_token_info_captures_expiry_when_present(self): + class MockTokenWithExpiry: + token = SAMPLE_TOKEN + expires_on = 1_900_000_000 + + class MockCredWithExpiry: + def get_token(self, scope): + return MockTokenWithExpiry() + + import sys + + azure_identity = sys.modules["azure.identity"] + original = azure_identity.DefaultAzureCredential + azure_identity.DefaultAzureCredential = MockCredWithExpiry + try: + info = AADAuth.get_token_info("default") + assert info.expires_on == 1_900_000_000 + finally: + azure_identity.DefaultAzureCredential = original + + def test_get_auth_token_info_returns_none_without_auth_type(self): + assert get_auth_token_info(None) is None + assert get_auth_token_info("") is None + + def test_get_auth_token_info_default(self): + info = get_auth_token_info("default") + assert isinstance(info, TokenInfo) + assert isinstance(info.token_struct, bytes) + + def test_get_auth_token_info_windows_interactive_returns_none(self): + # Windows Interactive auth is handled natively by the driver, so the + # Python layer must not acquire a token (auth.py line 478). + with patch("mssql_python.auth.platform.system", return_value="Windows"): + assert get_auth_token_info("interactive") is None + + def test_get_auth_token_info_raises_on_runtime_error(self): + # A token-acquisition RuntimeError must propagate (fail closed) so the + # caller never falls through to a token-less connect. + with patch( + "mssql_python.auth.AADAuth.get_token_info", + side_effect=RuntimeError("token endpoint unreachable"), + ): + with pytest.raises(RuntimeError, match="token endpoint unreachable"): + get_auth_token_info("default") + + def test_get_auth_token_info_raises_on_value_error(self): + # Same fail-closed contract for ValueError (unsupported auth type, etc.). + with patch( + "mssql_python.auth.AADAuth.get_token_info", + side_effect=ValueError("invalid credential kwargs"), + ): + with pytest.raises(ValueError, match="invalid credential kwargs"): + get_auth_token_info("default") + + +class TestDefaultCredentialPoolingWarning: + """The one-time DefaultAzureCredential pooling warning uses double-checked + locking so concurrent callers emit it at most once.""" + + def test_inner_double_check_returns_without_warning(self, monkeypatch): + # Simulate a race: the outer check sees the flag unset, but another + # thread sets it while we wait on the lock, so the inner check returns + # (auth.py line 96) and no warning is emitted by this caller. + import mssql_python.auth as auth_mod + + monkeypatch.setattr(auth_mod, "_dac_pooling_warned", False, raising=False) + mock_logger = MagicMock() + monkeypatch.setattr(auth_mod, "logger", mock_logger) + + class RacingLock: + def __enter__(self): + auth_mod._dac_pooling_warned = True + return self + + def __exit__(self, *exc): + return False + + monkeypatch.setattr(auth_mod, "_dac_pooling_warned_lock", RacingLock()) + + auth_mod._warn_default_credential_pooling() + + mock_logger.warning.assert_not_called() + + def test_warning_emitted_once_when_not_yet_warned(self, monkeypatch): + # Baseline: with the flag unset and a normal lock, the warning fires. + import threading + import mssql_python.auth as auth_mod + + monkeypatch.setattr(auth_mod, "_dac_pooling_warned", False, raising=False) + monkeypatch.setattr(auth_mod, "_dac_pooling_warned_lock", threading.Lock()) + mock_logger = MagicMock() + monkeypatch.setattr(auth_mod, "logger", mock_logger) + + auth_mod._warn_default_credential_pooling() + auth_mod._warn_default_credential_pooling() + + assert mock_logger.warning.call_count == 1 + + +class TestComputeIdentityKey: + """The per-identity discriminator that makes the native pool key + identity-aware.""" + + def test_none_when_no_auth_type(self): + # No token auth => pool key stays the bare connection string. + assert compute_identity_key(None) is None + assert compute_identity_key("") is None + + def test_msi_user_assigned_uses_client_id(self): + key = compute_identity_key("msi", {"client_id": "abc-123"}) + assert key == "msi:abc-123" + + def test_msi_system_assigned(self): + assert compute_identity_key("msi") == "msi:system" + assert compute_identity_key("msi", {}) == "msi:system" + + def test_msi_key_derived_without_token(self): + # No token_struct supplied, yet MSI still yields a key -> enables + # skipping token acquisition on a pool hit. + assert compute_identity_key("msi", {"client_id": "cid"}) == "msi:cid" + + def test_token_hash_fallback_for_default(self): + token = b"\x00\x01sometokenbytes" + key = compute_identity_key("default", token_struct=token) + assert key is not None and key.startswith("tok:") + assert len(key) == len("tok:") + 64 # sha256 hex + + def test_different_tokens_produce_different_keys(self): + k1 = compute_identity_key("default", token_struct=b"token-a") + k2 = compute_identity_key("default", token_struct=b"token-b") + assert k1 != k2 + + def test_same_token_produces_stable_key(self): + k1 = compute_identity_key("default", token_struct=b"same") + k2 = compute_identity_key("default", token_struct=b"same") + assert k1 == k2 + + def test_none_when_token_required_but_absent(self): + # default/interactive/devicecode need the token to form the key; + # without it, signal the caller to acquire then recompute. + assert compute_identity_key("default") is None + assert compute_identity_key("interactive") is None + assert compute_identity_key("devicecode") is None + + def test_interactive_uses_home_account_id(self): + # A stable signed-in account id keys the pool on the account, so a + # silent refresh reuses it while a different account gets its own pool. + assert compute_identity_key("interactive", home_account_id="acct-xyz") == "acct:acct-xyz" + + def test_devicecode_uses_home_account_id(self): + assert compute_identity_key("devicecode", home_account_id="dev-acct") == "acct:dev-acct" + + def test_account_id_preferred_over_token_hash(self): + # When both are available the account id wins (stable across refreshes). + key = compute_identity_key("interactive", token_struct=b"tok", home_account_id="acct-1") + assert key == "acct:acct-1" + + +class TestConnectionPoolKey: + """Part B: Connection.__init__ must build an identity-aware pool key and + forward it to the native Connection so distinct Entra identities never + share a pooled connection, while non-token auth keeps the legacy + bare-connection-string key (backward compatible).""" + + @patch("mssql_python.connection.ddbc_bindings.Connection") + def test_no_pool_key_for_non_token_auth(self, mock_ddbc_conn): + # Plain SQL auth (no Authentication=) => identity key None => pool key + # stays empty so the native layer keys on the connection string. + mock_ddbc_conn.return_value = MagicMock() + from mssql_python import connect + + conn = connect("Server=test;Database=testdb;UID=sa;PWD=secret") + assert conn._pool_key == "" + # Native Connection received the empty pool key as the 4th positional arg. + args, _ = mock_ddbc_conn.call_args + assert args[3] == "" + conn.close() + + @patch("mssql_python.connection.ddbc_bindings.Connection") + def test_system_assigned_msi_pool_key(self, mock_ddbc_conn): + mock_ddbc_conn.return_value = MagicMock() + from mssql_python import connect + + conn = connect("Server=test;Database=testdb;Authentication=ActiveDirectoryMSI") + assert conn._pool_key == conn.connection_str + "\x00msi:system" + args, _ = mock_ddbc_conn.call_args + assert args[3] == conn._pool_key + conn.close() + + @patch("mssql_python.connection.ddbc_bindings.Connection") + def test_user_assigned_msi_pool_key_uses_client_id(self, mock_ddbc_conn): + mock_ddbc_conn.return_value = MagicMock() + from mssql_python import connect + + client_id = "11111111-2222-3333-4444-555555555555" + conn = connect( + f"Server=test;Database=testdb;Authentication=ActiveDirectoryMSI;UID={client_id}" + ) + assert conn._pool_key == conn.connection_str + f"\x00msi:{client_id}" + conn.close() + + @patch("mssql_python.connection.get_auth_token_info") + @patch("mssql_python.connection.ddbc_bindings.Connection") + def test_default_auth_pool_key_hashes_token(self, mock_ddbc_conn, mock_get_token_info): + # When a token is present, the pool key must include its hash so that + # two different tokens (identities) never collide on the same pool. + mock_ddbc_conn.return_value = MagicMock() + mock_get_token_info.return_value = TokenInfo( + token_struct=b"\x00\x01fake-token-bytes", expires_on=None + ) + from mssql_python import connect + + conn = connect("Server=test;Database=testdb;Authentication=ActiveDirectoryDefault") + assert conn._pool_key.startswith(conn.connection_str + "\x00tok:") + args, _ = mock_ddbc_conn.call_args + assert args[3] == conn._pool_key + conn.close() + + @patch("mssql_python.connection.get_auth_token_info") + @patch("mssql_python.connection.ddbc_bindings.Connection") + def test_different_tokens_yield_different_pool_keys(self, mock_ddbc_conn, mock_get_token_info): + mock_ddbc_conn.return_value = MagicMock() + from mssql_python import connect + + cs = "Server=test;Database=testdb;Authentication=ActiveDirectoryDefault" + + mock_get_token_info.return_value = TokenInfo( + token_struct=b"identity-A-token", expires_on=None + ) + conn_a = connect(cs) + key_a = conn_a._pool_key + conn_a.close() + + mock_get_token_info.return_value = TokenInfo( + token_struct=b"identity-B-token", expires_on=None + ) + conn_b = connect(cs) + key_b = conn_b._pool_key + conn_b.close() + + # Same connection string, different identity => different pool key. + assert key_a != key_b + assert conn_a.connection_str == conn_b.connection_str + + @patch("mssql_python.connection.get_auth_token_info") + @patch("mssql_python.connection.ddbc_bindings.Connection") + def test_default_auth_no_token_fails_closed(self, mock_ddbc_conn, mock_get_token_info): + # Token acquisition yielded nothing usable: there is no identity to key + # on and no token to attach, so connect() must fail closed rather than + # open a token-less connection on the bare connection-string key. + mock_ddbc_conn.return_value = MagicMock() + mock_get_token_info.return_value = None + from mssql_python import connect + + with pytest.raises(InterfaceError): + connect("Server=test;Database=testdb;Authentication=ActiveDirectoryDefault") + + +class TestTokenFactoryLazyAcquisition: + """Part C: when the pool key is known *without* a token (MSI, whose + identity is the client id), token acquisition is deferred to a callback the + native layer invokes only when it opens a physical connection. A same- + identity pool hit therefore never pays for a token. This is an internal + ``token_factory`` (5th positional arg to native Connection) and is distinct + from the public ``token_provider=`` credential API.""" + + _TOKEN_ATTR = ConstantsDDBC.SQL_COPT_SS_ACCESS_TOKEN.value + + @patch("mssql_python.connection.get_auth_token_info") + @patch("mssql_python.connection.ddbc_bindings.Connection") + def test_msi_defers_token_acquisition(self, mock_ddbc_conn, mock_get_token_info): + # MSI identity is known from the client id, so the token must NOT be + # acquired during __init__ — it is deferred to the token_factory. + mock_ddbc_conn.return_value = MagicMock() + from mssql_python import connect + + conn = connect("Server=test;Database=testdb;Authentication=ActiveDirectoryMSI") + try: + mock_get_token_info.assert_not_called() + # token_factory is passed as the 5th positional arg and is callable. + args, _ = mock_ddbc_conn.call_args + assert callable(args[4]) + # The token is not eagerly placed into attrs_before either. + assert self._TOKEN_ATTR not in conn._attrs_before + finally: + conn.close() + + @patch("mssql_python.connection.get_auth_token_info") + @patch("mssql_python.connection.ddbc_bindings.Connection") + def test_token_factory_materializes_token_on_call(self, mock_ddbc_conn, mock_get_token_info): + # Invoking the factory (as native does on a pool miss) acquires the + # token and returns ``(attrs, expires_on)`` with the access token merged + # into attrs. + mock_ddbc_conn.return_value = MagicMock() + mock_get_token_info.return_value = TokenInfo( + token_struct=b"materialized-token-bytes", expires_on=1893456000 + ) + from mssql_python import connect + + conn = connect("Server=test;Database=testdb;Authentication=ActiveDirectoryMSI") + try: + attrs, expires_on = conn._token_factory() + mock_get_token_info.assert_called_once() + assert attrs[self._TOKEN_ATTR] == b"materialized-token-bytes" + assert expires_on == 1893456000 + finally: + conn.close() + + @patch("mssql_python.connection.get_auth_token_info") + @patch("mssql_python.connection.ddbc_bindings.Connection") + def test_token_dependent_auth_has_no_factory(self, mock_ddbc_conn, mock_get_token_info): + # DefaultAzureCredential keys on the token hash, so the token must be + # acquired eagerly (to compute the key) and there is no deferred factory. + mock_ddbc_conn.return_value = MagicMock() + mock_get_token_info.return_value = TokenInfo( + token_struct=b"eager-token-bytes", expires_on=None + ) + from mssql_python import connect + + conn = connect("Server=test;Database=testdb;Authentication=ActiveDirectoryDefault") + try: + mock_get_token_info.assert_called_once() + assert conn._token_factory is None + args, _ = mock_ddbc_conn.call_args + assert args[4] is None + assert conn._attrs_before[self._TOKEN_ATTR] == b"eager-token-bytes" + finally: + conn.close() + + @patch("mssql_python.connection.ddbc_bindings.Connection") + def test_non_token_auth_has_no_factory(self, mock_ddbc_conn): + # Plain SQL auth: no identity, no deferred token. + mock_ddbc_conn.return_value = MagicMock() + from mssql_python import connect + + conn = connect("Server=test;Database=testdb;UID=sa;PWD=secret") + try: + assert conn._token_factory is None + args, _ = mock_ddbc_conn.call_args + assert args[4] is None + finally: + conn.close() + + @patch("mssql_python.connection.get_auth_token_info") + @patch("mssql_python.connection.ddbc_bindings.Connection") + def test_account_keyed_auth_defers_to_factory(self, mock_ddbc_conn, mock_get_token_info): + # Device-code yields a stable home_account_id, so the pool key is + # ``acct:`` and token acquisition is deferred to the factory rather + # than pinned into attrs_before. + mock_ddbc_conn.return_value = MagicMock() + mock_get_token_info.return_value = TokenInfo( + token_struct=b"dc-token-bytes", + expires_on=1893456000, + home_account_id="dc-acct-1", + ) + from mssql_python import connect + + conn = connect("Server=test;Database=testdb;Authentication=ActiveDirectoryDeviceCode") + try: + assert conn._pool_key.endswith("\x00acct:dc-acct-1") + assert conn._token_factory is not None + assert self._TOKEN_ATTR not in conn._attrs_before + args, _ = mock_ddbc_conn.call_args + assert callable(args[4]) + finally: + conn.close() + + @patch("mssql_python.connection.get_auth_token_info") + @patch("mssql_python.connection.ddbc_bindings.Connection") + def test_factory_raises_when_token_unavailable(self, mock_ddbc_conn, mock_get_token_info): + # If the deferred token acquisition yields nothing, the factory must + # fail closed (raise) so native never opens a pooled connection with + # Authentication= stripped and no token attached. + mock_ddbc_conn.return_value = MagicMock() + mock_get_token_info.return_value = TokenInfo( + token_struct=b"dc-token-bytes", + expires_on=1893456000, + home_account_id="dc-acct-2", + ) + from mssql_python import connect + + conn = connect("Server=test;Database=testdb;Authentication=ActiveDirectoryDeviceCode") + try: + # Simulate the provider returning nothing on the deferred call. + mock_get_token_info.return_value = None + with pytest.raises(InterfaceError): + conn._token_factory() + finally: + conn.close() + + +class TestRawAccessTokenPoolKey: + """A raw SQL_COPT_SS_ACCESS_TOKEN in attrs_before (no ``Authentication=`` + keyword — the documented msodbcsql raw-token / token_provider pattern) must + still isolate the pool per token, or two callers with different tokens + against the same server would share a pooled connection.""" + + _TOKEN_ATTR = ConstantsDDBC.SQL_COPT_SS_ACCESS_TOKEN.value + + @pytest.mark.parametrize( + "bad_token", + [ + "a-str-token", # the exact str bypass called out in review + 12345, # int + memoryview(b"\x04\x00\x00\x00tok"), # memoryview is not bytes/bytearray + ["not", "bytes"], # list + ], + ids=["str", "int", "memoryview", "list"], + ) + @patch("mssql_python.connection.ddbc_bindings.Connection") + def test_non_binary_access_token_fails_closed(self, mock_ddbc_conn, bad_token): + """A non-bytes access token is rejected up front with InterfaceError. + + Adversarial case: a caller supplies attr 1256 as a ``str`` (or any + non-binary type). The Python identity-key logic only hashes + bytes/bytearray, so without this guard a str token would fall through + with the bare connStr pool key and two callers with different str tokens + against the same server could share a pooled connection. The guard fails + closed *before* any physical connect, so nothing is ever pooled under a + shared key. The native ``setAttribute`` enforces the same rule. + """ + from mssql_python import connect + + with pytest.raises(InterfaceError, match="SQL_COPT_SS_ACCESS_TOKEN"): + connect( + "Server=test;Database=testdb;UID=sa;PWD=secret", + attrs_before={self._TOKEN_ATTR: bad_token}, + ) + # Fail closed: the guard raises before the native Connection is built. + mock_ddbc_conn.assert_not_called() + + @patch("mssql_python.connection.ddbc_bindings.Connection") + def test_raw_token_sets_token_hash_pool_key(self, mock_ddbc_conn): + import hashlib + + mock_ddbc_conn.return_value = MagicMock() + from mssql_python import connect + + token = b"raw-access-token-bytes" + conn = connect( + "Server=test;Database=testdb;UID=sa;PWD=secret", + attrs_before={self._TOKEN_ATTR: token}, + ) + try: + expected = "tok:" + hashlib.sha256(token).hexdigest() + assert conn._pool_key.endswith("\x00" + expected) + assert conn._pool_key != "" + finally: + conn.close() + + @patch("mssql_python.connection.ddbc_bindings.Connection") + def test_different_raw_tokens_get_distinct_pool_keys(self, mock_ddbc_conn): + mock_ddbc_conn.return_value = MagicMock() + from mssql_python import connect + + conn_a = connect( + "Server=test;Database=testdb;UID=sa;PWD=secret", + attrs_before={self._TOKEN_ATTR: b"token-A"}, + ) + conn_b = connect( + "Server=test;Database=testdb;UID=sa;PWD=secret", + attrs_before={self._TOKEN_ATTR: b"token-B"}, + ) + try: + assert conn_a._pool_key != conn_b._pool_key + finally: + conn_a.close() + conn_b.close() + + @patch("mssql_python.connection.ddbc_bindings.Connection") + def test_same_raw_token_gets_stable_pool_key(self, mock_ddbc_conn): + mock_ddbc_conn.return_value = MagicMock() + from mssql_python import connect + + conn_a = connect( + "Server=test;Database=testdb;UID=sa;PWD=secret", + attrs_before={self._TOKEN_ATTR: b"same-token"}, + ) + conn_b = connect( + "Server=test;Database=testdb;UID=sa;PWD=secret", + attrs_before={self._TOKEN_ATTR: b"same-token"}, + ) + try: + assert conn_a._pool_key == conn_b._pool_key + finally: + conn_a.close() + conn_b.close() + + @patch("mssql_python.connection.ddbc_bindings.Connection") + def test_bytearray_raw_token_is_handled(self, mock_ddbc_conn): + import hashlib + + mock_ddbc_conn.return_value = MagicMock() + from mssql_python import connect + + token = bytearray(b"bytearray-token") + conn = connect( + "Server=test;Database=testdb;UID=sa;PWD=secret", + attrs_before={self._TOKEN_ATTR: token}, + ) + try: + expected = "tok:" + hashlib.sha256(bytes(token)).hexdigest() + assert conn._pool_key.endswith("\x00" + expected) + finally: + conn.close() + + @patch("mssql_python.connection.ddbc_bindings.Connection") + def test_no_token_leaves_pool_key_empty(self, mock_ddbc_conn): + # Plain SQL auth with no token: pool key stays empty (legacy connStr + # keying), so the raw-token guard does not disturb existing behavior. + mock_ddbc_conn.return_value = MagicMock() + from mssql_python import connect + + conn = connect( + "Server=test;Database=testdb;UID=sa;PWD=secret", + attrs_before={1256000: 30}, # some unrelated non-token attribute + ) + try: + assert conn._pool_key == "" + finally: + conn.close() diff --git a/tests/test_009_pooling.py b/tests/test_009_pooling.py index 4d790f2c..b32483f6 100644 --- a/tests/test_009_pooling.py +++ b/tests/test_009_pooling.py @@ -407,6 +407,234 @@ def session_identity(conn): ) +def test_idle_identity_pool_is_evicted_by_later_acquire(conn_str): + """Regression test for lazy eviction of idle identity pools. + + David's ask verbatim: "connect as A, wait past idle timeout, connect as B, + assert A's pool is gone (today it won't be)." + + Distinct connection strings (differing APP name) key to distinct pools — + the same mechanism identity-aware pooling uses (connStr + identity). The old + ``canEvict()`` required ``_current_size == 0 && _pool.empty()``, so a pool + holding a single *idle* connection (``_current_size == 1``) was never + reclaimed: connecting as B would leave A's pool alive forever. The fix makes + ``canEvict()`` evaluate the idle timeout, so acquiring B sweeps the + aged-out pool A away. + + Run in a subprocess so this test's ``pooling(idle_timeout=1)`` is the first + (and therefore effective) call in a clean process, and so ``_pool_count`` + observes only pools this test created. + """ + _run_in_subprocess( + """ + import os, re, sys, time + from mssql_python import connect, pooling, ddbc_bindings + + base = os.environ["DB_CONNECTION_STRING"] + # Two distinct pool keys against the same server. The pool key is derived + # from the processed connection string, so switching the target database + # (master vs tempdb -- both always present) yields two separate pools, + # the same way two different identities would. + if re.search(r"(?i)database=", base): + conn_a = re.sub(r"(?i)database=[^;]*", "Database=master", base) + conn_b = re.sub(r"(?i)database=[^;]*", "Database=tempdb", base) + else: + trimmed = base.rstrip(";") + conn_a = trimmed + ";Database=master" + conn_b = trimmed + ";Database=tempdb" + + pooling(max_size=2, idle_timeout=1) + + # Identity A: open and return -> pool A now holds one idle connection + # (_current_size == 1). Under the old canEvict() this pool could never + # be reclaimed. + a = connect(conn_a) + a.cursor().execute("SELECT 1") + a.close() + assert ddbc_bindings._pool_count() == 1, ( + f"expected exactly pool A, got {ddbc_bindings._pool_count()}" + ) + + # Let pool A's idle connection age past the 1s idle timeout. + time.sleep(3) + + # Identity B: acquiring on a different key triggers the manager's lazy + # eviction sweep, which must now reclaim the aged-out pool A. + b = connect(conn_b) + b.cursor().execute("SELECT 1") + b.close() + + count = ddbc_bindings._pool_count() + assert count == 1, ( + f"pool A was not evicted: _pool_count()={count} (expected 1, " + f"which is only pool B). Under the old canEvict() this would be 2." + ) + """, + conn_str, + ) + + +def test_pool_full_raises_when_max_size_reached(conn_str): + """Acquiring past ``max_size`` on a busy pool raises rather than blocking. + + Covers the "pool size limit reached" throw in ``ConnectionPool::acquire``: + with ``max_size == 1`` and the single slot already checked out, a second + acquire on the same key finds the pool empty and no free capacity, so it + raises immediately (the pool never blocks waiting for a return). + + Run in a subprocess so ``pooling(max_size=1, ...)`` is the first (and thus + effective) call in a clean process — the C++ pool config is locked in via + ``std::call_once`` for the lifetime of a process. + """ + _run_in_subprocess( + """ + import os, sys + from mssql_python import connect, pooling + + conn_str = os.environ["DB_CONNECTION_STRING"] + pooling(max_size=1, idle_timeout=30) + + # Fill the single slot and keep it checked out. + held = connect(conn_str) + held.cursor().execute("SELECT 1") + + # A second acquire on the same pool key has no free capacity and no + # idle connection to hand back -> it must raise, not hang. + try: + second = connect(conn_str) + except Exception as exc: + assert "pool size limit reached" in str(exc).lower() or "pool" in str(exc).lower(), ( + f"unexpected error for a full pool: {exc!r}" + ) + else: + second.close() + held.close() + raise AssertionError("expected a full-pool error, but the acquire succeeded") + + held.close() + """, + conn_str, + ) + + +def test_checked_out_pool_is_not_evicted(conn_str): + """A pool with a checked-out connection is never swept, even when aged out. + + Covers the "checked out" guard in ``ConnectionPool::canEvict``: reserved + capacity beyond what is sitting idle in the pool means a caller still holds a + connection, so the eviction sweep must leave that pool alone regardless of + the idle timeout. + + Run in a subprocess so ``pooling(idle_timeout=1)`` is the effective config + and ``_pool_count`` observes only this test's pools. + """ + _run_in_subprocess( + """ + import os, re, sys, time + from mssql_python import connect, pooling, ddbc_bindings + + base = os.environ["DB_CONNECTION_STRING"] + # Two distinct pool keys against the same server (see the idle-eviction + # test): switching the target database yields two separate pools. + if re.search(r"(?i)database=", base): + conn_a = re.sub(r"(?i)database=[^;]*", "Database=master", base) + conn_b = re.sub(r"(?i)database=[^;]*", "Database=tempdb", base) + else: + trimmed = base.rstrip(";") + conn_a = trimmed + ";Database=master" + conn_b = trimmed + ";Database=tempdb" + + pooling(max_size=2, idle_timeout=1) + + # Identity A: open and KEEP OPEN. The connection is checked out, so pool + # A's reserved capacity exceeds its idle contents (_current_size == 1, + # _pool empty) — canEvict() must refuse to reclaim it. + a = connect(conn_a) + a.cursor().execute("SELECT 1") + + # Age past the 1s idle timeout, then acquire on a different key to run + # the manager's eviction sweep over pool A. + time.sleep(3) + b = connect(conn_b) + b.cursor().execute("SELECT 1") + + count = ddbc_bindings._pool_count() + assert count == 2, ( + f"checked-out pool A must survive the sweep: _pool_count()={count} " + f"(expected 2: pools A and B)." + ) + + a.close() + b.close() + """, + conn_str, + ) + + +def test_disable_during_concurrent_connects_leaves_no_pool(conn_str): + """Disabling pooling amid concurrent connects leaves no escaped pool. + + Regression test for the disable()-vs-connect() race. A connect reads the + (unlocked) enabled flag; before the fix, one that read it as True could + create a native pool just after ``disable()`` cleared the map, leaving a + pool alive for the life of the process. The native manager now disarms + new-pool creation under ``_manager_mutex`` before closing, so a raced + connect either had its pool reaped by the close or is declined and falls + back to a non-pooled connection. Once the storm settles, ``_pool_count()`` + must be 0. + + Run in a subprocess so the enable -> concurrent connect/disable lifecycle is + isolated from other tests (pooling config is process-global). + """ + _run_in_subprocess( + """ + import os, sys, threading, time + from mssql_python import connect, pooling, ddbc_bindings + from mssql_python.pooling import PoolingManager + + conn_str = os.environ["DB_CONNECTION_STRING"] + pooling(max_size=8, idle_timeout=30) + + stop = threading.Event() + errors = [] + + def worker(): + while not stop.is_set(): + try: + c = connect(conn_str) + c.cursor().execute("SELECT 1") + c.close() + except Exception as exc: + errors.append(repr(exc)) + return + + threads = [threading.Thread(target=worker) for _ in range(4)] + for t in threads: + t.start() + + # Let pools populate, then disable while the workers are hammering so + # some connects race the disable. + time.sleep(0.2) + PoolingManager.disable() + + # Keep the workers running after the disable so late connects exercise + # the disarmed path (they must fall back to non-pooled, creating no pool). + time.sleep(0.3) + stop.set() + for t in threads: + t.join(timeout=10) + + assert not errors, f"workers hit connection errors: {errors}" + count = ddbc_bindings._pool_count() + assert count == 0, ( + f"a pool escaped disable(): _pool_count()={count} (expected 0). " + f"A connect that raced the disable resurrected a pool." + ) + """, + conn_str, + ) + + # ============================================================================= # Error Handling and Recovery Tests # ============================================================================= @@ -668,6 +896,34 @@ def test_pooling_enable_disable_cycle(conn_str): print("All enable/disable cycles completed successfully") +def test_reenable_rearms_disable_guard(conn_str): + """After enable -> disable -> enable, a second disable() must actually + disarm the native manager (call disable_pooling) rather than being skipped. + + Regression for the state-sync bug where enable() did not reset + _pools_closed, so the second disable() saw _pools_closed=True and skipped + ddbc_bindings.disable_pooling(), leaving the native manager accepting while + the Python side believed pooling was off. + """ + from unittest.mock import patch + + PoolingManager._reset_for_testing() + try: + PoolingManager.enable() + PoolingManager.disable() + # Re-enable: this must re-arm the guard (reset _pools_closed to False). + PoolingManager.enable() + assert PoolingManager._pools_closed is False + + with patch("mssql_python.pooling.ddbc_bindings.disable_pooling") as mock_disable: + PoolingManager.disable() + mock_disable.assert_called_once() + assert PoolingManager.is_enabled() is False + assert PoolingManager._pools_closed is True + finally: + PoolingManager._reset_for_testing() + + def test_pooling_state_consistency(conn_str): """Test that pooling state remains consistent across operations.""" print("Testing pooling state consistency...") @@ -693,3 +949,433 @@ def test_pooling_state_consistency(conn_str): assert PoolingManager.is_initialized(), "Should remain initialized after disable call" print("Pooling state consistency verified") + + +# ============================================================================= +# Native token-factory (lazy token acquisition) integration tests +# ============================================================================= +# +# These white-box tests drive the native ``ddbc_bindings.Connection`` +# constructor directly with a ``token_factory`` callable. The regular Python +# unit tests mock the native module, so they never execute the C++ lazy-token +# branches. Exercising them requires a live server, so these tests are guarded +# on ``conn_str`` and only run in the integration environment (where they also +# provide C++ line coverage for the token-factory paths). +# +# For a normal SQL-auth connection the credentials live in the connection +# string, so a factory that returns an empty attrs dict is sufficient to open +# a real connection while still forcing the C++ code down the ``token_factory`` +# branch. + + +class TestNativeTokenFactory: + """Integration tests for the native lazy token-factory path.""" + + @pytest.fixture(autouse=True) + def _drain_native_pool(self): + """Drain the native connection pool after each test. + + These tests exercise ``ddbc_bindings.Connection`` directly with + ``use_pool=True``, bypassing the high-level ``mssql_python.connect()`` + path that auto-enables :class:`PoolingManager`. Because pooling is never + enabled here, the module's ``atexit`` drain (``shutdown_pooling`` in + ``pooling.py``, guarded by ``PoolingManager._enabled``) does not run, so + pooled connections would otherwise linger in the process-lifetime native + pool singleton. Its C++ static destructor runs after the interpreter is + finalized, where touching the GIL/ODBC blocks and the process hangs on + exit. Draining here (the same call the product makes at exit) keeps the + native pool empty so the interpreter can shut down cleanly. + """ + from mssql_python import ddbc_bindings + + # These tests drive ``ddbc_bindings.Connection`` with ``use_pool=True`` + # directly, so the native manager must be armed to accept new pools. An + # earlier test may have called ``pooling(enabled=False)``, which now + # disarms the native manager (the disable-vs-connect race fix); re-arm + # here so pooled reuse works. ``enable_pooling`` re-arms accepting + # without re-running the one-time size/idle-timeout configuration. + ddbc_bindings.enable_pooling(100, 600) + yield + ddbc_bindings.close_pooling() + + def test_pooled_factory_invoked_on_miss_and_skipped_on_reuse(self, conn_str): + """The factory runs on a pool miss but not on a same-key pool reuse. + + Covers connection_pool.cpp (token_factory() invocation on the pool-miss + connect path) and the pooled acquireConnection branch in connection.cpp. + """ + if not conn_str: + pytest.skip("Live database connection required") + + from mssql_python import ddbc_bindings + + calls = {"n": 0} + + def factory(): + calls["n"] += 1 + return {} # SQL-auth creds are in conn_str; no attrs needed + + # Unique pool key so this test never collides with other pools. + pool_key = conn_str + "\x00mssql_test_659_pooled" + + # First open -> pool miss -> factory materializes the (empty) attrs. + conn1 = ddbc_bindings.Connection(conn_str, True, {}, pool_key, factory) + assert calls["n"] == 1, "Factory should be invoked once on a pool miss" + conn1.close() # returns the connection to the pool under pool_key + + # Second open with the same key -> pool hit -> factory is NOT called. + conn2 = ddbc_bindings.Connection(conn_str, True, {}, pool_key, factory) + assert calls["n"] == 1, "Factory must be skipped on a same-key pool reuse" + conn2.close() + + def test_non_pooled_factory_invoked(self, conn_str): + """A non-pooled connection still honors the factory. + + Covers the non-pool token_factory branch in connection.cpp. + """ + if not conn_str: + pytest.skip("Live database connection required") + + from mssql_python import ddbc_bindings + + calls = {"n": 0} + + def factory(): + calls["n"] += 1 + return {} + + conn = ddbc_bindings.Connection(conn_str, False, {}, "", factory) + assert calls["n"] == 1, "Factory should be invoked for a non-pooled connect" + conn.close() + + def test_non_pooled_without_factory_uses_attrs_before(self, conn_str): + """A non-pooled connection with no factory connects via attrs_before. + + Covers the non-pool ``else`` (no factory) branch in connection.cpp. + """ + if not conn_str: + pytest.skip("Live database connection required") + + from mssql_python import ddbc_bindings + + conn = ddbc_bindings.Connection(conn_str, False, {}, "", None) + assert conn is not None + conn.close() + + def test_distinct_pool_keys_do_not_share_connections(self, conn_str): + """Different identity keys must never reuse each other's pooled connection. + + This is the cross-identity isolation guarantee: the native + pool is keyed on the (identity-aware) pool key, so a connection opened + under identity A's key must not be handed out to identity B. We prove it + without real Entra tokens by counting factory invocations: a reuse skips + the factory, a miss calls it. If B were wrongly served A's pooled + connection, B's factory would never run. + + It also exercises the embedded-NUL separator (``\\x00``) that joins the + connection string and the identity discriminator: both keys contain a + NUL yet remain distinct, confirming the separator survives the + Python ``str`` -> pybind11 ``std::u16string`` conversion and that the + full key (not a NUL-truncated prefix) is used for pool lookup. + """ + if not conn_str: + pytest.skip("Live database connection required") + + from mssql_python import ddbc_bindings + + calls = {"a": 0, "b": 0} + + def factory_a(): + calls["a"] += 1 + return {} + + def factory_b(): + calls["b"] += 1 + return {} + + key_a = conn_str + "\x00identityA" + key_b = conn_str + "\x00identityB" + + # Open + close under identity A: pool miss -> factory A runs once, then + # the connection is returned to pool A. + conn_a1 = ddbc_bindings.Connection(conn_str, True, {}, key_a, factory_a) + assert calls["a"] == 1 + conn_a1.close() + + # Open under identity B (different key): must be a miss, not a reuse of + # A's pooled connection -> factory B runs. + conn_b1 = ddbc_bindings.Connection(conn_str, True, {}, key_b, factory_b) + assert ( + calls["b"] == 1 + ), "distinct identity key must not reuse another identity's pooled connection" + conn_b1.close() + + # Re-open under identity A: its own pooled connection is still there -> + # pool hit -> factory A is NOT called again. + conn_a2 = ddbc_bindings.Connection(conn_str, True, {}, key_a, factory_a) + assert calls["a"] == 1, "same identity key should reuse its own pooled connection" + conn_a2.close() + + def test_pooled_factory_accepts_tuple_return(self, conn_str): + """The factory may return ``(attrs, expires_on)``. + + Covers the tuple-unpacking path in Connection::invokeTokenFactory and + confirms a token whose expiry is far in the future is still reused on a + same-key pool hit (i.e. expiry-aware checkout does not discard it). + """ + if not conn_str: + pytest.skip("Live database connection required") + + import time + + from mssql_python import ddbc_bindings + + calls = {"n": 0} + far_future = int(time.time()) + 3600 # well beyond the 300s threshold + + def factory(): + calls["n"] += 1 + return {}, far_future # (attrs, expires_on) + + pool_key = conn_str + "\x00mssql_test_tuple_return" + + conn1 = ddbc_bindings.Connection(conn_str, True, {}, pool_key, factory) + assert calls["n"] == 1, "Factory (tuple return) should run once on a pool miss" + conn1.close() + + conn2 = ddbc_bindings.Connection(conn_str, True, {}, pool_key, factory) + assert calls["n"] == 1, "Non-expiring token must be reused on a same-key pool hit" + conn2.close() + + def test_near_expiry_token_refreshed_on_checkout(self, conn_str): + """A pooled connection with a near-expiry token is refreshed on checkout. + + Covers the expiry-aware checkout branch: when the + factory reports an expiry within the refresh threshold, the pooled + candidate is discarded and a fresh connection is opened, so the factory + runs again on reuse instead of being skipped. + """ + if not conn_str: + pytest.skip("Live database connection required") + + import time + + from mssql_python import ddbc_bindings + + calls = {"n": 0} + # Expiry inside the 300s refresh threshold => treated as near-expiry. + near_expiry = int(time.time()) + 60 + + def factory(): + calls["n"] += 1 + return {}, near_expiry + + pool_key = conn_str + "\x00mssql_test_651_nearexpiry" + + conn1 = ddbc_bindings.Connection(conn_str, True, {}, pool_key, factory) + assert calls["n"] == 1, "Factory should run once on the initial pool miss" + conn1.close() # returned to the pool, but its token is near expiry + + # Same key, but the pooled connection's token is near expiry, so it must + # be discarded and reopened -> factory runs a second time. + conn2 = ddbc_bindings.Connection(conn_str, True, {}, pool_key, factory) + assert calls["n"] == 2, "Near-expiry pooled connection must be refreshed on checkout" + conn2.close() + + def test_near_expiry_factory_token_bytes_are_extracted(self, conn_str): + """A refreshed factory that returns token bytes drives token extraction. + + Covers the ``extractAccessToken`` loop body in connection_pool.cpp: the + earlier near-expiry tests hand back an empty attrs dict, so the loop that + scans the dict for the ``SQL_COPT_SS_ACCESS_TOKEN`` (1256) entry never + iterates. Here the refresh call returns a dict *with* that token key, so + the C++ extracts the bytes and takes the token-rotation branch (fresh + token != the pooled connection's empty token). The reopen with a bogus + token against a SQL-auth server is expected to fail; we only need the + extraction + rotation branch to execute, which it does before the reopen. + """ + if not conn_str: + pytest.skip("Live database connection required") + + import time + + from mssql_python import ddbc_bindings + + # ConstantsDDBC.SQL_COPT_SS_ACCESS_TOKEN — the attr id the C++ scans for. + SQL_COPT_SS_ACCESS_TOKEN = 1256 + calls = {"n": 0} + near_expiry = int(time.time()) + 60 # inside the 300s refresh threshold + + def factory(): + calls["n"] += 1 + if calls["n"] == 1: + # Initial pool miss: SQL-auth creds are in conn_str, so an empty + # attrs dict opens a real connection (its token stays empty). + return {}, near_expiry + # Refresh on checkout: hand back a token so extractAccessToken() + # iterates its loop body and pulls the bytes out. + return {SQL_COPT_SS_ACCESS_TOKEN: b"rotated-token-bytes"}, near_expiry + + pool_key = conn_str + "\x00mssql_test_token_bytes" + + conn1 = ddbc_bindings.Connection(conn_str, True, {}, pool_key, factory) + assert calls["n"] == 1, "Factory should run once on the initial pool miss" + conn1.close() # returned to the pool with an empty token, near expiry + + # Same key: the pooled token is near expiry, so the factory is invoked + # again and now returns real token bytes. The C++ extracts them, sees a + # rotated (different) token, and reopens with it — which the SQL-auth + # server rejects. The extraction + rotation branch has already run. + with pytest.raises(Exception): + ddbc_bindings.Connection(conn_str, True, {}, pool_key, factory) + assert calls["n"] == 2, "Near-expiry checkout must re-invoke the factory" + + def test_factory_token_unknown_expiry_is_reused(self, conn_str): + """A factory reporting no expiry (``None``) or epoch ``0`` leaves the + pooled connection with an *unknown* expiry, which the native checkout + treats as "not near expiry" and therefore reuses. + + This pins down the current epoch-0 semantics (``isTokenNearExpiry`` + returns ``false`` for a 0/unknown expiry, connection.cpp): the factory + runs once on the miss and is skipped on the same-key hit. Treating an + unknown expiry as *unsafe* for token-backed factory pools is tracked as + an optional hardening follow-up, not current behavior. + """ + if not conn_str: + pytest.skip("Live database connection required") + + from mssql_python import ddbc_bindings + + for expiry in (None, 0): + calls = {"n": 0} + + def factory(_expiry=expiry): + calls["n"] += 1 + return {}, _expiry + + pool_key = conn_str + f"\x00mssql_test_unknown_expiry_{expiry}" + + conn1 = ddbc_bindings.Connection(conn_str, True, {}, pool_key, factory) + assert calls["n"] == 1, "Factory should run once on the initial pool miss" + conn1.close() + + conn2 = ddbc_bindings.Connection(conn_str, True, {}, pool_key, factory) + assert calls["n"] == 1, "Unknown/zero expiry is currently treated as safe and reused" + conn2.close() + + def test_factory_raises_on_near_expiry_checkout_recovers_pool(self, conn_str): + """A factory that raises while refreshing a near-expiry pooled connection + must not leak the reserved pool slot. + + Covers the exception path of the expiry-aware checkout: prime the pool + with a near-expiry connection, then make the refresh factory raise on + the next checkout. The connect fails (exception propagates), but the + reserved slot is released under the pool mutex, so a subsequent connect + under the same key still succeeds instead of finding the pool wedged or + exhausted. + """ + if not conn_str: + pytest.skip("Live database connection required") + + import time + + from mssql_python import ddbc_bindings + + near_expiry = int(time.time()) + 60 # inside the 300s refresh threshold + state = {"raise_on_checkout": False, "n": 0} + + def factory(): + state["n"] += 1 + if state["raise_on_checkout"]: + raise RuntimeError("token refresh failed mid-checkout") + return {}, near_expiry + + pool_key = conn_str + "\x00mssql_test_factory_raise_recover" + + # Prime the pool with a near-expiry pooled connection. + conn1 = ddbc_bindings.Connection(conn_str, True, {}, pool_key, factory) + assert state["n"] == 1 + conn1.close() + + # Next checkout is near-expiry -> factory runs to refresh, but raises. + state["raise_on_checkout"] = True + with pytest.raises(Exception): + ddbc_bindings.Connection(conn_str, True, {}, pool_key, factory) + + # The reserved slot must have been released: a healthy factory can still + # open a connection under the same key (pool is neither wedged nor full). + state["raise_on_checkout"] = False + conn3 = ddbc_bindings.Connection(conn_str, True, {}, pool_key, factory) + assert conn3 is not None, "pool slot must be recovered after a factory failure" + conn3.close() + + def test_orphaned_connection_is_disconnected_on_return(self, conn_str): + """Returning a connection whose pool was evicted disconnects it cleanly. + + Covers the orphan branch of ``ConnectionPoolManager::returnConnection``: + when a checked-out connection is returned but no pool is registered under + its key (here because ``close_pooling()`` cleared the pool map while the + connection was still out), the manager must disconnect the orphan rather + than leak the ODBC handle. + """ + if not conn_str: + pytest.skip("Live database connection required") + + from mssql_python import ddbc_bindings + + def factory(): + return {} + + pool_key = conn_str + "\x00mssql_test_orphan_return" + + # Check a connection out of the pool, then drop the whole pool map while + # it is still held. The pool under pool_key no longer exists. + conn = ddbc_bindings.Connection(conn_str, True, {}, pool_key, factory) + ddbc_bindings.close_pooling() + + # Returning it now finds no pool -> the orphan-disconnect path runs. + conn.close() + + def test_disabled_manager_creates_no_pool(self, conn_str): + """A pooled-style connect after disable_pooling() creates no pool. + + Regression test for the disable()-vs-connect() race: ``disable_pooling()`` + disarms the native manager under ``_manager_mutex`` before clearing the + pool map, so any ``acquireConnection`` serialized after it declines + (returns nullptr) and the connection transparently falls back to a + non-pooled one. A connect can therefore never resurrect a pool after a + disable, so ``_pool_count()`` stays 0. The non-pooled fallback still + honors the token factory. + """ + if not conn_str: + pytest.skip("Live database connection required") + + from mssql_python import ddbc_bindings + + try: + # Arm, then disable: disarm new-pool creation and close everything. + ddbc_bindings.enable_pooling(10, 600) + ddbc_bindings.disable_pooling() + assert ddbc_bindings._pool_count() == 0, "disable must leave no pools" + + calls = {"n": 0} + + def factory(): + calls["n"] += 1 + return {} + + pool_key = conn_str + "\x00mssql_test_disabled_nopool" + + # use_pool=True, but the manager is disarmed -> non-pooled fallback. + conn = ddbc_bindings.Connection(conn_str, True, {}, pool_key, factory) + assert calls["n"] == 1, "non-pooled fallback must still invoke the factory" + assert ( + ddbc_bindings._pool_count() == 0 + ), "a disabled manager must not create a pool for a racing connect" + conn.close() + assert ( + ddbc_bindings._pool_count() == 0 + ), "returning the fallback connection must not create a pool" + finally: + # Re-arm so sibling tests pool normally again. + ddbc_bindings.enable_pooling(10, 600)