FIX: Add ActiveDirectoryMSI support for bulk copy#573
Open
bewithgaurav wants to merge 1 commit into
Open
Conversation
Adds Authentication=ActiveDirectoryMSI to the auth pipeline: - Zero-arg ManagedIdentityCredential() for system-assigned MSI. - ManagedIdentityCredential(client_id=UID) for user-assigned MSI, matching ODBC's convention where UID carries the identity's client_id under MSI. - Threads optional credential_kwargs through get_auth_token / get_raw_token / _acquire_token so future auth methods that need constructor args (e.g. ClientSecretCredential) can plug in via the same channel. - Cache key remains a plain string for zero-arg auth types and becomes a tuple when kwargs are present, so different client_ids get separate cached credentials. Partial fix for #534. ServicePrincipal and Password to follow as separate PRs.
Contributor
There was a problem hiding this comment.
Pull request overview
Adds support for Authentication=ActiveDirectoryMSI (managed identity) token acquisition for the bulk copy (mssql-py-core) path, including user-assigned MSI via UID=<client_id>.
Changes:
- Add
AuthType.MSI(activedirectorymsi) and map it toManagedIdentityCredential. - Thread optional
credential_kwargsthrough token acquisition and update the credential-instance cache key to incorporate kwargs. - Attempt to extract MSI
client_idfrom the connection string and pass it into bulk copy token acquisition; add tests and a changelog entry.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
mssql_python/constants.py |
Adds AuthType.MSI enum value. |
mssql_python/auth.py |
Adds MSI credential support, credential kwargs plumbing, and cache key changes; adds extract_credential_kwargs. |
mssql_python/cursor.py |
Bulk copy now tries to extract MSI credential kwargs and pass them to get_raw_token. |
tests/test_008_auth.py |
Adds tests for MSI auth type parsing, credential construction, cache isolation, and connection-string processing. |
CHANGELOG.md |
Documents MSI support for bulk copy. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+2915
to
+2923
| from mssql_python.auth import AADAuth, extract_credential_kwargs | ||
|
|
||
| credential_kwargs = extract_credential_kwargs( | ||
| self.connection.connection_str, self.connection._auth_type | ||
| ) | ||
| try: | ||
| raw_token = AADAuth.get_raw_token(self.connection._auth_type) | ||
| raw_token = AADAuth.get_raw_token( | ||
| self.connection._auth_type, credential_kwargs or None | ||
| ) |
Comment on lines
+302
to
+313
| def extract_credential_kwargs( | ||
| connection_string: str, auth_type: Optional[str] | ||
| ) -> Dict[str, str]: | ||
| """Extract credential constructor kwargs for the given auth type. | ||
|
|
||
| For ActiveDirectoryMSI: returns ``{"client_id": uid}`` when UID is | ||
| set (user-assigned MSI) and ``{}`` for system-assigned MSI. | ||
| """ | ||
| if auth_type != "msi": | ||
| return {} | ||
| client_id = _extract_msi_client_id(connection_string.split(";")) | ||
| return {"client_id": client_id} if client_id else {} |
Comment on lines
+528
to
+529
|
|
||
|
|
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.
What
Adds
Authentication=ActiveDirectoryMSIsupport to bulk copy. Partial fix for #534.ManagedIdentityCredential()for system-assigned MSIManagedIdentityCredential(client_id=UID)for user-assigned MSI — matches ODBC's convention whereUIDcarries the identity'sclient_idunder MSIcredential_kwargsthroughget_auth_token/get_raw_token/_acquire_tokenso future auth methods that need constructor args plug in via the same channel(auth_type, sorted_kwargs)when kwargs are present, so differentclient_ids get separate cached credentialsWhy
mssql-py-core(the Rust path used by bulk copy) doesn't itself acquire Entra tokens — Python has to pre-acquire a JWT and pass it asaccess_token. Today this works forDefault,DeviceCode, andInteractive. MSI was missing, which is the most common Azure-hosted-service auth method.ServicePrincipal and Password are explicitly out of scope for this PR — they need different design work.
Files
mssql_python/constants.pyAuthType.MSI = "activedirectorymsi"mssql_python/auth.pyManagedIdentityCredentialto credential map, threadcredential_kwargsthrough_acquire_token/get_token/get_raw_token/get_auth_token, add_extract_msi_client_idand publicextract_credential_kwargshelpers, update cache keymssql_python/cursor.pyclient_idfrom connection string and pass toget_raw_tokentests/test_008_auth.pyManagedIdentityCredentialin fixture, 10 new tests covering enum, parser, credential construction, cache isolation, and end-to-endprocess_connection_stringCHANGELOG.mdTest
10 new tests:
test_auth_type_constantsextended forAuthType.MSITestProcessAuthParameters::test_msi_auth,test_msi_auth_case_insensitiveTestExtractAuthType::test_msiTestManagedIdentity::test_get_token_system_assigned_msiTestManagedIdentity::test_get_raw_token_system_assigned_msiTestManagedIdentity::test_get_token_user_assigned_msiTestManagedIdentity::test_msi_separate_cache_entries_per_client_idTestManagedIdentity::test_extract_credential_kwargs_system_assigned/_user_assigned/_non_msi/_empty_uidTestManagedIdentity::test_process_connection_string_msi_strips_uidConnection string examples