From f04f6cbb233ac294798414f777921b49f402e907 Mon Sep 17 00:00:00 2001
From: Satvik Ramaprasad <satvik.ramaprasad@databricks.com>
Date: Wed, 6 Nov 2024 12:58:48 +0530
Subject: [PATCH] [Feature] DatabricksConfig: Add clone() support (#376)

## 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
---
 .../databricks/sdk/core/DatabricksConfig.java | 33 +++++++++++--------
 .../sdk/core/DatabricksConfigTest.java        | 18 ++++++++++
 2 files changed, 38 insertions(+), 13 deletions(-)

diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java
index 3b692860..20e7f883 100644
--- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java
+++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java
@@ -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<>(
@@ -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);
   }
 }
diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java
index 3ec891c6..e552a142 100644
--- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java
+++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java
@@ -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");
+  }
 }