Adding custom User-Agent to Java SDK (#29)

spiceai · Nov 21, 2024 · 46a25f9 · 46a25f9
1 parent b734850
commit 46a25f9
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 5 deletions.
diff --git a/src/main/java/ai/spice/Config.java b/src/main/java/ai/spice/Config.java
@@ -123,7 +123,7 @@ public static String getUserAgent() {
             osVersion = osVersion.split(" ")[0];
         }
 
-        return "spice-java " + Version.SPICE_JAVA_VERSION + " (" + osName + "/"
+        return "spice-java/" + Version.SPICE_JAVA_VERSION + " (" + osName + "/"
                 + System.getProperty("os.version") + " "
                 + osArch + ")";
     }

diff --git a/src/main/java/ai/spice/SpiceClient.java b/src/main/java/ai/spice/SpiceClient.java
@@ -92,7 +92,7 @@ public static SpiceClientBuilder builder() throws URISyntaxException {
      * 
      * @param maxRetries    the maximum number of connection retries for the client
      */
-    public SpiceClient(String appId, String apiKey, URI flightAddress, URI httpAddress, int maxRetries) {
+    public SpiceClient(String appId, String apiKey, URI flightAddress, URI httpAddress, int maxRetries, String userAgent) {
         this.appId = appId;
         this.apiKey = apiKey;
         this.maxRetries = maxRetries;
@@ -117,7 +117,13 @@ public SpiceClient(String appId, String apiKey, URI flightAddress, URI httpAddre
 
         // prepare additional headers to insert into Flight requests
         Map<String, String> headers = new HashMap<>();
-        headers.put("X-Spice-User-Agent", Config.getUserAgent());
+        String uaString;
+        if (Strings.isNullOrEmpty(userAgent)) {
+            uaString = Config.getUserAgent();
+        } else {
+            uaString = userAgent;
+        }
+        headers.put("User-Agent", uaString);
 
         final ClientIncomingAuthHeaderMiddleware.Factory authFactory = new ClientIncomingAuthHeaderMiddleware.Factory(
                 new ClientBearerHeaderHandler());

diff --git a/src/main/java/ai/spice/SpiceClientBuilder.java b/src/main/java/ai/spice/SpiceClientBuilder.java
@@ -34,6 +34,7 @@ public class SpiceClientBuilder {
 
     private String appId;
     private String apiKey;
+    private String userAgent;
     private URI flightAddress;
     private URI httpAddress;
     private int maxRetries = 3;
@@ -98,6 +99,20 @@ public SpiceClientBuilder withApiKey(String apiKey) {
         return this;
     }
 
+    /**
+     * Sets the client's custom User-Agent string
+     * 
+     * @param userAgent The User-Agent string
+     * @return The current instance of SpiceClientBuilder for method chaining.
+     */
+    public SpiceClientBuilder withUserAgent(String userAgent) {
+        if (Strings.isNullOrEmpty(userAgent)) {
+            throw new IllegalArgumentException("userAgent can't be null or empty");
+        }
+        this.userAgent = userAgent;
+        return this;
+    }
+
     /**
      * Sets the client's flight address to default Spice Cloud address.
      *
@@ -130,6 +145,6 @@ public SpiceClientBuilder withMaxRetries(int maxRetries) {
      * @return The SpiceClient instance
      */
     public SpiceClient build() {
-        return new SpiceClient(appId, apiKey, flightAddress, httpAddress, maxRetries);
+        return new SpiceClient(appId, apiKey, flightAddress, httpAddress, maxRetries, userAgent);
     }
 }
diff --git a/src/test/java/ai/spice/UserAgentTest.java b/src/test/java/ai/spice/UserAgentTest.java
@@ -33,7 +33,7 @@ public class UserAgentTest
         extends TestCase {
     public void testUserAgent() throws ExecutionException, InterruptedException {
         // use a regex to match the expected user agent
-        String regex = "spice-java \\d+\\.\\d+\\.\\d+ \\((Linux|Windows|Darwin)/[\\d\\w\\.\\-\\_]+ (x86_64|aarch64|i386)\\)";
+        String regex = "spice-java/\\d+\\.\\d+\\.\\d+ \\((Linux|Windows|Darwin)/[\\d\\w\\.\\-\\_]+ (x86_64|aarch64|i386)\\)";
         Pattern pattern = Pattern.compile(regex);
 
         String userAgent = Config.getUserAgent();