From 4e2a6c0505c23a2b156603e11afb6628b17325ed Mon Sep 17 00:00:00 2001
From: "zongsi.zhang" <krishnazongsi@gmail.com>
Date: Tue, 17 Sep 2024 17:35:26 +0800
Subject: [PATCH 1/2] chore: use new runtime ids, add user agent header

---
 lib/src/main/java/com/wherobots/db/Runtime.java | 14 +++++++-------
 .../wherobots/db/jdbc/WherobotsJdbcDriver.java  | 17 +++++++++++++++--
 2 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/lib/src/main/java/com/wherobots/db/Runtime.java b/lib/src/main/java/com/wherobots/db/Runtime.java
index fe68c87..0a98366 100644
--- 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
index c1ca124..12a1e10 100644
--- 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);

From af7b5b51cb9ac1e672b88b2460acb981420be885 Mon Sep 17 00:00:00 2001
From: "zongsi.zhang" <krishnazongsi@gmail.com>
Date: Fri, 20 Sep 2024 17:29:29 +0800
Subject: [PATCH 2/2] feat: Add test case for userAgentHeader

---
 lib/build.gradle                              |  7 +++++++
 .../db/jdbc/WherobotsJdbcDriverTest.java      | 20 +++++++++++++++++++
 2 files changed, 27 insertions(+)
 create mode 100644 lib/src/test/java/com/wherobots/db/jdbc/WherobotsJdbcDriverTest.java

diff --git a/lib/build.gradle b/lib/build.gradle
index bafe764..133e17f 100644
--- 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/test/java/com/wherobots/db/jdbc/WherobotsJdbcDriverTest.java b/lib/src/test/java/com/wherobots/db/jdbc/WherobotsJdbcDriverTest.java
new file mode 100644
index 0000000..8b92c1d
--- /dev/null
+++ 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");
+    }
+}
\ No newline at end of file