FEAT: Add "token_provider= " parameter for custom Azure Identity credential support by jahnvi480 · Pull Request #603 · microsoft/mssql-python

jahnvi480 · 2026-05-27T14:11:20Z

Work Item / Issue Reference

AB#45121

GitHub Issue: #577

Summary

This pull request introduces a new token_provider parameter for Microsoft Entra ID (Azure AD) authentication, enabling the use of any credential object with a .get_token(scope) method (such as those from azure-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:

Added a token_provider parameter to connect() and the Connection class, accepting any object with a get_token(scope) method for Microsoft Entra ID authentication (e.g., DefaultAzureCredential, AzureCliCredential, ManagedIdentityCredential). This is mutually exclusive with Authentication= in the connection string and with pre-acquired tokens in attrs_before. Bulk copy operations now re-acquire a fresh token from the provider for each operation.
Introduced the TokenProvider protocol to define the expected interface for custom credential objects, ensuring compatibility with both minimal implementations and all azure-identity credentials.

Token Acquisition and Error Handling:

Added _get_token_from_credential, acquire_token_from_credential, and acquire_raw_token_from_credential utility 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 .token attribute) and warn if an already-expired token is returned.
Improved documentation and logging for the Azure SQL Database OAuth scope, clarifying that only the Azure commercial cloud is supported via the token_provider path.

Connection Pooling and Identity Key Computation:

Refined the logic for computing connection pool identity keys: managed identities now use distinct prefixes (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.
Added compute_token_identity to encapsulate the hashing logic for token-based identity keys.

Type Checking and Coverage:

Updated .coveragerc to 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.

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

Copilot

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 credential parameter to connect() / Connection to accept objects implementing .get_token(scope).
Implemented credential-based token acquisition helpers in auth.py and integrated credential token usage into Cursor.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.

github-actions · 2026-05-27T14:46:07Z

📊 Code Coverage Report

🔥 Diff Coverage 98%	🎯 Overall Coverage 82%	📈 Total Lines Covered: `7361` out of `8973` 📁 Project: `mssql-python`

Diff Coverage

Diff: main...HEAD, staged and unstaged changes

mssql_python/init.py (100%)
mssql_python/auth.py (97.9%): Missing lines 727
mssql_python/connection.py (98.1%): Missing lines 84
mssql_python/constants.py (100%)
mssql_python/cursor.py (100%)
mssql_python/db_connection.py (100%)

Summary

Total: 117 lines
Missing: 2 lines
Coverage: 98%

mssql_python/auth.py

Lines 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.py

Lines 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.2%
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: 82.7%
mssql_python.pybind.connection.connection.cpp: 83.7%
mssql_python.logging.py: 85.5%

🔗 Quick Links

⚙️ Build Summary	📋 Coverage Details
View Azure DevOps Build	Browse Full Coverage Report

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

…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).

…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.

…#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.

sumitmsft · 2026-07-31T07:45:38Z

Looks good.. have some minor comments but not very critical to solve and can be ignored.

Please re-review

bewithgaurav

added comments inline, the PR description needs revision - requesting changes
can skip some minor ones if needed

…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.

After merging the Arrow bulk copy feature (#665), _build_pycore_context now reads self.connection._token_provider before the _auth_type dispatch. The test_024 _cursor_with_conn helper builds a bare MagicMock connection, which auto-vivifies a truthy _token_provider and routed every case through the custom-credential token path -> InterfaceError (12 failures in CI). Real Connection objects always initialize _token_provider to None; set it explicitly on the mock so it is faithful.

jahnvi480 requested review from bewithgaurav, Copilot, saurabh500 and sumitmsft May 27, 2026 14:11

Copilot started reviewing on behalf of jahnvi480 May 27, 2026 14:11 View session

github-actions Bot added the pr-size: medium Moderate update size label May 27, 2026

Copilot AI reviewed May 27, 2026

View reviewed changes

Comment thread mssql_python/connection.py Outdated

Comment thread mssql_python/db_connection.py

jahnvi480 added 2 commits May 27, 2026 19:51

jahnvi480 and others added 3 commits May 27, 2026 22:00

Merge branch 'main' into jahnvi/custom-credential-support

819549b

Merge branch 'main' into jahnvi/custom-credential-support

5370672

jahnvi480 changed the title ~~FEAT: Add "credential= " parameter for custom Azure Identity credential support~~ FEAT: Add "token_provider= " parameter for custom Azure Identity credential support May 29, 2026

jahnvi480 added 2 commits May 29, 2026 22:45

Merge branch 'main' into jahnvi/custom-credential-support

b1c8a59

Merge branch 'main' into jahnvi/custom-credential-support

09b7198

jahnvi480 marked this pull request as draft June 2, 2026 03:17

jahnvi480 and others added 4 commits June 23, 2026 13:49

Merge branch 'main' into jahnvi/custom-credential-support

e35270e

Merge branch 'main' into jahnvi/custom-credential-support

8e4e04a

Merge remote-tracking branch 'origin/main' into jahnvi/custom-credent…

42f976b

…ial-support # Conflicts: # mssql_python/cursor.py # tests/test_008_auth.py

github-actions Bot added pr-size: large Substantial code update and removed pr-size: medium Moderate update size labels Jun 25, 2026

jahnvi480 and others added 6 commits June 26, 2026 10:05

Merge branch 'main' into jahnvi/custom-credential-support

b49daa8

Merge branch 'main' into jahnvi/custom-credential-support

bc59232

jahnvi480 requested a review from sdebruyn July 27, 2026 10:02

jahnvi480 added 2 commits July 28, 2026 12:24

Merge branch 'main' into jahnvi/custom-credential-support

c452ff7

Merge branch 'main' into jahnvi/custom-credential-support

a81f3c9

sumitmsft previously approved these changes Jul 31, 2026

View reviewed changes

jahnvi480 added 2 commits August 3, 2026 15:38

Merge branch 'main' into jahnvi/custom-credential-support

6a74206

Merge branch 'main' into jahnvi/custom-credential-support

272bf67

bewithgaurav requested changes Aug 4, 2026

View reviewed changes

jahnvi480 dismissed sumitmsft’s stale review via 2fd1d4f August 4, 2026 08:49

jahnvi480 requested review from bewithgaurav and sumitmsft August 4, 2026 08:58

jahnvi480 added 2 commits August 4, 2026 14:39

sumitmsft previously approved these changes Aug 4, 2026

View reviewed changes

Merge branch 'main' into jahnvi/custom-credential-support

7c36299

bewithgaurav previously approved these changes Aug 4, 2026

View reviewed changes

Merge branch 'main' into jahnvi/custom-credential-support

87b8048

jahnvi480 dismissed stale reviews from bewithgaurav and sumitmsft via 87b8048 August 4, 2026 11:50

bewithgaurav previously approved these changes Aug 4, 2026

View reviewed changes

sumitmsft previously approved these changes Aug 4, 2026

View reviewed changes

jahnvi480 dismissed stale reviews from sumitmsft and bewithgaurav via bd5e5f7 August 4, 2026 14:17

jahnvi480 requested review from bewithgaurav and sumitmsft August 4, 2026 14:31

Navigation Menu

Uh oh!

FEAT: Add "token_provider= " parameter for custom Azure Identity credential support - #603

FEAT: Add "token_provider= " parameter for custom Azure Identity credential support#603
jahnvi480 wants to merge 34 commits into
mainfrom
jahnvi/custom-credential-support

jahnvi480 commented May 27, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

github-actions Bot commented May 27, 2026 •

edited

Loading

🔥 Diff Coverage

98%

🎯 Overall Coverage

82%

Uh oh!

sumitmsft commented Jul 31, 2026

Uh oh!

bewithgaurav left a comment •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Uh oh!

Conversation

jahnvi480 commented May 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Work Item / Issue Reference

Summary

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

github-actions Bot commented May 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

📊 Code Coverage Report

🔥 Diff Coverage

98%

🎯 Overall Coverage

82%

Diff Coverage

Diff: main...HEAD, staged and unstaged changes

Summary

mssql_python/auth.py

mssql_python/connection.py

📋 Files Needing Attention

🔗 Quick Links

Uh oh!

sumitmsft commented Jul 31, 2026

Uh oh!

bewithgaurav left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

jahnvi480 commented May 27, 2026 •

edited

Loading

github-actions Bot commented May 27, 2026 •

edited

Loading

bewithgaurav left a comment •

edited

Loading