[PECOBLR-587] Azure Service Principal Credential Provider #621

jprakash-db · 2025-06-30T18:31:12Z

Description

This pull request introduces support for Azure Service Principal M2M (SP) authentication to the PySQL Connector

Key Changes

AzureServicePrincipalCredentialProvider, this is the provider that will deal with getting the creds using the azure credentials
ClientCredentialsTokenSource, is responsible for managing the token lifecycle for the flow where credentials are obtained using the grant_type: client_credentials
Introduced a common HTTP client called DatabricksHttpClient that can be used to unify the Http logic across the connector, to ensure standard client level behaviour

New dependencies

Introduced library PyJWT which is required for handling JWT parsing functionalities

Tests

Expanded unit tests to cover:

AzureServicePrincipalCredentialProvider.
ClientCredentialsTokenSource

Manual Testing

Tested the Azure M2M SP flow by creating a SP using Azure Entra ID and then mapping it to a workspace SP. Further queries were made to the workspace using the Azure Entra ID credentials to test it on Databricks Workspace

github-actions · 2025-06-30T18:31:22Z

Thanks for your contribution! To satisfy the DCO policy in our contributing guide every commit message must include a sign-off message. One or more of your commits is missing this message. You can reword previous commit messages with an interactive rebase (git rebase -i main).

github-actions · 2025-07-01T08:19:29Z

Thanks for your contribution! To satisfy the DCO policy in our contributing guide every commit message must include a sign-off message. One or more of your commits is missing this message. You can reword previous commit messages with an interactive rebase (git rebase -i main).

github-actions · 2025-07-02T18:06:43Z

Thanks for your contribution! To satisfy the DCO policy in our contributing guide every commit message must include a sign-off message. One or more of your commits is missing this message. You can reword previous commit messages with an interactive rebase (git rebase -i main).

github-actions · 2025-07-03T05:14:00Z

Thanks for your contribution! To satisfy the DCO policy in our contributing guide every commit message must include a sign-off message. One or more of your commits is missing this message. You can reword previous commit messages with an interactive rebase (git rebase -i main).

github-actions · 2025-07-03T05:16:31Z

Thanks for your contribution! To satisfy the DCO policy in our contributing guide every commit message must include a sign-off message. One or more of your commits is missing this message. You can reword previous commit messages with an interactive rebase (git rebase -i main).

github-actions · 2025-07-03T05:18:13Z

Thanks for your contribution! To satisfy the DCO policy in our contributing guide every commit message must include a sign-off message. One or more of your commits is missing this message. You can reword previous commit messages with an interactive rebase (git rebase -i main).

github-actions · 2025-07-03T05:34:21Z

Thanks for your contribution! To satisfy the DCO policy in our contributing guide every commit message must include a sign-off message. One or more of your commits is missing this message. You can reword previous commit messages with an interactive rebase (git rebase -i main).

github-actions · 2025-07-03T05:37:40Z

Thanks for your contribution! To satisfy the DCO policy in our contributing guide every commit message must include a sign-off message. One or more of your commits is missing this message. You can reword previous commit messages with an interactive rebase (git rebase -i main).

github-actions · 2025-07-03T05:48:26Z

Thanks for your contribution! To satisfy the DCO policy in our contributing guide every commit message must include a sign-off message. One or more of your commits is missing this message. You can reword previous commit messages with an interactive rebase (git rebase -i main).

github-actions · 2025-07-03T05:51:34Z

Thanks for your contribution! To satisfy the DCO policy in our contributing guide every commit message must include a sign-off message. One or more of your commits is missing this message. You can reword previous commit messages with an interactive rebase (git rebase -i main).

vikrantpuppala · 2025-07-04T04:58:10Z

pyproject.toml


 [tool.poetry.extras]
 pyarrow = ["pyarrow"]

-[tool.poetry.dev-dependencies]
+[tool.poetry.group.dev.dependencies]


does the poetry update have to go in the same PR?

vikrantpuppala · 2025-07-04T05:03:02Z

src/databricks/sql/auth/auth.py

+
+    if auth_type == AuthType.AZURE_SP_M2M.value:
+        pass
+    else:


why not just do if if auth_type != AuthType.AZURE_SP_M2M.value:

vikrantpuppala · 2025-07-04T05:06:29Z

src/databricks/sql/auth/auth.py

        oauth_redirect_port_range=[kwargs["oauth_redirect_port"]]
-        if kwargs.get("oauth_client_id") and kwargs.get("oauth_redirect_port")
+        if client_id and kwargs.get("oauth_redirect_port")


this is behaviour change from before i.e. earlier we required client_id to come in from kwargs and now it is ok even if we derive it from get_client_id_and_redirect_port, is this intended?

vikrantpuppala · 2025-07-04T05:12:54Z

src/databricks/sql/auth/authenticators.py


 # Private API: this is an evolving interface and it will change in the future.
 # Please must not depend on it in your applications.
 from databricks.sql.experimental.oauth_persistence import OAuthToken, OAuthPersistence

+logger = logging.getLogger(__name__)


don't think this is being used, can we add logging?

vikrantpuppala · 2025-07-04T05:15:39Z

src/databricks/sql/auth/common.py

+            return app_id
+
+    # default databricks resource id
+    return "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d"


can we add these IDs as constants at the top so they're in one place

vikrantpuppala · 2025-07-04T05:17:45Z

src/databricks/sql/auth/oauth.py

+
+    @abstractmethod
+    def refresh(self) -> Token:
+        pass


can we add a comment here that we have duplicate code here with the sdk (with the code pointer) and in the long term we should try to unify?

vikrantpuppala · 2025-07-04T05:26:23Z

src/databricks/sql/auth/oauth.py

+            self.token = self.refresh()
+        return self.token
+
+    def refresh(self) -> Token:


how is the refresh mechanism being handled for the existing credential providers? is there opportunity to dedup/reuse?

vikrantpuppala · 2025-07-04T05:27:19Z

src/databricks/sql/common/http.py

+
+# Singleton class for common Http Client
+class DatabricksHttpClient:
+    ## TODO: Unify all the http clients in the PySQL Connector


why can't this be done right now? can we not use the existing http client?

Copilot

Pull Request Overview

This PR adds support for Azure Service Principal M2M authentication to the PySQL Connector by introducing a shared HTTP client, a client-credentials token source, and a new credentials provider for Azure SP.

Introduce DatabricksHttpClient for unified HTTP logic
Add Token and ClientCredentialsTokenSource to manage OAuth client-credentials flow
Implement AzureServicePrincipalCredentialProvider and wire it through get_auth_provider

Reviewed Changes

Copilot reviewed 8 out of 9 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
tests/unit/test_thrift_field_ids.py	Formatting cleanup and consistent quote style
tests/unit/test_auth.py	Added tests for `ClientCredentialsTokenSource` and SP credential provider, JWT fixtures
src/databricks/sql/common/http.py	New singleton `DatabricksHttpClient` with retry logic
src/databricks/sql/auth/oauth.py	Introduced `Token`, `RefreshableTokenSource`, and `ClientCredentialsTokenSource`
src/databricks/sql/auth/common.py	Extended `AuthType` and helper for mapping Azure login app IDs
src/databricks/sql/auth/authenticators.py	Added `AzureServicePrincipalCredentialProvider`
src/databricks/sql/auth/auth.py	Updated `ClientContext`, `get_auth_provider`, and auth provider resolution for SP
pyproject.toml	Added `pyjwt`, moved dev dependencies under `[tool.poetry.group.dev.dependencies]`

Copilot · 2025-07-04T05:30:18Z

src/databricks/sql/common/http.py

+        with self.session.request(method.value, url, **kwargs) as response:
+            return response


The requests.Response object is not a context manager. Remove the with and instead assign response = self.session.request(...) before returning it.

Suggested change

with self.session.request(method.value, url, **kwargs) as response:

return response

response = self.session.request(method.value, url, **kwargs)

return response

Copilot · 2025-07-04T05:30:18Z

src/databricks/sql/auth/oauth.py

+            return exp_time and (exp_time - buffer_time) <= current_time
+        except Exception as e:
+            logger.error("Failed to decode token: %s", e)
+            return e


is_expired should return a boolean, not an exception. Consider returning True on decode failure or re-raising the error to avoid returning a non-boolean value.

Suggested change

return e

return True

Copilot · 2025-07-04T05:30:19Z

tests/unit/test_auth.py

 from databricks.sql.auth.auth import (
    AccessTokenAuthProvider,
    AuthProvider,
    ExternalAuthProvider,
    AuthType,
 )
+import time
+from datetime import datetime, timedelta


[nitpick] The imports datetime and timedelta are not used in this test file. Consider removing them to keep imports clean.

Suggested change

from datetime import datetime, timedelta

Copilot · 2025-07-04T05:30:19Z

src/databricks/sql/auth/auth.py

+    redirect_port_range = kwargs.get("oauth_redirect_port_range")
+
+    if auth_type == AuthType.AZURE_SP_M2M.value:
+        pass


[nitpick] The pass under the Azure SP M2M branch appears to be a placeholder. Implement the handling for this branch or remove the pass if it is no longer needed.

Suggested change

pass

return ExternalAuthProvider(

AzureServicePrincipalCredentialProvider(

hostname=normalize_host_name(hostname),

oauth_client_id=kwargs.get("oauth_client_id"),

oauth_client_secret=kwargs.get("oauth_client_secret"),

azure_tenant_id=kwargs.get("azure_tenant_id"),

azure_workspace_resource_id=kwargs.get("azure_workspace_resource_id"),

)

)

basic setup

5e2b792

jprakash-db self-assigned this Jun 30, 2025

jprakash-db had a problem deploying to azure-prod June 30, 2025 18:31 — with GitHub Actions Failure

Nit

05ab3e8

jprakash-db had a problem deploying to azure-prod July 1, 2025 08:19 — with GitHub Actions Failure

working

bef7ac6

jprakash-db temporarily deployed to azure-prod July 2, 2025 18:06 — with GitHub Actions Inactive

moved pyjwt to code dependency

0877b6f

jprakash-db had a problem deploying to azure-prod July 3, 2025 05:13 — with GitHub Actions Failure

nit

83d6a9e

jprakash-db had a problem deploying to azure-prod July 3, 2025 05:16 — with GitHub Actions Failure

nit

b21d015

jprakash-db had a problem deploying to azure-prod July 3, 2025 05:18 — with GitHub Actions Failure

nit

6d85d19

jprakash-db temporarily deployed to azure-prod July 3, 2025 05:34 — with GitHub Actions Inactive

nit

f518085

jprakash-db had a problem deploying to azure-prod July 3, 2025 05:37 — with GitHub Actions Failure

nit

814d1cb

jprakash-db had a problem deploying to azure-prod July 3, 2025 05:47 — with GitHub Actions Failure

nit

ce4543c

jprakash-db temporarily deployed to azure-prod July 3, 2025 05:51 — with GitHub Actions Inactive

jprakash-db marked this pull request as ready for review July 3, 2025 05:55

jprakash-db requested review from deeksha-db, samikshya-db, jackyhu-db, madhav-db, gopalldb, jayantsing-db, vikrantpuppala and shivam2680 as code owners July 3, 2025 05:55

jprakash-db changed the title ~~Azure Service Principal Credential Provider~~ [PECOBLR-587] Azure Service Principal Credential Provider Jul 3, 2025

jprakash-db removed request for jackyhu-db, madhav-db, jayantsing-db and deeksha-db July 4, 2025 04:53

vikrantpuppala reviewed Jul 4, 2025

View reviewed changes

vikrantpuppala requested a review from Copilot July 4, 2025 05:28

Copilot AI reviewed Jul 4, 2025

View reviewed changes

		with self.session.request(method.value, url, **kwargs) as response:
		return response

-        pass
+        return ExternalAuthProvider(
+            AzureServicePrincipalCredentialProvider(
+                hostname=normalize_host_name(hostname),
+                oauth_client_id=kwargs.get("oauth_client_id"),
+                oauth_client_secret=kwargs.get("oauth_client_secret"),
+                azure_tenant_id=kwargs.get("azure_tenant_id"),
+                azure_workspace_resource_id=kwargs.get("azure_workspace_resource_id"),
+            )
+        )

[PECOBLR-587] Azure Service Principal Credential Provider #621

Are you sure you want to change the base?

[PECOBLR-587] Azure Service Principal Credential Provider #621

Conversation

jprakash-db commented Jun 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Key Changes

New dependencies

Tests

Uh oh!

github-actions bot commented Jun 30, 2025

Uh oh!

github-actions bot commented Jul 1, 2025

Uh oh!

github-actions bot commented Jul 2, 2025

Uh oh!

github-actions bot commented Jul 3, 2025

Uh oh!

github-actions bot commented Jul 3, 2025

Uh oh!

github-actions bot commented Jul 3, 2025

Uh oh!

github-actions bot commented Jul 3, 2025

Uh oh!

github-actions bot commented Jul 3, 2025

Uh oh!

github-actions bot commented Jul 3, 2025

Uh oh!

github-actions bot commented Jul 3, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Copilot AI Jul 4, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jul 4, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jul 4, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jul 4, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

jprakash-db commented Jun 30, 2025 •

edited

Loading