Skip to content

Commit 2ec6df9

Browse files
Divyansh-dbDivyansh-db
committed
add logging and demo file
1 parent 222a66e commit 2ec6df9

File tree

2 files changed

+165
-0
lines changed

2 files changed

+165
-0
lines changed

‎databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,15 +66,19 @@ public synchronized HeaderFactory configure(DatabricksConfig config) {
6666
&& !provider.authType().equals(config.getAuthType())) {
6767
LOG.info(
6868
"Ignoring {} auth, because {} is preferred", provider.authType(), config.getAuthType());
69+
System.out.println("[credential-chain] SKIP " + provider.authType() + " (not preferred)");
6970
continue;
7071
}
7172
try {
7273
LOG.info("Trying {} auth", provider.authType());
74+
System.out.println("[credential-chain] TRYING " + provider.authType() + " ...");
7375
HeaderFactory headerFactory = provider.configure(config);
7476
if (headerFactory == null) {
77+
System.out.println("[credential-chain] " + provider.authType() + " -> not applicable");
7578
continue;
7679
}
7780
authType = provider.authType();
81+
System.out.println("[credential-chain] " + provider.authType() + " -> SUCCESS ✓");
7882
return headerFactory;
7983
} catch (DatabricksException e) {
8084
throw new DatabricksException(

‎examples/docs/src/main/java/com/databricks/example/AzureMsiAuthDemo.java‎

Lines changed: 161 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,161 @@
1+
package com.databricks.example;
2+
3+
import com.databricks.sdk.AccountClient;
4+
import com.databricks.sdk.WorkspaceClient;
5+
import com.databricks.sdk.core.DatabricksConfig;
6+
import com.databricks.sdk.service.compute.ClusterDetails;
7+
import com.databricks.sdk.service.compute.ListClustersRequest;
8+
import com.databricks.sdk.service.iam.Group;
9+
import com.databricks.sdk.service.iam.ListAccountGroupsRequest;
10+
import com.fasterxml.jackson.core.type.TypeReference;
11+
import com.fasterxml.jackson.databind.ObjectMapper;
12+
import java.io.InputStream;
13+
import java.nio.file.Files;
14+
import java.nio.file.Paths;
15+
import java.util.Map;
16+
17+
/**
18+
* Demo for Azure MSI authentication.
19+
*
20+
* <p>This example loads environment variables from ~/.databricks/debug-env.json (the "workspace" or
21+
* "account" section) and authenticates using Azure Managed Service Identity.
22+
*
23+
* <p>Prerequisites:
24+
*
25+
* <ul>
26+
* <li>Running on an Azure VM (or compute) with a managed identity assigned
27+
* <li>~/.databricks/debug-env.json configured with the appropriate section
28+
* </ul>
29+
*
30+
* <p>Expected debug-env.json format:
31+
*
32+
* <pre>{@code
33+
* {
34+
* "workspace": {
35+
* "DATABRICKS_HOST": "https://adb-xxxx.azuredatabricks.net",
36+
* "ARM_USE_MSI": "true",
37+
* "ARM_CLIENT_ID": "<optional-user-assigned-identity-client-id>"
38+
* },
39+
* "account": {
40+
* "DATABRICKS_HOST": "https://accounts.azuredatabricks.net",
41+
* "DATABRICKS_ACCOUNT_ID": "<account-id>",
42+
* "DATABRICKS_AZURE_RESOURCE_ID": "/subscriptions/.../workspaces/...",
43+
* "ARM_USE_MSI": "true",
44+
* "ARM_CLIENT_ID": "<optional-user-assigned-identity-client-id>"
45+
* }
46+
* }
47+
* }</pre>
48+
*
49+
* <p>Usage: Comment/uncomment the workspace or account demo in main() as needed.
50+
*/
51+
public class AzureMsiAuthDemo {
52+
53+
public static void main(String[] args) throws Exception {
54+
// ── Uncomment ONE of the following ──
55+
56+
workspaceDemo();
57+
// accountDemo();
58+
}
59+
60+
/** Authenticate via Azure MSI and list clusters in the workspace. */
61+
static void workspaceDemo() throws Exception {
62+
Map<String, String> env = loadDebugEnv("workspace");
63+
64+
DatabricksConfig config =
65+
new DatabricksConfig()
66+
.setHost(env.get("DATABRICKS_HOST"))
67+
.setAzureUseMsi(true)
68+
.setAuthType("azure-msi");
69+
70+
// Set optional user-assigned identity client ID.
71+
if (env.containsKey("ARM_CLIENT_ID")) {
72+
config.setAzureClientId(env.get("ARM_CLIENT_ID"));
73+
}
74+
75+
System.out.println("=== Azure MSI Workspace Demo ===");
76+
System.out.println("Host: " + config.getHost());
77+
System.out.println("Auth type: azure-msi");
78+
System.out.println();
79+
80+
WorkspaceClient workspace = new WorkspaceClient(config);
81+
82+
// This call will trigger authentication — the credential chain runs here.
83+
System.out.println("Authenticating and listing clusters...");
84+
System.out.println();
85+
86+
int count = 0;
87+
for (ClusterDetails c : workspace.clusters().list(new ListClustersRequest())) {
88+
System.out.println(" Cluster: " + c.getClusterName() + " (" + c.getState() + ")");
89+
count++;
90+
}
91+
System.out.println();
92+
System.out.println("Total clusters: " + count);
93+
System.out.println("Auth succeeded with type: " + config.getAuthType());
94+
}
95+
96+
/** Authenticate via Azure MSI and list groups in the account. */
97+
static void accountDemo() throws Exception {
98+
Map<String, String> env = loadDebugEnv("account");
99+
100+
DatabricksConfig config =
101+
new DatabricksConfig()
102+
.setHost(env.get("DATABRICKS_HOST"))
103+
.setAccountId(env.get("DATABRICKS_ACCOUNT_ID"))
104+
.setAzureUseMsi(true)
105+
.setAuthType("azure-msi");
106+
107+
if (env.containsKey("DATABRICKS_AZURE_RESOURCE_ID")) {
108+
config.setAzureWorkspaceResourceId(env.get("DATABRICKS_AZURE_RESOURCE_ID"));
109+
}
110+
if (env.containsKey("ARM_CLIENT_ID")) {
111+
config.setAzureClientId(env.get("ARM_CLIENT_ID"));
112+
}
113+
114+
System.out.println("=== Azure MSI Account Demo ===");
115+
System.out.println("Host: " + config.getHost());
116+
System.out.println("Account ID: " + config.getAccountId());
117+
System.out.println("Auth type: azure-msi");
118+
System.out.println();
119+
120+
AccountClient account = new AccountClient(config);
121+
122+
System.out.println("Authenticating and listing account groups...");
123+
System.out.println();
124+
125+
int count = 0;
126+
for (Group g : account.groups().list(new ListAccountGroupsRequest())) {
127+
System.out.println(" Group: " + g.getDisplayName());
128+
count++;
129+
}
130+
System.out.println();
131+
System.out.println("Total groups: " + count);
132+
System.out.println("Auth succeeded with type: " + config.getAuthType());
133+
}
134+
135+
/**
136+
* Loads environment variables from ~/.databricks/debug-env.json for the given section.
137+
*
138+
* @param section The section name (e.g., "workspace", "account")
139+
* @return Map of environment variable names to values
140+
*/
141+
private static Map<String, String> loadDebugEnv(String section) throws Exception {
142+
String path =
143+
String.format("%s/.databricks/debug-env.json", System.getProperty("user.home"));
144+
ObjectMapper mapper = new ObjectMapper();
145+
try (InputStream in = Files.newInputStream(Paths.get(path))) {
146+
Map<String, Map<String, String>> all =
147+
mapper.readValue(in, new TypeReference<Map<String, Map<String, String>>>() {});
148+
Map<String, String> env = all.get(section);
149+
if (env == null) {
150+
throw new RuntimeException(
151+
"Section '"
152+
+ section
153+
+ "' not found in "
154+
+ path
155+
+ ". Available: "
156+
+ all.keySet());
157+
}
158+
return env;
159+
}
160+
}
161+
}

0 commit comments

Comments
 (0)