Repository navigation
FEAT: Entra id authentication support for AQE - #870
Subrata (subrata-ms) wants to merge 4 commits into
Conversation
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
PR Performance Report✅ No regression detectedNo consistent slowdowns detected across all 2 environments. 0 IMPROVEMENTS 0 SLOWDOWNS 2/2 ENVIRONMENTS Coverage: 2 of 2 environments completed. Advisory result; does not block merging. Performance diagnosticsPhase times are inclusive diagnostics and must not be added together. They identify where measured time changed, not why it changed. No affected phases or call-count changes were recorded. All database tasks and timingsUnix / SQL Server 2022
Unix / SQL Server 2025
Build and measurement detailsPR head:
A consistent change requires more than 20% median paired movement, at least 1 ms between the median runtimes, and at least 80% of pairs exceeding the relative threshold in the same direction. A slowdown without enough pair agreement is reported as inconsistent. The displayed change is the median of paired before-and-after ratios. It is not recalculated from the two displayed median runtimes. Both revisions use profiling-enabled builds on the same agent and database, with alternating order and discarded warmups. Results are diagnostic and do not represent production-wheel latency. Raw samples and logs are attached to the ADO run as |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
There was a problem hiding this comment.
🔵 Needs a closer look
Entra authentication depends on native Rust handshake and callback behavior that could not be runtime-validated with the pinned extension and a live Entra-enabled SQL endpoint.
No open issues 🎉
What changed in this PR
Adds Microsoft Entra ID authentication to the Rust-backed async connection layer, supporting built-in Azure Identity credentials and custom synchronous token providers.
Changes:
- Adds credential caching, off-loop token acquisition, and service-principal callback setup.
- Extends async connection options and validation, documenting authentication modes and limitations.
- Adds tests for authentication flows, error handling, cancellation, credential reuse, and native integration.
| File | Description |
|---|---|
tests/AsyncTest/test_008_async_entraid_auth.py |
Adds async Entra authentication and integration coverage. |
tests/AsyncTest/test_002_async_connection.py |
Updates context validation expectations to accept Entra authentication. |
mssql_python/async_query/async_connection.py |
Adds token_provider and prepares authentication before native login. |
mssql_python/async_query/_connection_context.py |
Allows supported Entra modes through connection validation. |
mssql_python/async_query/_auth.py |
Implements credential caching, token acquisition, and service-principal setup. |
Review effort: Balanced
Subrata (subrata-ms)
changed the title
📊 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.performance_counter.hpp: 0.7%
mssql_python.pybind.logger_bridge.cpp: 57.9%
mssql_python.pybind.ddbc_bindings.h: 62.6%
mssql_python.pybind.logger_bridge.hpp: 70.8%
mssql_python.pybind.connection.connection_pool.cpp: 81.3%
mssql_python.pybind.connection.connection.cpp: 84.1%
mssql_python.logging.py: 86.2%
mssql_python.pooling.py: 90.1%
mssql_python.pybind.fetch_temporal.hpp: 92.1%
mssql_python.cursor.py: 92.5%🔗 Quick Links
|
Subrata (subrata-ms)
changed the title
There was a problem hiding this comment.
🔵 Needs a closer look
The native authentication implementation is unavailable, and existing mocks cannot establish Entra SQL handshake success, callback threading, or cancellation behavior without live validation.
No open issues 🎉
Review effort: Balanced
| result = await asyncio.to_thread( | ||
| core._invoke_entra_id_token_factory, |
Sumit Sarabhai (sumitmsft)
left a comment
Sumit Sarabhai (sumitmsft)
left a comment
There was a problem hiding this comment.
Entra async-authentication review against the supplied AQE design and Azure Identity implementation. Posting the four approved inline comments; no approval or request-changes decision is submitted.
| client_id = None | ||
| if auth_type == _AuthInternal.MSI: | ||
| client_id = (context.get("user_name") or "").strip() or None | ||
| raw_token = await asyncio.to_thread(_acquire_builtin_token, auth_type, client_id) |
There was a problem hiding this comment.
§4 of the AQE design explicitly requires azure.identity.aio credentials and await credential.get_token(...). Azure Identity provides async DefaultAzureCredential and ManagedIdentityCredential, but this path runs their synchronous counterparts through asyncio's shared executor. The custom-provider branch also reuses the sync helper that rejects async credentials, so callers cannot supply the credential model required by the design.
Please use async acquisition for modes with supported SDK implementations and accept async token providers, preserving token validation and DB-API error mapping. Give driver-owned async credentials an explicit event-loop-compatible lifetime and await their cleanup; leave caller-owned providers caller-owned. Any necessary synchronous compatibility paths should be separate and explicitly documented rather than making every mode thread-backed.
| raw_token, _ = await asyncio.to_thread(acquire_raw_token_from_credential, token_provider) | ||
| elif auth_type == _AuthInternal.SERVICE_PRINCIPAL: | ||
| try: | ||
| authenticated["entra_id_token_factory"] = ServicePrincipalAuth.make_token_factory( |
There was a problem hiding this comment.
This wires the synchronous bulk-copy token factory into the new async ServicePrincipal path. That factory calls synchronous ClientSecretCredential.get_token(), although Azure Identity provides aio.ClientSecretCredential. Offloading the factory prevents event-loop blocking, but still occupies a worker throughout token acquisition.
Please make this path await async acquisition after the server supplies the tenant/SPN, coordinating any required native callback-contract and dependency-pin changes. The inspected native bridge expects immediate bytes, so simply changing the Python factory to async def would not work. Preserve the existing synchronous bulk-copy behavior. If an awaitable callback cannot be supported in this slice, explicitly gate and document ServicePrincipal as an approved blocking exception rather than claiming end-to-end async compliance.
The native bridge was inspected at upstream commit bf415d10; this does not establish the exact repository-pinned wheel's implementation. The Python factory's synchronous token acquisition is directly verifiable here.
|
|
||
|
|
||
| @pytest.mark.asyncio | ||
| async def test_live_configured_authentication_login_and_query(async_connection): |
There was a problem hiding this comment.
This test accepts the generic DB_CONNECTION_STRING, so it can pass with SQL-password authentication without exercising an Entra handshake. The closed-port native tests and separately invoked callback hook do not establish successful FedAuth login.
Please add explicitly Entra-configured integration coverage for the pre-acquired-token and ServicePrincipal callback paths. Add async-provider acceptance, cancellation and credential-cleanup tests as well—the current test at line 457 instead requires async credentials to be rejected. For modes intended to be genuinely async, verify that the driver does not offload token acquisition to an executor; an event-loop heartbeat alone also passes synchronous acquisition moved to a worker.
|
|
||
| Entra modes: ActiveDirectoryDefault, ActiveDirectoryMsi (optional UID | ||
| selects a user-assigned identity), ActiveDirectoryDeviceCode, | ||
| ActiveDirectoryInteractive, and client-secret ActiveDirectoryServicePrincipal. |
There was a problem hiding this comment.
Please document a per-mode async capability matrix before describing this feature as end-to-end async.
Azure Identity's public async API supports Default, Managed Identity and client-secret authentication, but does not expose aio.DeviceCodeCredential or aio.InteractiveBrowserCredential in the checked versions (1.12.0 and 1.24.0). Those interactive modes need an explicitly approved compatibility exception or a clear unsupported status; they should not prevent genuinely async implementation of the other modes.
The matrix should also cover Default credential-chain behavior on supported event loops: the SDK's CLI credential has a synchronous fallback on Windows Selector loops. Please ensure the selected chain remains nonblocking, or explicitly approve and document that limitation.
Finally, declare the async HTTP transport requirement and credential ownership/cleanup policy. Azure Identity does not install an async transport automatically. Confirm how that prerequisite fits the design's “no new PyPI dep” requirement rather than relying on an incidental development-environment installation.
Work Item / Issue Reference
Summary
This pull request introduces Entra (Azure Active Directory) authentication support for asynchronous connections in the
mssql_pythonpackage. It adds a new authentication helper, updates the async connection logic to handle various Entra authentication modes and token providers, and improves validation and error handling for authentication options. The tests have also been updated to reflect the new authentication flow.Entra Authentication Support and Async Connection Updates:
mssql_python/async_query/_auth.pythat provides Entra authentication helpers, credential caching, and async token acquisition for async connections. It supports built-in Azure Identity credentials and custom token providers, with careful handling of multi-user scenarios.mssql_python/async_query/async_connection.pyto:token_providerparameter for supplying a custom token provider.prepare_async_authenticationfunction to handle authentication and token acquisition before establishing a connection.connectmethod's docstring to explain Entra authentication modes, token provider usage, and limitations. [1] [2] [3] [4]Authentication Validation Improvements:
mssql_python/async_query/_connection_context.pyto accept supported Entra authentication modes and provide clearer error messages for unsupported modes. [1] [2]Testing Updates:
tests/AsyncTest/test_002_async_connection.pyto check that Entra authentication modes are accepted in the connection context without immediately acquiring a token, reflecting the new authentication flow.