Skip to content

Commit

Permalink
[Feature] DatabricksConfig: Add clone() support (#376)
Browse files Browse the repository at this point in the history
## Changes
Adds support for cloning DatabricksConfig(), this is needed because we
need a way to set configurations per API call such as
1. timeout
2. httpClient configuration
3. debugHeaders

However, we still want to use the cached oauth token etc and we would
like the header factory to be a common object across workspace clients
and databricks configs

Ideally this should be supported by the SDK itself natively, however
since its not supported and it will take significant migration to
achieve that, this is a work around to achieve the same.


## Tests
Added UT
  • Loading branch information
satviksr-db authored Nov 6, 2024
1 parent a1460c5 commit f04f6cb
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 13 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -631,6 +631,25 @@ public DatabricksEnvironment getDatabricksEnvironment() {
return DatabricksEnvironment.getEnvironmentFromHostname(this.host);
}

private DatabricksConfig clone(Set<String> fieldsToSkip) {
DatabricksConfig newConfig = new DatabricksConfig();
for (Field f : DatabricksConfig.class.getDeclaredFields()) {
if (fieldsToSkip.contains(f.getName())) {
continue;
}
try {
f.set(newConfig, f.get(this));
} catch (IllegalAccessException e) {
throw new RuntimeException(e);
}
}
return newConfig;
}

public DatabricksConfig clone() {
return clone(new HashSet<>());
}

public DatabricksConfig newWithWorkspaceHost(String host) {
Set<String> fieldsToSkip =
new HashSet<>(
Expand All @@ -645,18 +664,6 @@ public DatabricksConfig newWithWorkspaceHost(String host) {
// don't cache the
// header factory.
"headerFactory"));
DatabricksConfig newConfig = new DatabricksConfig();
for (Field f : DatabricksConfig.class.getDeclaredFields()) {
if (fieldsToSkip.contains(f.getName())) {
continue;
}
try {
f.set(newConfig, f.get(this));
} catch (IllegalAccessException e) {
throw new RuntimeException(e);
}
}
newConfig.setHost(host);
return newConfig;
return clone(fieldsToSkip).setHost(host);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -177,4 +177,22 @@ public void testNewWithWorkspaceHost() {
assert newWorkspaceConfig.getClientId().equals("my-client-id");
assert newWorkspaceConfig.getClientSecret().equals("my-client-secret");
}

@Test
public void testClone() {
DatabricksConfig config =
new DatabricksConfig()
.setAuthType("oauth-m2m")
.setClientId("my-client-id")
.setClientSecret("my-client-secret")
.setAccountId("account-id")
.setHost("https://account.cloud.databricks.com");

DatabricksConfig newWorkspaceConfig = config.clone();

assert newWorkspaceConfig.getHost().equals("https://account.cloud.databricks.com");
assert newWorkspaceConfig.getAuthType().equals("oauth-m2m");
assert newWorkspaceConfig.getClientId().equals("my-client-id");
assert newWorkspaceConfig.getClientSecret().equals("my-client-secret");
}
}

0 comments on commit f04f6cb

Please sign in to comment.