{{ message }}
FEAT: Add "token_provider= " parameter for custom Azure Identity credential support#603
Draft
jahnvi480 wants to merge 19 commits into
Draft
FEAT: Add "token_provider= " parameter for custom Azure Identity credential support#603jahnvi480 wants to merge 19 commits into
jahnvi480 wants to merge 19 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
💡 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
📋 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.5%
mssql_python.row.py: 76.9%
mssql_python.__init__.py: 77.3%
mssql_python.pybind.connection.connection.cpp: 77.3%
mssql_python.ddbc_bindings.py: 79.6%
mssql_python.logging.py: 85.5%
mssql_python.connection.py: 86.8%🔗 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
…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.
…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.
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 support for custom Microsoft Entra ID (Azure AD) authentication in the connection layer by adding a new
token_providerparameter toconnect()and theConnectionclass. This enables users to supply any credential object (such as those fromazure-identity) that exposes a.get_token(scope)method, facilitating advanced authentication scenarios. The implementation includes robust error handling, disables connection pooling for access-token connections to prevent security risks, and documents the new feature thoroughly.The most important changes are:
Custom Token Provider for Entra ID Authentication:
Added a
token_providerparameter toconnect()and theConnectionclass, allowing users to supply any object with a.get_token(scope)method (e.g.,DefaultAzureCredential,AzureCliCredential,ManagedIdentityCredential) for Microsoft Entra ID authentication. This is mutually exclusive withAuthentication=andattrs_before[SQL_COPT_SS_ACCESS_TOKEN].Introduced the
TokenProviderprotocol inauth.pyto define the expected interface for custom credentials, and added detailed documentation and type annotations.Token Acquisition and Validation Logic:
Implemented
_get_token_from_credential,acquire_token_from_credential, andacquire_raw_token_from_credentialfunctions to centralize token acquisition, error handling, and diagnostics (including expiry checks and async credential detection).Updated token acquisition paths to use a shared constant for the Azure commercial cloud scope (
https://database.windows.net/.default), clarifying that sovereign clouds are out of scope for this feature.Connection Pooling and Security:
token_provider, built-inAuthentication=ActiveDirectory*, or rawattrs_beforetokens) to prevent identity confusion and privilege escalation.Other Improvements and Documentation:
Improved parameter documentation and error messages for the new authentication flow, and ensured that user dictionaries are not mutated when injecting tokens.
Updated
.coveragercto exclude type-checking-only imports from coverage.These changes collectively enable secure, flexible, and user-friendly integration with Microsoft Entra ID using custom credentials, while addressing important security and usability concerns.