Skip to content

Commit 4bd5abc

Browse files
committed
exclude entitlements from CurrentWorkspaceID Me probe
CurrentWorkspaceID issues `GET /api/2.0/preview/scim/v2/Me` to read the X-Databricks-Org-Id response header. The Me endpoint computes User.Entitlements by default - a scan whose cost scales with the total number of entitlement grants in the workspace - and CurrentWorkspaceID never reads the body at all. Ask the server to skip the attribute.
1 parent 04b7dda commit 4bd5abc

2 files changed

Lines changed: 25 additions & 0 deletions

File tree

workspace_functions.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ func (w *WorkspaceClient) CurrentWorkspaceID(ctx context.Context) (int64, error)
1919
var workspaceIdStr string
2020
opts := []httpclient.DoOption{
2121
httpclient.WithResponseHeader("X-Databricks-Org-Id", &workspaceIdStr),
22+
httpclient.WithRequestData(map[string]string{"excludedAttributes": "entitlements"}),
2223
}
2324
if w.Config != nil && w.Config.WorkspaceID != "" {
2425
opts = append(opts, httpclient.WithRequestHeader("X-Databricks-Org-Id", w.Config.WorkspaceID))

workspace_functions_test.go

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,30 @@ func TestCurrentWorkspaceIDSendsOrgIdHeaderWhenConfigHasWorkspaceID(t *testing.T
4242
assert.Equal(t, "7474644166319138", gotOrgIdHeader)
4343
}
4444

45+
func TestCurrentWorkspaceIDExcludesEntitlements(t *testing.T) {
46+
var gotRawQuery string
47+
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
48+
if r.URL.Path == "/api/2.0/preview/scim/v2/Me" {
49+
gotRawQuery = r.URL.RawQuery
50+
w.Header().Set("X-Databricks-Org-Id", "7474644166319138")
51+
w.Write([]byte(`{}`))
52+
return
53+
}
54+
http.NotFound(w, r)
55+
}))
56+
defer server.Close()
57+
58+
w, err := NewWorkspaceClient(&Config{
59+
Host: server.URL,
60+
Token: "token",
61+
})
62+
require.NoError(t, err)
63+
64+
_, err = w.CurrentWorkspaceID(t.Context())
65+
require.NoError(t, err)
66+
assert.Equal(t, "excludedAttributes=entitlements", gotRawQuery)
67+
}
68+
4569
func TestCurrentWorkspaceIDOmitsOrgIdHeaderWhenConfigMissingWorkspaceID(t *testing.T) {
4670
// On legacy workspace hosts the host itself identifies the workspace, so
4771
// no routing header is needed. When Config.WorkspaceID is empty we send

0 commit comments

Comments
 (0)