Skip to content

Add support for authentication through Azure MSI#760

Merged
Divyansh-db merged 4 commits intomainfrom
divyansh-vijayvergia_data/stack/azure-msi-support
Apr 17, 2026
Merged

Add support for authentication through Azure MSI#760
Divyansh-db merged 4 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.
  • Verified it works against a real Azure VM

private static boolean isAzureUseMsi(DatabricksConfig config) {
try {
return config.getAzureUseMsi();
} catch (NullPointerException e) {
Copy link
Copy Markdown
Contributor

@hectorcast-db hectorcast-db Apr 17, 2026

Choose a reason for hiding this comment

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

Causing a null-pointer on purpose and using a catch is very unorthodox. Better use:

Boolean flag = config.getAzureUseMsi();
return flag != null && flag;

Copy link
Copy Markdown
Contributor

@hectorcast-db hectorcast-db left a comment

Choose a reason for hiding this comment

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

Approved but change the issue above

@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: d13f052ed6e33c6bb04ef6115eb553753ecc5fc9

Checks will be approved automatically on success.

@Divyansh-db Divyansh-db enabled auto-merge April 17, 2026 10:41
@Divyansh-db Divyansh-db added this pull request to the merge queue Apr 17, 2026
Merged via the queue into main with commit a1e7fe6 Apr 17, 2026
15 of 17 checks passed
@Divyansh-db Divyansh-db deleted the divyansh-vijayvergia_data/stack/azure-msi-support branch April 17, 2026 10:46
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.

2 participants