Skip to content

Commit 733d3ae

Browse files
committed
Wire HostMetadata discovery into DatabricksConfig
Adds package-private `resolveHostMetadata()` to `DatabricksConfig`. When called, it fetches `/.well-known/databricks-config` and back-fills `accountId`, `workspaceId`, and `discoveryUrl` (with any `{account_id}` placeholder substituted) if those fields are not already set. Raises DatabricksException if `accountId` cannot be resolved or if no `oidc_endpoint` is present in the metadata. Mirrors `Config._resolve_host_metadata()` in the Python SDK.
1 parent 3b4e6d8 commit 733d3ae

File tree

2 files changed

+164
-0
lines changed

2 files changed

+164
-0
lines changed

databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -827,6 +827,71 @@ HostMetadata getHostMetadata() throws IOException {
827827
}
828828
}
829829

830+
/**
831+
* [Experimental] Fetch the raw Databricks well-known configuration from
832+
* {host}/.well-known/databricks-config.
833+
*
834+
* <p><b>Note:</b> This API is experimental and may change or be removed in future releases
835+
* without notice.
836+
*
837+
* @return Parsed {@link HostMetadata} as returned by the server.
838+
* @throws DatabricksException if the request fails or the server returns a non-200 status.
839+
*/
840+
HostMetadata getHostMetadata() throws IOException {
841+
String url = host + "/.well-known/databricks-config";
842+
try {
843+
Request request = new Request("GET", url);
844+
Response resp = getHttpClient().execute(request);
845+
if (resp.getStatusCode() != 200) {
846+
throw new DatabricksException(
847+
"Failed to fetch host metadata from " + url + ": HTTP " + resp.getStatusCode());
848+
}
849+
return new ObjectMapper().readValue(resp.getBody(), HostMetadata.class);
850+
} catch (IOException e) {
851+
throw new DatabricksException(
852+
"Failed to fetch host metadata from " + url + ": " + e.getMessage(), e);
853+
}
854+
}
855+
856+
/**
857+
* [Experimental] Populate missing config fields from the host's /.well-known/databricks-config
858+
* discovery endpoint.
859+
*
860+
* <p>Fills in {@code accountId}, {@code workspaceId}, and {@code discoveryUrl} (derived from
861+
* {@code oidc_endpoint}, with any {@code {account_id}} placeholder substituted) if not already
862+
* set.
863+
*
864+
* <p><b>Note:</b> This API is experimental and may change or be removed in future releases
865+
* without notice.
866+
*
867+
* @throws DatabricksException if {@code accountId} cannot be resolved or {@code oidc_endpoint} is
868+
* missing from the host metadata.
869+
*/
870+
void resolveHostMetadata() throws IOException {
871+
if (host == null) {
872+
return;
873+
}
874+
HostMetadata meta = getHostMetadata();
875+
if (accountId == null && meta.getAccountId() != null) {
876+
accountId = meta.getAccountId();
877+
}
878+
if (accountId == null) {
879+
throw new DatabricksException(
880+
"account_id is not configured and could not be resolved from host metadata");
881+
}
882+
if (workspaceId == null && meta.getWorkspaceId() != null) {
883+
workspaceId = meta.getWorkspaceId();
884+
}
885+
if (discoveryUrl == null) {
886+
if (meta.getOidcEndpoint() != null && !meta.getOidcEndpoint().isEmpty()) {
887+
discoveryUrl = meta.getOidcEndpoint().replace("{account_id}", accountId);
888+
} else {
889+
throw new DatabricksException(
890+
"discovery_url is not configured and could not be resolved from host metadata");
891+
}
892+
}
893+
}
894+
830895
private OpenIDConnectEndpoints fetchOidcEndpointsFromDiscovery() {
831896
try {
832897
Request request = new Request("GET", discoveryUrl);

databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -480,6 +480,105 @@ public void testGetHostMetadataRaisesOnHttpError() throws IOException {
480480
}
481481
}
482482

483+
// --- resolveHostMetadata tests ---
484+
485+
@Test
486+
public void testResolveHostMetadataWorkspacePopulatesAllFields() throws IOException {
487+
String response =
488+
"{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\","
489+
+ "\"account_id\":\""
490+
+ DUMMY_ACCOUNT_ID
491+
+ "\","
492+
+ "\"workspace_id\":\""
493+
+ DUMMY_WORKSPACE_ID
494+
+ "\"}";
495+
try (FixtureServer server =
496+
new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
497+
DatabricksConfig config = new DatabricksConfig().setHost(server.getUrl());
498+
config.resolve(emptyEnv());
499+
config.resolveHostMetadata();
500+
assertEquals(DUMMY_ACCOUNT_ID, config.getAccountId());
501+
assertEquals(DUMMY_WORKSPACE_ID, config.getWorkspaceId());
502+
assertEquals("https://ws.databricks.com/oidc", config.getDiscoveryUrl());
503+
}
504+
}
505+
506+
@Test
507+
public void testResolveHostMetadataAccountSubstitutesAccountId() throws IOException {
508+
String response =
509+
"{\"oidc_endpoint\":\"https://acc.databricks.com/oidc/accounts/{account_id}\"}";
510+
try (FixtureServer server =
511+
new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
512+
DatabricksConfig config =
513+
new DatabricksConfig().setHost(server.getUrl()).setAccountId(DUMMY_ACCOUNT_ID);
514+
config.resolve(emptyEnv());
515+
config.resolveHostMetadata();
516+
assertEquals(
517+
"https://acc.databricks.com/oidc/accounts/" + DUMMY_ACCOUNT_ID, config.getDiscoveryUrl());
518+
}
519+
}
520+
521+
@Test
522+
public void testResolveHostMetadataDoesNotOverwriteExistingFields() throws IOException {
523+
String existingAccountId = "existing-account-id";
524+
String existingWorkspaceId = "existing-workspace-id";
525+
String response =
526+
"{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\","
527+
+ "\"account_id\":\"other-account\","
528+
+ "\"workspace_id\":\"other-ws\"}";
529+
try (FixtureServer server =
530+
new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
531+
DatabricksConfig config =
532+
new DatabricksConfig()
533+
.setHost(server.getUrl())
534+
.setAccountId(existingAccountId)
535+
.setWorkspaceId(existingWorkspaceId);
536+
config.resolve(emptyEnv());
537+
config.resolveHostMetadata();
538+
assertEquals(existingAccountId, config.getAccountId());
539+
assertEquals(existingWorkspaceId, config.getWorkspaceId());
540+
}
541+
}
542+
543+
@Test
544+
public void testResolveHostMetadataRaisesWhenAccountIdUnresolvable() throws IOException {
545+
String response =
546+
"{\"oidc_endpoint\":\"https://acc.databricks.com/oidc/accounts/{account_id}\"}";
547+
try (FixtureServer server =
548+
new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
549+
DatabricksConfig config = new DatabricksConfig().setHost(server.getUrl());
550+
config.resolve(emptyEnv());
551+
DatabricksException ex =
552+
assertThrows(DatabricksException.class, () -> config.resolveHostMetadata());
553+
assertTrue(ex.getMessage().contains("account_id is not configured"));
554+
}
555+
}
556+
557+
@Test
558+
public void testResolveHostMetadataRaisesWhenOidcEndpointMissing() throws IOException {
559+
String response = "{\"account_id\":\"" + DUMMY_ACCOUNT_ID + "\"}";
560+
try (FixtureServer server =
561+
new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
562+
DatabricksConfig config = new DatabricksConfig().setHost(server.getUrl());
563+
config.resolve(emptyEnv());
564+
DatabricksException ex =
565+
assertThrows(DatabricksException.class, () -> config.resolveHostMetadata());
566+
assertTrue(ex.getMessage().contains("discovery_url is not configured"));
567+
}
568+
}
569+
570+
@Test
571+
public void testResolveHostMetadataRaisesOnHttpError() throws IOException {
572+
try (FixtureServer server =
573+
new FixtureServer().with("GET", "/.well-known/databricks-config", "{}", 500)) {
574+
DatabricksConfig config = new DatabricksConfig().setHost(server.getUrl());
575+
config.resolve(emptyEnv());
576+
DatabricksException ex =
577+
assertThrows(DatabricksException.class, () -> config.resolveHostMetadata());
578+
assertTrue(ex.getMessage().contains("Failed to fetch host metadata"));
579+
}
580+
}
581+
483582
// --- discoveryUrl / OIDC endpoint tests ---
484583

485584
@Test

0 commit comments

Comments
 (0)