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

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

Draft
jahnvi480 wants to merge 19 commits into
mainfrom
jahnvi/custom-credential-support
Draft

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

Conversation

@jahnvi480

@jahnvi480 jahnvi480 commented May 27, 2026

Copy link
Copy Markdown
Contributor

Work Item / Issue Reference

AB#45121

GitHub Issue: #577


Summary

This pull request introduces support for custom Microsoft Entra ID (Azure AD) authentication in the connection layer by adding a new token_provider parameter to connect() and the Connection class. This enables users to supply any credential object (such as those from azure-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_provider parameter to connect() and the Connection class, 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 with Authentication= and attrs_before[SQL_COPT_SS_ACCESS_TOKEN].

  • Introduced the TokenProvider protocol in auth.py to 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, and acquire_raw_token_from_credential functions 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:

  • Disabled connection pooling automatically for any access-token connection (including those using token_provider, built-in Authentication=ActiveDirectory*, or raw attrs_before tokens) 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 .coveragerc to 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.

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 AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread mssql_python/connection.py Outdated
Comment thread mssql_python/db_connection.py
jahnvi480 added 2 commits May 27, 2026 19:51
- 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

github-actions Bot commented May 27, 2026

Copy link
Copy Markdown

📊 Code Coverage Report

🔥 Diff Coverage

100%


🎯 Overall Coverage

81%


📈 Total Lines Covered: 6808 out of 8376
📁 Project: mssql-python


Diff Coverage

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

  • mssql_python/auth.py (100%)
  • mssql_python/connection.py (100%)
  • mssql_python/constants.py (100%)
  • mssql_python/db_connection.py (100%)

Summary

  • Total: 81 lines
  • Missing: 0 lines
  • Coverage: 100%

📋 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

⚙️ Build Summary 📋 Coverage Details

View Azure DevOps Build

Browse Full Coverage Report

jahnvi480 and others added 3 commits May 27, 2026 22:00
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 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 jahnvi480 marked this pull request as draft June 2, 2026 03:17
jahnvi480 and others added 4 commits June 23, 2026 13:49
…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 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
…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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

pr-size: large Substantial code update

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants