FEAT: Add "token_provider= " parameter for custom Azure Identity credential support - #603
Open
jahnvi480 wants to merge 32 commits into
Open
FEAT: Add "token_provider= " parameter for custom Azure Identity credential support#603jahnvi480 wants to merge 32 commits into
jahnvi480 wants to merge 32 commits into
Conversation
Add a new 'credential' parameter to connect() that accepts any object following the Azure TokenCredential protocol (.get_token() method). This allows users to authenticate with any azure-identity credential class without being limited to the driver's hardcoded credential map. Changes: - auth.py: Add _get_token_from_credential() shared helper, acquire_token_from_credential(), acquire_raw_token_from_credential() - db_connection.py: Add credential=None parameter to connect() - connection.py: Validate credential, acquire token, store for bulk copy token refresh. Mutually exclusive with Authentication= - cursor.py: Check _custom_credential before _auth_type in bulk copy - constants.py: Unify _KEY_* constants with _ALLOWED_CONNECTION_STRING_PARAMS to use single source of truth (_CONNECTION_STRING_*_KEY pattern) - test_008_auth.py: Add 12 new tests for custom credential flow
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds a new credential= parameter to the public connection API to support custom Azure Identity (Entra ID) credential objects for token acquisition, and wires that credential through to bulk copy so fresh tokens can be acquired when needed.
Changes:
- Added
credentialparameter toconnect()/Connectionto accept objects implementing.get_token(scope). - Implemented credential-based token acquisition helpers in
auth.pyand integrated credential token usage intoCursor.bulkcopy(). - Refactored connection-string key constants and added test coverage for the new credential flows.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
tests/test_008_auth.py |
Adds unit tests for the new credential token helpers and connect(..., credential=...) behaviors. |
mssql_python/db_connection.py |
Extends the connect() API to accept and forward the new credential parameter. |
mssql_python/cursor.py |
Updates bulk copy to acquire a fresh token from a user-supplied custom credential when present. |
mssql_python/constants.py |
Refactors connection-string key constants/aliases used by auth/connection code. |
mssql_python/connection.py |
Implements the new credential parameter behavior (validation, token acquisition, mutual exclusivity with Authentication=). |
mssql_python/auth.py |
Adds centralized helpers for acquiring raw tokens / ODBC token structs from custom credentials. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
- Strip UID/PWD/Trusted_Connection from connection_str when credential= is used (same as Authentication= path) to avoid leaking unused secrets - Add credential= parameter to Connection.__init__ and connect() in mssql_python.pyi type stubs
The _make_cursor helper uses MagicMock for the connection, which auto-creates truthy attributes. Without explicitly setting _custom_credential = None, the bulk copy code takes the custom credential path instead of the expected _auth_type path.
📊 Code Coverage Report
Diff CoverageDiff: main...HEAD, staged and unstaged changes
Summary
mssql_python/auth.pyLines 723-731 723 frame = frame.f_back
724 level += 1
725 # Every frame was internal (not expected in practice); fall back to the
726 # outermost real frame rather than blaming this helper.
! 727 return max(level - 1, 1)
728
729
730 def _get_token_from_credential(
731 credential: "TokenCredential",mssql_python/connection.pyLines 80-88 80 """
81
82 def get_token(self, scope: str) -> Any:
83 """Return an object with a ``.token`` attribute for ``scope``."""
! 84 ...
85
86
87 # Add SQL_WMETADATA constant for metadata decoding configuration
88 SQL_WMETADATA: int = -99 # Special flag for column name decoding📋 Files Needing Attention📉 Files with overall lowest coverage (click to expand)mssql_python.pybind.logger_bridge.cpp: 59.2%
mssql_python.pybind.ddbc_bindings.h: 59.9%
mssql_python.pybind.logger_bridge.hpp: 70.8%
mssql_python.pybind.ddbc_bindings.cpp: 76.3%
mssql_python.__init__.py: 77.6%
mssql_python.row.py: 77.6%
mssql_python.ddbc_bindings.py: 79.6%
mssql_python.pybind.connection.connection_pool.cpp: 81.4%
mssql_python.pybind.connection.connection.cpp: 83.7%
mssql_python.logging.py: 85.5%🔗 Quick Links
|
Rename the public API parameter from 'credential' to 'token_provider' to reduce ambiguity in our multi-auth-path context. 'credential' could be confused with SQL auth username/password; 'token_provider' clearly signals token-based Entra ID auth. - Rename parameter: credential -> token_provider (connect, Connection) - Rename internal attr: _custom_credential -> _token_provider - Update error messages, docstrings, comments, .pyi stubs - Improve docstring with usage example and explicit guidance - All 97 tests pass
jahnvi480
changed the title
jahnvi480
marked this pull request as draft
June 2, 2026 03:17
…ial-support # Conflicts: # mssql_python/cursor.py # tests/test_008_auth.py
…lk-copy token branch C2: capture token expires_on from custom credential and store on connection. C3: raise DB-API InterfaceError/OperationalError instead of ValueError/TypeError for token_provider misuse and acquisition failures. Add unit tests covering the cursor bulk-copy token_provider branch (success, get_token failure, invalid token).
github-actions
Bot
added
pr-size: large
…ol typing; fix docstring error type - auth.py: type credential params as TokenProvider Protocol; hard-code commercial-cloud scope - connection.py: warn on ignored UID/PWD/Trusted_Connection when token_provider set; validate get_token arity; document token lifecycle limitations - db_connection.py: note sovereign clouds out of scope - test_008_auth.py: cover arity validation and dropped-credential warning
…al collision The native connection pool keys on the sanitized connection string only, and the access token lives in attrs_before (applied once on a new physical connection, never re-applied on reuse). Two different principals sharing the same Server/Database collapsed into one pool bucket, so one caller could be handed another's authenticated connection (silent identity confusion). Fix: Connection.__init__ disables pooling whenever SQL_COPT_SS_ACCESS_TOKEN is present in attrs_before. One condition covers all access-token paths: raw attrs_before token, built-in Authentication=ActiveDirectory* (token-injecting), and token_provider=. Driver-native paths (e.g. ServicePrincipal) keep creds in the connection string and remain poolable. Adds regression tests in TestTokenProviderPooling.
- Remove unnecessary f-prefix from two non-interpolated SQL_WCHAR error strings in connection.py (flagged by flake8-no-fstring-u style linters). - Type token_provider as Optional[TokenProvider] (was Optional[object]) in the .pyi stubs for Connection.__init__ and connect(), matching the runtime.
- Drop the 'See docs/DESIGN_TOKEN_PROVIDER_SUPPORT.md' comment in connection.py (that file is not part of the PR). - The expired-token warning in _get_token_from_credential is reached via two call chains at different depths (connect vs bulk-copy), so a fixed stacklevel cannot point at user code for both. Compute the stacklevel dynamically via _stacklevel_to_caller(), which walks out of the package to the first external frame. Works across all supported Python versions.
…token warning The _stacklevel_to_caller() helper walked frames on every expired-token warning and relied on __file__ (fragile under zipimport/frozen deploys). The expired case is rare and the message is self-explanatory, so a fixed stacklevel=2 is sufficient. Removes the helper, _PACKAGE_DIR, and the now-unused os import.
jahnvi480
marked this pull request as ready for review
June 29, 2026 05:15
…m TokenProvider Protocol
…y to test requirements
jahnvi480
marked this pull request as draft
June 29, 2026 09:30
jahnvi480
marked this pull request as ready for review
June 30, 2026 03:11
sdebruyn
previously requested changes
Jul 2, 2026
jahnvi480
added a commit
that referenced
this pull request
Jul 3, 2026
…#659) For MSI auth the identity (client id) is known without a token, so the pool key can be computed up front. Acquisition of the access token is deferred to an internal token_factory callback that the native layer invokes only when it opens a physical connection. A same-identity pool hit therefore reuses the pooled connection and never pays for a token. Token-dependent identities (DefaultAzureCredential / raw token / interactive / device-code) still acquire eagerly because the pool key is the token hash. The internal token_factory is intentionally distinct from the public token_provider= credential API proposed in PR #603 (issue #577).
Integrate identity-aware connection pooling (#660) from main and remove the now-obsolete access-token pooling-disable guard. With the identity-aware pool key (connStr + msi:/acct:/tok: identity), access-token connections (token_provider=, built-in Authentication=ActiveDirectory*, and raw SQL_COPT_SS_ACCESS_TOKEN) are safe to pool: distinct principals land in distinct pool buckets, while same-identity reuse benefits from pooling. - connection.py: drop the disable-pooling guard; fix the 'if token_provider / elif Authentication' chain broken by the auto-merge; update the pooling docstring note. - test_008_auth.py: TestTokenProviderPooling now asserts pooling is enabled with an identity-aware pool key for every access-token path. - auth.py / test_020: resolve merge conflicts; consolidate the OAuth scope constant on _SQL_SCOPE.
Contributor
|
Looks good.. have some minor comments but not very critical to solve and can be ignored. |
sumitmsft
previously approved these changes
Jul 31, 2026
bewithgaurav
requested changes
Aug 4, 2026
…on, MSI disjoint encoding, acct account-binding) Revert token_provider= pooling from credential object-identity back to token-hash (tok:) keying: object identity assumes a credential object maps to a stable principal, but mutable credentials (e.g. AzureCliCredential) can mint different principals from the same object, so a pool hit could hand back a stale-principal connection. Token-hash re-derives identity from the actual token every connect. Additional hardening: reject embedded NUL in connection string/kwargs at the Python boundary; encode MSI identities disjointly (msi:client:<id> vs msi:system); bind the deferred acct: token factory to the pooled home_account_id; add compute_token_identity helper; freeze bytearray raw tokens to immutable bytes; derive warning stacklevel dynamically.
The .pyi stub is not shipped (not in setup.py package_data) and is misnamed for a package stub, so type checkers never read it. Real type info for token_provider comes from inline hints in db_connection.py plus the shipped py.typed marker. Reverting keeps the PR surgical; ship-vs-delete of the stub file is tracked as a separate backlog item.
- Re-export TokenProvider from the package root so rom mssql_python import
TokenProvider works for type annotations (added to __all__).
- Normalize MSI client_id (strip {braces}, lowercase) in compute_identity_key
so the same managed identity keys one pool regardless of GUID casing/bracing.
- Downgrade the already-expired-token diagnostic from warnings.warn to
logger.warning so it is never promoted to an exception under -W error.
- Type _token_expires_on as Optional[float] (a custom provider may report a
float POSIX timestamp).
- CHANGELOG: note that NUL characters in connection strings/params are now
rejected with InterfaceError instead of being silently truncated.
Tests updated accordingly (log assertion instead of pytest.warns) plus new
cases for the re-export and client_id normalization.
sumitmsft
approved these changes
Aug 4, 2026
bewithgaurav
approved these changes
Aug 4, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Work Item / Issue Reference
Summary
This pull request introduces a new
token_providerparameter for Microsoft Entra ID (Azure AD) authentication, enabling the use of any credential object with a.get_token(scope)method (such as those fromazure-identity). It also adds robust support for custom credential objects, improves token acquisition error handling, and refines identity key computation for connection pooling. Several internal utilities and documentation comments have been added to clarify and harden the authentication process.Authentication and Token Provider Support:
token_providerparameter toconnect()and theConnectionclass, accepting any object with aget_token(scope)method for Microsoft Entra ID authentication (e.g.,DefaultAzureCredential,AzureCliCredential,ManagedIdentityCredential). This is mutually exclusive withAuthentication=in the connection string and with pre-acquired tokens inattrs_before. Bulk copy operations now re-acquire a fresh token from the provider for each operation.TokenProviderprotocol to define the expected interface for custom credential objects, ensuring compatibility with both minimal implementations and allazure-identitycredentials.Token Acquisition and Error Handling:
_get_token_from_credential,acquire_token_from_credential, andacquire_raw_token_from_credentialutility functions to centralize token acquisition, handle errors, and log token details, including expiry. These functions provide clear error messages for common mistakes (e.g., async credentials, missing.tokenattribute) and warn if an already-expired token is returned.token_providerpath.Connection Pooling and Identity Key Computation:
msi:client:<client_id>vs.msi:system), and token-based identities use a SHA-256 hash of the token. The computation is now expiry-aware and avoids collisions between system- and user-assigned identities.compute_token_identityto encapsulate the hashing logic for token-based identity keys.Type Checking and Coverage:
.coveragercto exclude type-checking-only imports from coverage, improving test accuracy.These changes significantly improve the flexibility, reliability, and clarity of Microsoft Entra ID authentication support in the codebase.