Merge pull request #4 from wherobots/use-new-runtime-ids

chore: use new runtime ids, add user agent header
wherobots · Sep 26, 2024 · 38a3721 · 38a3721
2 parents 53daca1 + af7b5b5
commit 38a3721
Show file tree

Hide file tree

Showing 4 changed files with 49 additions and 9 deletions.
diff --git a/lib/build.gradle b/lib/build.gradle
@@ -30,6 +30,9 @@ dependencies {
     implementation 'org.apache.arrow:arrow-memory-netty:16.1.0'
     implementation 'io.github.resilience4j:resilience4j-retry:2.2.0'
     implementation 'org.slf4j:slf4j-simple:2.0.13'
+    // Test dependencies
+    testImplementation 'org.junit.jupiter:junit-jupiter-api:5.8.1'
+    testRuntimeOnly 'org.junit.jupiter:junit-jupiter-engine:5.8.1'
 }
 
 // Apply a specific Java toolchain to ease working on different environments.
@@ -51,3 +54,7 @@ tasks.named('jar') {
 tasks.withType(Jar) {
     archiveBaseName.set('wherobots-jdbc')
 }
+
+test {
+    useJUnitPlatform()
+}
diff --git a/lib/src/main/java/com/wherobots/db/Runtime.java b/lib/src/main/java/com/wherobots/db/Runtime.java
@@ -1,16 +1,16 @@
 package com.wherobots.db;
 
 public enum Runtime {
-    SEDONA("TINY"),
-    SAN_FRANCISCO("SMALL"),
-    NEW_YORK("MEDIUM"),
-    CAIRO("LARGE"),
-    DELHI("XLARGE"),
-    TOKYO("XXLARGE"),
+    SEDONA("tiny"),
+    SAN_FRANCISCO("small"),
+    NEW_YORK("medium"),
+    CAIRO("large"),
+    DELHI("x-large"),
+    TOKYO("2x-large"),
 
     NEW_YORK_HIMEM("medium-himem"),
     CAIRO_HIMEM("large-himem"),
-    DEHLI_HIMEM("xlarge-himem"),
+    DEHLI_HIMEM("x-large-himem"),
     TOKYO_HIMEM("2x-large-himem"),
     ATLANTIS_HIMEM("4x-large-himem"),
 

diff --git a/lib/src/main/java/com/wherobots/db/jdbc/WherobotsJdbcDriver.java b/lib/src/main/java/com/wherobots/db/jdbc/WherobotsJdbcDriver.java
@@ -14,6 +14,7 @@
 import java.sql.DriverPropertyInfo;
 import java.sql.SQLException;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.Map;
 import java.util.Properties;
 
@@ -47,6 +48,18 @@ public class WherobotsJdbcDriver implements Driver {
     public static final Runtime DEFAULT_RUNTIME = Runtime.SEDONA;
     public static final Region DEFAULT_REGION = Region.AWS_US_WEST_2;
 
+    public Map<String, String> getUserAgentHeader() {
+        String javaVersion = System.getProperty("java.version");
+        String osName = System.getProperty("os.name");
+        String packageVersion = getClass().getPackage().getImplementationVersion();
+        if (packageVersion == null) {
+            packageVersion = "unknown";
+        }
+        String userAgent = String.format("wherobots-jdbc-driver/%s os/%s java/%s",
+                packageVersion, osName, javaVersion);
+        return Map.of("User-Agent", userAgent);
+    }
+
     @Override
     public Connection connect(String url, Properties info) throws SQLException {
         String host = DEFAULT_ENDPOINT;
@@ -70,8 +83,8 @@ public Connection connect(String url, Properties info) throws SQLException {
         if (StringUtils.isNotBlank(regionName)) {
             region = Region.valueOf(regionName);
         }
-
-        Map<String, String> headers = getAuthHeaders(info);
+        Map<String, String> headers = new HashMap<>(getAuthHeaders(info));
+        headers.putAll(getUserAgentHeader());
         WherobotsSession session;
 
         String wsUriString = info.getProperty(WS_URI_PROP);

diff --git a/lib/src/test/java/com/wherobots/db/jdbc/WherobotsJdbcDriverTest.java b/lib/src/test/java/com/wherobots/db/jdbc/WherobotsJdbcDriverTest.java
@@ -0,0 +1,20 @@
+package com.wherobots.db.jdbc;
+
+
+import org.junit.jupiter.api.Test;
+
+import java.util.Map;
+
+class WherobotsJdbcDriverTest {
+
+    @Test
+    void getUserAgentHeader() {
+        System.setProperty("java.version", "1");
+        System.setProperty("os.name", "os1");
+        WherobotsJdbcDriver driver = new WherobotsJdbcDriver();
+        Map<String, String> header = driver.getUserAgentHeader();
+        assert header.containsKey("User-Agent");
+        String user_agent = header.get("User-Agent");
+        assert user_agent.equals("wherobots-jdbc-driver/unknown os/os1 java/1");
+    }
+}