Skip to content

Add support for authentication through Azure MSI#760

Open
Divyansh-db wants to merge 2 commits intomainfrom
divyansh-vijayvergia_data/stack/azure-msi-support
Open

Add support for authentication through Azure MSI#760
Divyansh-db wants to merge 2 commits intomainfrom
divyansh-vijayvergia_data/stack/azure-msi-support

Conversation

@Divyansh-db
Copy link
Copy Markdown
Contributor

@Divyansh-db Divyansh-db commented Apr 15, 2026

Summary

Adds Azure Managed Service Identity (MSI) authentication to the Java SDK, porting the azure-msi credential strategy from the Go SDK (https://github.com/databricks/databricks-sdk-go/blob/main/config/auth_azure_msi.go).

Why

The Java SDK had the azureUseMsi config field (ARM_USE_MSI) wired up but no provider behind it. Users on Azure VMs, App Service, or Functions with managed identities had to fall back to service principal secrets or Azure CLI. The Go SDK has supported this since early on — this PR closes the gap.

What changed

Interface changes

  • AzureMsiCredentialsProvider — New CredentialsProvider with auth type "azure-msi". Activated when isAzure() &&
    ARM_USE_MSI=true && (host or resourceId is set).
  • AzureMsiTokenSource — TokenSource that fetches tokens from the Azure IMDS endpoint
    (http://169.254.169.254/metadata/identity/oauth2/token). Supports system-assigned and user-assigned identities (via ARM_CLIENT_ID).

Behavioral changes

  • Default credential chain now includes azure-msi between azure-github-oidc and azure-client-secret, matching Go SDK ordering.
  • Produces Authorization, X-Databricks-Azure-SP-Management-Token, and X-Databricks-Azure-Workspace-Resource-Id headers.

Internal changes

  • MsiTokenResponse POJO for IMDS JSON deserialization (mirrors Go's msiToken struct), handling expires_on (epoch
    seconds).
  • Null-safe isAzureUseMsi() helper for the boxed Boolean config field.
  • Reuses AzureUtils.ensureHostPresent() for workspace URL resolution, consistent with other Azure providers.

How is this tested?

20 unit tests in AzureMsiCredentialsProviderTest across two groups:

  • TokenSourceTests (12) — IMDS request shape, response parsing, table-driven client_id handling (3 cases),
    table-driven invalid responses (5 cases: missing fields, bad expiry), HTTP error paths.
  • ProviderTests (8) — Mirrors Go SDK's TestMsiHappyFlow, TestMsiHappyFlowWithHostAndNoResourceID,
    TestMsiFailsOnResolveWorkspace, TestMsiTokenNotFound, plus table-driven null-return preconditions.

@github-actions
Copy link
Copy Markdown
Contributor

If integration tests don't run automatically, an authorized user can run them manually by following the instructions below:

Trigger:
go/deco-tests-run/sdk-java

Inputs:

  • PR number: 760
  • Commit SHA: 222a66e1ee3114585b4094f74264f7db521802db

Checks will be approved automatically on success.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant