Skip to content

Commit aa1af2f

Browse files
tanmay-dbtanmay-db
authored
[Internal] Add HostType to HostMetadata (#1596)
## Summary Parses the `host_type` field from the `/.well-known/databricks-config` discovery endpoint into `HostMetadata` and stores it on `Config`, so the SDK knows whether a host is a workspace, account, or unified host without parsing hostnames. ## Why Today, `HostType()` infers the host type by prefix-matching the hostname against `accounts.` and `accounts-dod.` — everything else is classified as `WorkspaceHost`. This cannot detect unified hosts, which share the same hostname pattern as workspace hosts. The server-side `host_type` field is being added behind a feature flag to the `/.well-known/databricks-config` endpoint, which returns `"workspace"`, `"account"`, or `"unified"`. This PR threads that value into the SDK so it's available for future use when `HostType()` is updated to prefer endpoint-resolved values over URL-based inference. ## What changed ### Interface changes - **`HostTypeUnknown HostType = ""`** — new sentinel constant for an unset/empty host type, analogous to `environment.CloudUnknown`. - **`HostType.UnmarshalJSON`** — case-insensitive JSON deserialization that normalizes API values (`workspace`, `account`, `unified`) to the canonical SDK constants (`WORKSPACE_HOST`, `ACCOUNT_HOST`, `UNIFIED_HOST`). Unrecognized values return `HostTypeUnknown`. - **`HostMetadata.HostType HostType`** — new field parsed from the `host_type` key in the `/.well-known/databricks-config` response. ### Behavioral changes - `resolveHostMetadata` now populates `Config.resolvedHostType` from the discovery endpoint if not already set. `HostType()` is **not yet modified** to use this value — that will be a follow-up change once the server-side flag is rolled out. ### Internal changes - `Config.resolvedHostType` (unexported) stores the host type resolved during `resolveHostMetadata`. An exported field cannot be used because Go does not allow a struct field and method with the same name. ## How is this tested? Unit tests covering: - `TestConfig_ResolveHostMetadata_PopulatesHostTypeFromAPI` — resolved from endpoint response. - `TestConfig_ResolveHostMetadata_HostTypeEmptyWhenNotInResponse` — remains `HostTypeUnknown` when endpoint omits `host_type`. - `TestConfig_ResolveHostMetadata_DoesNotOverwriteExistingHostType` — existing value is not overwritten. - `TestConfig_ResolveHostMetadata_HostTypes` — table-driven: all case variants, uppercase, unknown values, empty. - `TestGetHostMetadata_WithHostTypeField` — `HostMetadata` deserialization for all API values and missing field. NO_CHANGELOG=true
1 parent d64a777 commit aa1af2f

4 files changed

Lines changed: 228 additions & 0 deletions

File tree

‎config/config.go‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package config
22

33
import (
44
"context"
5+
"encoding/json"
56
"errors"
67
"fmt"
78
"net/http"
@@ -50,8 +51,33 @@ const (
5051
AccountHost HostType = "ACCOUNT_HOST"
5152
// UnifiedHost supports both workspace-level and account-level APIs.
5253
UnifiedHost HostType = "UNIFIED_HOST"
54+
// HostTypeUnknown is the zero value for an unset/empty host type.
55+
HostTypeUnknown HostType = ""
5356
)
5457

58+
// UnmarshalJSON automatically normalizes the host type string when parsing JSON.
59+
func (h *HostType) UnmarshalJSON(data []byte) error {
60+
var rawString string
61+
if err := json.Unmarshal(data, &rawString); err != nil {
62+
return err
63+
}
64+
*h = normalizeHostType(rawString)
65+
return nil
66+
}
67+
68+
func normalizeHostType(hostType string) HostType {
69+
switch strings.ToLower(hostType) {
70+
case "workspace":
71+
return WorkspaceHost
72+
case "account":
73+
return AccountHost
74+
case "unified":
75+
return UnifiedHost
76+
default:
77+
return HostTypeUnknown
78+
}
79+
}
80+
5581
// ConfigType represents the type of API this config is valid for.
5682
type ConfigType string
5783

@@ -222,6 +248,10 @@ type Config struct {
222248

223249
Loaders []Loader
224250

251+
// resolvedHostType is the host type resolved from the /.well-known/databricks-config
252+
// discovery endpoint.
253+
resolvedHostType HostType
254+
225255
// marker for configuration resolving
226256
resolved bool
227257

@@ -702,6 +732,10 @@ func (c *Config) resolveHostMetadata(ctx context.Context) {
702732
c.Cloud = c.Environment().Cloud
703733
logger.Debugf(ctx, "Resolved cloud from hostname: %q", c.Cloud)
704734
}
735+
if c.resolvedHostType == HostTypeUnknown && meta.HostType != HostTypeUnknown {
736+
logger.Debugf(ctx, "Resolved host_type from host metadata: %q", meta.HostType)
737+
c.resolvedHostType = meta.HostType
738+
}
705739
if c.TokenAudience == "" && meta.DefaultOIDCAudience != "" {
706740
logger.Debugf(ctx, "Resolved token_audience from host metadata default_oidc_audience: %q", meta.DefaultOIDCAudience)
707741
c.TokenAudience = meta.DefaultOIDCAudience

‎config/config_test.go‎

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1012,3 +1012,134 @@ func TestConfig_ResolveHostMetadata_Clouds(t *testing.T) {
10121012
})
10131013
}
10141014
}
1015+
1016+
func TestConfig_ResolveHostMetadata_PopulatesHostTypeFromAPI(t *testing.T) {
1017+
noopLoader := mockLoader(func(cfg *Config) error { return nil })
1018+
cfg := &Config{
1019+
Host: testHMHost,
1020+
Loaders: []Loader{noopLoader},
1021+
HTTPTransport: fixtures.SliceTransport{
1022+
{
1023+
Method: "GET",
1024+
Resource: "/.well-known/databricks-config",
1025+
ReuseRequest: true,
1026+
Status: 200,
1027+
Response: `{"oidc_endpoint": "` + testHMHost + `/oidc", "account_id": "` + testHMAccountID + `", "host_type": "unified"}`,
1028+
},
1029+
},
1030+
}
1031+
err := cfg.EnsureResolved()
1032+
require.NoError(t, err)
1033+
assert.Equal(t, UnifiedHost, cfg.resolvedHostType)
1034+
}
1035+
1036+
func TestConfig_ResolveHostMetadata_HostTypeEmptyWhenNotInResponse(t *testing.T) {
1037+
noopLoader := mockLoader(func(cfg *Config) error { return nil })
1038+
cfg := &Config{
1039+
Host: "https://accounts.cloud.databricks.com",
1040+
Loaders: []Loader{noopLoader},
1041+
HTTPTransport: fixtures.SliceTransport{
1042+
{
1043+
Method: "GET",
1044+
Resource: "/.well-known/databricks-config",
1045+
ReuseRequest: true,
1046+
Status: 200,
1047+
Response: `{"oidc_endpoint": "https://accounts.cloud.databricks.com/oidc", "account_id": "` + testHMAccountID + `"}`,
1048+
},
1049+
},
1050+
}
1051+
err := cfg.EnsureResolved()
1052+
require.NoError(t, err)
1053+
assert.Equal(t, HostTypeUnknown, cfg.resolvedHostType)
1054+
}
1055+
1056+
func TestConfig_ResolveHostMetadata_DoesNotOverwriteExistingHostType(t *testing.T) {
1057+
noopLoader := mockLoader(func(cfg *Config) error { return nil })
1058+
cfg := &Config{
1059+
Host: testHMHost,
1060+
Loaders: []Loader{noopLoader},
1061+
HTTPTransport: fixtures.SliceTransport{
1062+
{
1063+
Method: "GET",
1064+
Resource: "/.well-known/databricks-config",
1065+
ReuseRequest: true,
1066+
Status: 200,
1067+
Response: `{"oidc_endpoint": "` + testHMHost + `/oidc", "account_id": "` + testHMAccountID + `", "host_type": "account"}`,
1068+
},
1069+
},
1070+
}
1071+
cfg.resolvedHostType = UnifiedHost
1072+
err := cfg.EnsureResolved()
1073+
require.NoError(t, err)
1074+
assert.Equal(t, UnifiedHost, cfg.resolvedHostType)
1075+
}
1076+
1077+
func TestConfig_ResolveHostMetadata_HostTypes(t *testing.T) {
1078+
tests := []struct {
1079+
name string
1080+
hostTypeJSON string
1081+
wantHostType string
1082+
}{
1083+
{
1084+
name: "workspace",
1085+
hostTypeJSON: "workspace",
1086+
wantHostType: "WORKSPACE_HOST",
1087+
},
1088+
{
1089+
name: "account",
1090+
hostTypeJSON: "account",
1091+
wantHostType: "ACCOUNT_HOST",
1092+
},
1093+
{
1094+
name: "unified",
1095+
hostTypeJSON: "unified",
1096+
wantHostType: "UNIFIED_HOST",
1097+
},
1098+
{
1099+
name: "Workspace uppercase",
1100+
hostTypeJSON: "WORKSPACE",
1101+
wantHostType: "WORKSPACE_HOST",
1102+
},
1103+
{
1104+
name: "Account uppercase",
1105+
hostTypeJSON: "ACCOUNT",
1106+
wantHostType: "ACCOUNT_HOST",
1107+
},
1108+
{
1109+
name: "Unified uppercase",
1110+
hostTypeJSON: "UNIFIED",
1111+
wantHostType: "UNIFIED_HOST",
1112+
},
1113+
{
1114+
name: "Unknown host type string returns HostTypeUnknown",
1115+
hostTypeJSON: "CUSTOM_HOST",
1116+
wantHostType: "",
1117+
},
1118+
{
1119+
name: "Empty host type returns HostTypeUnknown",
1120+
hostTypeJSON: "",
1121+
wantHostType: "",
1122+
},
1123+
}
1124+
noopLoader := mockLoader(func(cfg *Config) error { return nil })
1125+
for _, tc := range tests {
1126+
t.Run(tc.name, func(t *testing.T) {
1127+
cfg := &Config{
1128+
Host: testHMHost,
1129+
Loaders: []Loader{noopLoader},
1130+
HTTPTransport: fixtures.SliceTransport{
1131+
{
1132+
Method: "GET",
1133+
Resource: "/.well-known/databricks-config",
1134+
ReuseRequest: true,
1135+
Status: 200,
1136+
Response: `{"oidc_endpoint": "` + testHMHost + `/oidc", "account_id": "` + testHMAccountID + `", "host_type": "` + tc.hostTypeJSON + `"}`,
1137+
},
1138+
},
1139+
}
1140+
err := cfg.EnsureResolved()
1141+
require.NoError(t, err)
1142+
assert.Equal(t, tc.wantHostType, string(cfg.resolvedHostType))
1143+
})
1144+
}
1145+
}

‎config/host_metadata.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,9 @@ type HostMetadata struct {
2323
// Cloud is the cloud provider for this Databricks deployment (AWS, Azure, or GCP).
2424
Cloud environment.Cloud `json:"cloud"`
2525

26+
// HostType is the type of host (WORKSPACE_HOST, ACCOUNT_HOST, or UNIFIED_HOST).
27+
HostType HostType `json:"host_type"`
28+
2629
// DefaultOIDCAudience is the default OIDC audience for token requests.
2730
// For workspace hosts: "https://<workspace_host>/oidc/v1/token"
2831
// For account/unified hosts: the resolved account ID.

‎config/host_metadata_test.go‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ func TestGetHostMetadata_WorkspaceStaticOIDCEndpoint(t *testing.T) {
3232
"account_id": testHMAccountID,
3333
"workspace_id": testHMWorkspaceID,
3434
"cloud": "AWS",
35+
"host_type": "workspace",
3536
},
3637
},
3738
})
@@ -44,6 +45,7 @@ func TestGetHostMetadata_WorkspaceStaticOIDCEndpoint(t *testing.T) {
4445
AccountID: testHMAccountID,
4546
WorkspaceID: testHMWorkspaceID,
4647
Cloud: "AWS",
48+
HostType: WorkspaceHost,
4749
}
4850
if diff := cmp.Diff(want, meta); diff != "" {
4951
t.Errorf("mismatch (-want +got):\n%s", diff)
@@ -86,6 +88,64 @@ func TestGetHostMetadata_HTTPError(t *testing.T) {
8688
}
8789
}
8890

91+
func TestGetHostMetadata_WithHostTypeField(t *testing.T) {
92+
tests := []struct {
93+
name string
94+
hostType string
95+
wantHostType HostType
96+
}{
97+
{
98+
name: "workspace",
99+
hostType: "workspace",
100+
wantHostType: WorkspaceHost,
101+
},
102+
{
103+
name: "account",
104+
hostType: "account",
105+
wantHostType: AccountHost,
106+
},
107+
{
108+
name: "unified",
109+
hostType: "unified",
110+
wantHostType: UnifiedHost,
111+
},
112+
{
113+
name: "Workspace uppercase",
114+
hostType: "WORKSPACE",
115+
wantHostType: WorkspaceHost,
116+
},
117+
{
118+
name: "missing host_type field",
119+
hostType: "",
120+
wantHostType: HostTypeUnknown,
121+
},
122+
}
123+
for _, tc := range tests {
124+
t.Run(tc.name, func(t *testing.T) {
125+
response := map[string]string{
126+
"oidc_endpoint": testHMHost + "/oidc",
127+
"account_id": testHMAccountID,
128+
}
129+
if tc.hostType != "" {
130+
response["host_type"] = tc.hostType
131+
}
132+
client := newTestAPIClient(fixtures.MappingTransport{
133+
"GET /.well-known/databricks-config": {
134+
Status: 200,
135+
Response: response,
136+
},
137+
})
138+
meta, err := getHostMetadata(context.Background(), testHMHost, client)
139+
if err != nil {
140+
t.Fatal(err)
141+
}
142+
if meta.HostType != tc.wantHostType {
143+
t.Errorf("HostType field mismatch: got %q, want %q", meta.HostType, tc.wantHostType)
144+
}
145+
})
146+
}
147+
}
148+
89149
func TestGetHostMetadata_WithDefaultOIDCAudience(t *testing.T) {
90150
tests := []struct {
91151
name string

0 commit comments

Comments
 (0)