Support partners in SDK

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

@@ -1,14 +1,33 @@
 package com.databricks.sdk.core;
 
-import java.util.HashMap;
-import java.util.Map;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 
 public class UserAgent {
   private static String product = "unknown";
   private static String productVersion = "0.0.0";
 
-  private static final Map<String, String> otherInfo = new HashMap<>();
+  private static class Info {
+    private String key;
+    private String value;
+
+    public Info(String key, String value) {
+      this.key = key;
+      this.value = value;
+    }
+
+    public String getKey() {
+      return key;
+    }
+
+    public String getValue() {
+      return value;
+    }
+  }
+
+  private static final ArrayList<Info> otherInfo = new ArrayList<>();
 
   // TODO: check if reading from
   // /META-INF/maven/com.databricks/databrics-sdk-java/pom.properties
@@ -20,8 +39,62 @@ public static void withProduct(String product, String productVersion) {
     UserAgent.productVersion = productVersion;
   }
 
+  // Regular expression copied from https://semver.org/.
+  private static final String semVerCore = "(0|[1-9]\\d*)\\.(0|[1-9]\\d*)\\.(0|[1-9]\\d*)";
+  private static final String semVerPrerelease =
+      "(?:-((?:0|[1-9]\\d*|\\d*[a-zA-Z-][0-9a-zA-Z-]*)(?:\\.(?:0|[1-9]\\d*|\\d*[a-zA-Z-][0-9a-zA-Z-]*))*))?";
+  private static final String semVerBuildmetadata = "(?:\\+([0-9a-zA-Z-]+(?:\\.[0-9a-zA-Z-]+)*))?";
+  private static final Pattern regexpSemVer =
+      Pattern.compile("^" + semVerCore + semVerPrerelease + semVerBuildmetadata + "$");
+
+  // Sanitize replaces all non-alphanumeric characters with a hyphen. Use this to
+  // ensure that the user agent value is valid. This is useful when the value is not
+  // ensured to be valid at compile time.
+  //
+  // Example: You want to avoid having '/' and ' ' in the value because it will
+  // make downstream applications fail.
+  //
+  // Note: Semver strings are comprised of alphanumeric characters, hyphens, periods
+  // and plus signs. This function will not remove these characters.
+  // see:
+  // 1. https://semver.org/#spec-item-9
+  // 2. https://semver.org/#spec-item-10
+  private static final Pattern regexpAlphanum = Pattern.compile("^[0-9A-Za-z_\\.\\+-]+$");
+  private static final Pattern regexpAlphanumInverse = Pattern.compile("[^0-9A-Za-z_\\.\\+-]");
+
+  public static String sanitize(String s) {
+    return regexpAlphanumInverse.matcher(s).replaceAll("-");
+  }
+
+  public static boolean matchSemVer(String s) throws IllegalArgumentException {
+    if (regexpSemVer.matcher(s).matches()) {
+      return true;
+    }
+    throw new IllegalArgumentException("Invalid semver string: " + s);
+  }
+
+  public static boolean matchAlphanum(String s) throws IllegalArgumentException {
+    if (regexpAlphanum.matcher(s).matches()) {
+      return true;
+    }
+    throw new IllegalArgumentException("Invalid alphanumeric string: " + s);
+  }
+
+  public static boolean matchAlphanumOrSemVer(String s) throws IllegalArgumentException {
+    if (regexpAlphanum.matcher(s).matches() || regexpSemVer.matcher(s).matches()) {
+      return true;
+    }
+    throw new IllegalArgumentException("Invalid alphanumeric or semver string: " + s);
+  }
+
+  public static void withPartner(String partner) {
+    withOtherInfo("partner", partner);
+  }
+
   public static void withOtherInfo(String key, String value) {
-    otherInfo.put(key, value);
+    matchAlphanum(key);
+    matchAlphanumOrSemVer(value);
+    otherInfo.add(new Info(key, value));
   }
 
   private static String osName() {
@@ -43,12 +116,15 @@ private static String jvmVersion() {
   }
 
   public static String asString() {
-    String otherInfo =
-        UserAgent.otherInfo.entrySet().stream()
-            .map(e -> String.format(" %s/%s", e.getKey(), e.getValue()))
-            .collect(Collectors.joining());
-    return String.format(
-        "%s/%s databricks-sdk-java/%s jvm/%s os/%s%s",
-        product, productVersion, version, jvmVersion(), osName(), otherInfo);
+    List<String> segments = new ArrayList<>();
+    segments.add(String.format("%s/%s", product, productVersion));
+    segments.add(String.format("databricks-sdk-java/%s", version));
+    segments.add(String.format("jvm/%s", jvmVersion()));
+    segments.add(String.format("os/%s", osName()));
+    segments.addAll(
+        otherInfo.stream()
+            .map(e -> String.format("%s/%s", e.getKey(), e.getValue()))
+            .collect(Collectors.toSet()));
+    return segments.stream().collect(Collectors.joining(" "));
   }
 }

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

@@ -0,0 +1,59 @@
+package com.databricks.sdk.core;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+public class UserAgentTest {
+  @Test
+  public void testUserAgent() {
+    UserAgent.withProduct("product", "productVersion");
+    String userAgent = UserAgent.asString();
+    Assertions.assertTrue(userAgent.contains("product/productVersion"));
+    Assertions.assertTrue(userAgent.contains("databricks-sdk-java/"));
+    Assertions.assertTrue(userAgent.contains("os/"));
+    Assertions.assertTrue(userAgent.contains("jvm/"));
+  }
+
+  @Test
+  public void testUserAgentWithPartner() {
+    UserAgent.withPartner("partner1");
+    UserAgent.withPartner("partner2");
+    String userAgent = UserAgent.asString();
+    Assertions.assertTrue(userAgent.contains("partner/partner1"));
+    Assertions.assertTrue(userAgent.contains("partner/partner2"));
+  }
+
+  @Test
+  public void testUserAgentWithOtherInfo() {
+    UserAgent.withOtherInfo("key1", "value1");
+    UserAgent.withOtherInfo("key2", "value2");
+    String userAgent = UserAgent.asString();
+    Assertions.assertTrue(userAgent.contains("key1/value1"));
+    Assertions.assertTrue(userAgent.contains("key2/value2"));
+  }
+
+  @Test
+  public void testUserAgentWithInvalidKey() {
+    Assertions.assertThrows(
+        IllegalArgumentException.class,
+        () -> {
+          UserAgent.withOtherInfo("key1!", "value1");
+        });
+  }
+
+  @Test
+  public void testUserAgentWithInvalidValue() {
+    Assertions.assertThrows(
+        IllegalArgumentException.class,
+        () -> {
+          UserAgent.withOtherInfo("key1", "value1!");
+        });
+  }
+
+  @Test
+  public void testUserAgentWithSemverValue() {
+    UserAgent.withOtherInfo("key1", "1.0.0-dev+metadata");
+    String userAgent = UserAgent.asString();
+    Assertions.assertTrue(userAgent.contains("key1/1.0.0-dev+metadata"));
+  }
+}