Skip to content

Commit 134c417

Browse files
Store WorkspaceID in the API Client (#647)
## What changes are proposed in this pull request? Store WorkspaceID in the API Client. This will be used to inject headers in Unified Mode. ## How is this tested? Added unit tests NO_CHANGELOG=true
1 parent b06a44c commit 134c417

File tree

2 files changed

+34
-0
lines changed

2 files changed

+34
-0
lines changed

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ public static class Builder {
4242
private Integer debugTruncateBytes;
4343
private HttpClient httpClient;
4444
private String accountId;
45+
private String workspaceId;
4546
private RetryStrategyPicker retryStrategyPicker;
4647
private boolean isDebugHeaders;
4748

@@ -52,6 +53,7 @@ public Builder withDatabricksConfig(DatabricksConfig config) {
5253
this.httpClient = config.getHttpClient();
5354
this.debugTruncateBytes = config.getDebugTruncateBytes();
5455
this.accountId = config.getAccountId();
56+
this.workspaceId = config.getWorkspaceId();
5557
this.isDebugHeaders = config.isDebugHeaders();
5658

5759
if (config.getDisableRetries()) {
@@ -114,6 +116,7 @@ public ApiClient build() {
114116
private final Function<Void, String> getHostFunc;
115117
private final Function<Void, String> getAuthTypeFunc;
116118
private final String accountId;
119+
private final String workspaceId;
117120
private final boolean isDebugHeaders;
118121
private static final String RETRY_AFTER_HEADER = "retry-after";
119122

@@ -125,6 +128,10 @@ public String configuredAccountID() {
125128
return accountId;
126129
}
127130

131+
public String workspaceId() {
132+
return workspaceId;
133+
}
134+
128135
public ApiClient(DatabricksConfig config) {
129136
this(config, new SystemTimer());
130137
}
@@ -143,6 +150,7 @@ private ApiClient(Builder builder) {
143150
this.getAuthTypeFunc = builder.getAuthTypeFunc != null ? builder.getAuthTypeFunc : v -> "";
144151
this.httpClient = builder.httpClient;
145152
this.accountId = builder.accountId;
153+
this.workspaceId = builder.workspaceId;
146154
this.retryStrategyPicker =
147155
builder.retryStrategyPicker != null
148156
? builder.retryStrategyPicker

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

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -483,4 +483,30 @@ void privateLinkRedirectBecomesPrivateLinkValidationError() throws MalformedURLE
483483
new Request("GET", req.getUri().getPath()), MyEndpointResponse.class));
484484
assertTrue(e.getMessage().contains("AWS PrivateLink"));
485485
}
486+
487+
@Test
488+
void testDefaultWorkspaceIdReturnsNullWhenNotSet() {
489+
Request req = getBasicRequest();
490+
DatabricksConfig config =
491+
new DatabricksConfig()
492+
.setHost("http://my.host")
493+
.setCredentialsProvider(new DummyCredentialsProvider());
494+
ApiClient client =
495+
getApiClient(config, req, Collections.singletonList(getSuccessResponse(req)));
496+
assertNull(client.workspaceId());
497+
}
498+
499+
@Test
500+
void testDefaultWorkspaceIdReturnsValueFromConfig() {
501+
String expectedWorkspaceId = "test-workspace-123";
502+
Request req = getBasicRequest();
503+
DatabricksConfig config =
504+
new DatabricksConfig()
505+
.setHost("http://my.host")
506+
.setWorkspaceId(expectedWorkspaceId)
507+
.setCredentialsProvider(new DummyCredentialsProvider());
508+
ApiClient client =
509+
getApiClient(config, req, Collections.singletonList(getSuccessResponse(req)));
510+
assertEquals(expectedWorkspaceId, client.workspaceId());
511+
}
486512
}

0 commit comments

Comments
 (0)