Add HostType to HostMetadata

Signed-off-by: Tanmay Rustagi <tanmay.rustagi@databricks.com>

Author: tanmay-db

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

@@ -162,6 +162,12 @@ public class DatabricksConfig {
 
   private DatabricksEnvironment databricksEnvironment;
 
+  /**
+   * The host type resolved from the /.well-known/databricks-config discovery endpoint. When set,
+   * this takes priority over URL-based host type detection in {@link #getHostType()}.
+   */
+  private HostType resolvedHostType;
+
   /**
    * When using Workload Identity Federation, the audience to specify when fetching an ID token from
    * the ID token supplier.
@@ -723,6 +729,17 @@ public DatabricksConfig setDisableOauthRefreshToken(boolean disable) {
     return this;
   }
 
+  /** Returns the host type resolved from host metadata, or {@code null} if not yet resolved. */
+  HostType getResolvedHostType() {
+    return resolvedHostType;
+  }
+
+  /** Sets the resolved host type. Package-private for testing. */
+  DatabricksConfig setResolvedHostType(HostType resolvedHostType) {
+    this.resolvedHostType = resolvedHostType;
+    return this;
+  }
+
   public boolean isAzure() {
     if (azureWorkspaceResourceId != null) {
       return true;
@@ -889,6 +906,13 @@ void resolveHostMetadata() throws IOException {
       LOG.debug("Resolved workspace_id from host metadata: \"{}\"", meta.getWorkspaceId());
       workspaceId = meta.getWorkspaceId();
     }
+    if (resolvedHostType == null && meta.getHostType() != null) {
+      HostType ht = HostType.fromApiValue(meta.getHostType());
+      if (ht != null) {
+        LOG.debug("Resolved host_type from host metadata: \"{}\"", ht);
+        resolvedHostType = ht;
+      }
+    }
     if (discoveryUrl == null) {
       if (meta.getOidcEndpoint() == null || meta.getOidcEndpoint().isEmpty()) {
         LOG.warn("Host metadata missing oidc_endpoint; skipping discovery URL resolution");

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

@@ -12,5 +12,25 @@ public enum HostType {
   ACCOUNTS,
 
   /** Unified host supporting both workspace and account operations. */
-  UNIFIED
+  UNIFIED;
+
+  /**
+   * Converts an API-level host type string (e.g. "workspace", "account", "unified") to the
+   * corresponding enum value. Returns {@code null} for unknown or empty values.
+   */
+  public static HostType fromApiValue(String value) {
+    if (value == null || value.isEmpty()) {
+      return null;
+    }
+    switch (value.toLowerCase()) {
+      case "workspace":
+        return WORKSPACE;
+      case "account":
+        return ACCOUNTS;
+      case "unified":
+        return UNIFIED;
+      default:
+        return null;
+    }
+  }
 }

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/HostMetadata.java

@@ -23,6 +23,9 @@ public class HostMetadata {
   @JsonProperty("cloud")
   private String cloud;
 
+  @JsonProperty("host_type")
+  private String hostType;
+
   public HostMetadata() {}
 
   public HostMetadata(String oidcEndpoint, String accountId, String workspaceId) {
@@ -53,4 +56,8 @@ public String getWorkspaceId() {
   public String getCloud() {
     return cloud;
   }
+
+  public String getHostType() {
+    return hostType;
+  }
 }

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

@@ -712,6 +712,94 @@ public void testEnsureResolvedHostMetadataMissingAccountIdWithPlaceholderNonFata
     }
   }
 
+  // --- resolveHostMetadata host type tests ---
+
+  @Test
+  public void testResolveHostMetadataPopulatesResolvedHostType() throws IOException {
+    String response =
+        "{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\","
+            + "\"account_id\":\""
+            + DUMMY_ACCOUNT_ID
+            + "\","
+            + "\"host_type\":\"workspace\"}";
+    try (FixtureServer server =
+        new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
+      DatabricksConfig config = new DatabricksConfig().setHost(server.getUrl());
+      config.resolve(emptyEnv());
+      config.resolveHostMetadata();
+      assertEquals(HostType.WORKSPACE, config.getResolvedHostType());
+    }
+  }
+
+  @Test
+  public void testResolveHostMetadataDoesNotOverwriteExistingHostType() throws IOException {
+    String response =
+        "{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\","
+            + "\"account_id\":\""
+            + DUMMY_ACCOUNT_ID
+            + "\","
+            + "\"host_type\":\"workspace\"}";
+    try (FixtureServer server =
+        new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
+      DatabricksConfig config = new DatabricksConfig().setHost(server.getUrl());
+      config.resolve(emptyEnv());
+      config.setResolvedHostType(HostType.UNIFIED);
+      config.resolveHostMetadata();
+      assertEquals(HostType.UNIFIED, config.getResolvedHostType());
+    }
+  }
+
+  @Test
+  public void testResolveHostMetadataUnknownHostTypeIgnored() throws IOException {
+    String response =
+        "{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\","
+            + "\"account_id\":\""
+            + DUMMY_ACCOUNT_ID
+            + "\","
+            + "\"host_type\":\"unknown_value\"}";
+    try (FixtureServer server =
+        new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
+      DatabricksConfig config = new DatabricksConfig().setHost(server.getUrl());
+      config.resolve(emptyEnv());
+      config.resolveHostMetadata();
+      assertNull(config.getResolvedHostType());
+    }
+  }
+
+  @Test
+  public void testResolveHostMetadataHostTypeAccount() throws IOException {
+    String response =
+        "{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\","
+            + "\"account_id\":\""
+            + DUMMY_ACCOUNT_ID
+            + "\","
+            + "\"host_type\":\"account\"}";
+    try (FixtureServer server =
+        new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
+      DatabricksConfig config = new DatabricksConfig().setHost(server.getUrl());
+      config.resolve(emptyEnv());
+      config.resolveHostMetadata();
+      assertEquals(HostType.ACCOUNTS, config.getResolvedHostType());
+    }
+  }
+
+  @Test
+  public void testResolveHostMetadataHostTypeUnified() throws IOException {
+    String response =
+        "{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\","
+            + "\"account_id\":\""
+            + DUMMY_ACCOUNT_ID
+            + "\","
+            + "\"host_type\":\"unified\"}";
+    try (FixtureServer server =
+        new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
+      DatabricksConfig config = new DatabricksConfig().setHost(server.getUrl());
+      config.resolve(emptyEnv());
+      config.resolveHostMetadata();
+      assertEquals(HostType.UNIFIED, config.getResolvedHostType());
+    }
+  }
+
   // --- discoveryUrl / OIDC endpoint tests ---
 
   @Test

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

@@ -0,0 +1,46 @@
+package com.databricks.sdk.core;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+import org.junit.jupiter.api.Test;
+
+public class HostTypeTest {
+
+  @Test
+  public void testFromApiValueWorkspace() {
+    assertEquals(HostType.WORKSPACE, HostType.fromApiValue("workspace"));
+  }
+
+  @Test
+  public void testFromApiValueAccount() {
+    assertEquals(HostType.ACCOUNTS, HostType.fromApiValue("account"));
+  }
+
+  @Test
+  public void testFromApiValueUnified() {
+    assertEquals(HostType.UNIFIED, HostType.fromApiValue("unified"));
+  }
+
+  @Test
+  public void testFromApiValueCaseInsensitive() {
+    assertEquals(HostType.WORKSPACE, HostType.fromApiValue("WORKSPACE"));
+    assertEquals(HostType.ACCOUNTS, HostType.fromApiValue("Account"));
+    assertEquals(HostType.UNIFIED, HostType.fromApiValue("UNIFIED"));
+  }
+
+  @Test
+  public void testFromApiValueNull() {
+    assertNull(HostType.fromApiValue(null));
+  }
+
+  @Test
+  public void testFromApiValueEmpty() {
+    assertNull(HostType.fromApiValue(""));
+  }
+
+  @Test
+  public void testFromApiValueUnknown() {
+    assertNull(HostType.fromApiValue("unknown"));
+    assertNull(HostType.fromApiValue("something_else"));
+  }
+}