Merge pull request #7 from wherobots/peter/multi-session

feat: add single and multi connection session type options

Author: peterfoldes

.gitignore

@@ -32,5 +32,8 @@ build
 # IntelliJ
 .idea/
 
+# VSCode
+.vscode/
+
 # Vim
 *.swp

lib/build.gradle

@@ -15,7 +15,7 @@ plugins {
 
 description = 'JDBC driver for the Wherobots Cloud Spatial SQL API'
 group = 'com.wherobots.jdbc'
-version = '0.1.0'
+version = '0.1.1'
 
 repositories {
     mavenCentral()
@@ -90,3 +90,16 @@ signing {
 test {
     useJUnitPlatform()
 }
+
+// Run this task with `./gradlew runSmokeTest -Papikey=<your_api_key>`
+task runSmokeTest {
+    group = 'SmokeTest'
+    doLast {
+        tasks.create('runSmokeTestExec', JavaExec) {
+            main = 'com.wherobots.db.jdbc.SmokeTest'
+            args = ["$apikey"]
+            classpath = sourceSets.test.runtimeClasspath
+            jvmArgs = ["--add-opens=java.base/java.nio=org.apache.arrow.memory.core,ALL-UNNAMED"]
+        }.exec()
+    }
+}

lib/src/main/java/com/wherobots/db/SessionType.java

@@ -0,0 +1,14 @@
+package com.wherobots.db;
+
+public enum SessionType {
+    // Allow only a single concurrent connection to the SQL Session
+    SINGLE("single"),
+    // Allow multiple concurrent connections to the SQL Session
+    MULTI("multi");
+
+    public final String name;
+
+    SessionType(String name) {
+        this.name = name;
+    }
+}

lib/src/main/java/com/wherobots/db/jdbc/WherobotsJdbcDriver.java

@@ -2,6 +2,7 @@
 
 import com.wherobots.db.Region;
 import com.wherobots.db.Runtime;
+import com.wherobots.db.SessionType;
 import com.wherobots.db.jdbc.session.WherobotsSession;
 import com.wherobots.db.jdbc.session.WherobotsSessionSupplier;
 import org.apache.commons.lang3.StringUtils;
@@ -31,7 +32,7 @@ public class WherobotsJdbcDriver implements Driver {
     public static final String TOKEN_PROP = "token";
     public static final String RUNTIME_PROP = "runtime";
     public static final String REGION_PROP = "region";
-    public static final String REUSE_SESSION_PROP = "reuseSession";
+    public static final String SESSION_TYPE_PROP = "sessionType";
     public static final String WS_URI_PROP = "wsUri";
 
     // Results format; one of {@link DataFormat}
@@ -48,7 +49,7 @@ public class WherobotsJdbcDriver implements Driver {
 
     public static final Runtime DEFAULT_RUNTIME = Runtime.TINY;
     public static final Region DEFAULT_REGION = Region.AWS_US_WEST_2;
-    public static final boolean DEFAULT_REUSE_SESSION = true;
+    public static final SessionType DEFAULT_SESSION_TYPE = SessionType.SINGLE;
 
     public Map<String, String> getUserAgentHeader() {
         String javaVersion = System.getProperty("java.version");
@@ -86,10 +87,10 @@ public Connection connect(String url, Properties info) throws SQLException {
             region = Region.valueOf(regionName);
         }
 
-        boolean reuse = DEFAULT_REUSE_SESSION;
-        String reuseSession = info.getProperty(REUSE_SESSION_PROP);
-        if (StringUtils.isNotBlank(reuseSession)) {
-            reuse = Boolean.parseBoolean(reuseSession);
+        SessionType sessionType = DEFAULT_SESSION_TYPE;
+        String sessionTypeName = info.getProperty(SESSION_TYPE_PROP);
+        if (StringUtils.isNotBlank(sessionTypeName)) {
+            sessionType = SessionType.valueOf(sessionTypeName);
         }
 
         Map<String, String> headers = new HashMap<>(getAuthHeaders(info));
@@ -105,7 +106,7 @@ public Connection connect(String url, Properties info) throws SQLException {
                 throw new SQLException("Invalid WebSocket URI: " + wsUriString, e);
             }
         } else {
-            session = WherobotsSessionSupplier.create(host, runtime, region, reuse, headers);
+            session = WherobotsSessionSupplier.create(host, runtime, region, sessionType, headers);
         }
 
         return new WherobotsJdbcConnection(session, info);

lib/src/main/java/com/wherobots/db/jdbc/session/WherobotsSessionSupplier.java

@@ -5,6 +5,7 @@
 import com.wherobots.db.AppStatus;
 import com.wherobots.db.Region;
 import com.wherobots.db.Runtime;
+import com.wherobots.db.SessionType;
 import com.wherobots.db.jdbc.serde.JsonUtil;
 import io.github.resilience4j.core.IntervalFunction;
 import io.github.resilience4j.core.functions.CheckedSupplier;
@@ -34,11 +35,14 @@ public abstract class WherobotsSessionSupplier {
 
     private static final Logger logger = LoggerFactory.getLogger(WherobotsSessionSupplier.class);
 
-    private static final String SQL_SESSION_ENDPOINT = "https://%s/sql/session?region=%s&reuse_session=%s";
+    private static final String SQL_SESSION_ENDPOINT = "https://%s/sql/session?region=%s";
     private static final String PROTOCOL_VERSION = "1.0.0";
 
     @JsonInclude(JsonInclude.Include.NON_NULL)
-    private record SqlSessionRequestPayload(String runtimeId, Integer shutdownAfterInactiveSeconds) {}
+    private record SqlSessionRequestPayload(
+        String runtimeId,
+        Integer shutdownAfterInactiveSeconds,
+        String sessionType) {}
     private record SqlSessionAppMeta(String url) {}
     private record SqlSessionResponsePayload(AppStatus status, SqlSessionAppMeta appMeta) {}
 
@@ -52,7 +56,7 @@ private record SqlSessionResponsePayload(AppStatus status, SqlSessionAppMeta app
      * @return
      * @throws SQLException
      */
-    public static WherobotsSession create(String host, Runtime runtime, Region region, boolean reuse, Map<String, String> headers)
+    public static WherobotsSession create(String host, Runtime runtime, Region region, SessionType sessionType, Map<String, String> headers)
         throws SQLException {
         HttpClient client = HttpClient.newBuilder()
                 .followRedirects(HttpClient.Redirect.NORMAL)
@@ -67,7 +71,7 @@ public static WherobotsSession create(String host, Runtime runtime, Region regio
         Retry retry = RetryRegistry.of(config).retry("session");
 
         try {
-            URI sessionIdUri = new SqlSessionSupplier(client, headers, host, runtime, region, reuse).get();
+            URI sessionIdUri = new SqlSessionSupplier(client, headers, host, runtime, region, sessionType).get();
             URI wsUri = Retry.decorateCheckedSupplier(retry, new SessionWsUriSupplier(client, headers, sessionIdUri)).get();
             return create(wsUri, headers);
         } catch (SQLException e) {
@@ -103,20 +107,23 @@ private record SqlSessionSupplier(HttpClient client,
                                       String host,
                                       Runtime runtime,
                                       Region region,
-                                      boolean reuse)
+                                      SessionType sessionType)
             implements CheckedSupplier<URI> {
 
         @Override
         public URI get() throws IOException, InterruptedException {
-            logger.info("{} {} runtime in {} from {}...",
-                    reuse ? "Recycling" : "Requesting", runtime.name, region.name, host);
+            logger.info("Requesting {} runtime using {} session type in {} from {}...",
+                    runtime.name, sessionType.name, region.name, host);
 
             HttpRequest.BodyPublisher body = HttpRequest.BodyPublishers.ofString(
-                    JsonUtil.serialize(new SqlSessionRequestPayload(runtime.name, null)));
+                    JsonUtil.serialize(new SqlSessionRequestPayload(
+                        runtime.name,
+                        null,
+                        sessionType.name)));
 
             HttpRequest.Builder request = HttpRequest.newBuilder()
                     .POST(body)
-                    .uri(URI.create(String.format(SQL_SESSION_ENDPOINT, host, region.name, reuse)))
+                    .uri(URI.create(String.format(SQL_SESSION_ENDPOINT, host, region.name)))
                     .header("Content-Type", "application/json");
             headers.forEach(request::header);
 

lib/src/test/java/com/wherobots/db/jdbc/SmokeTest.java

@@ -7,7 +7,14 @@
 import java.sql.Statement;
 import java.util.Properties;
 
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import com.wherobots.db.SessionType;
+
 public class SmokeTest {
+    private static final Logger logger = LoggerFactory.getLogger(SmokeTest.class);
+
     public static void main(String[] args) throws Exception {
         DriverManager.setLogWriter(new PrintWriter(System.out));
         DriverManager.registerDriver(new WherobotsJdbcDriver());
@@ -27,7 +34,9 @@ public static void main(String[] args) throws Exception {
 
         Properties props = new Properties();
         props.put(WherobotsJdbcDriver.API_KEY_PROP, args[0]);
-        props.put(WherobotsJdbcDriver.REUSE_SESSION_PROP, "true");
+        props.put(WherobotsJdbcDriver.SESSION_TYPE_PROP, SessionType.SINGLE.name());
+
+        logger.info("Connecting to Wherobots SQL API with properties: {}", props);
 
         try (Connection conn = DriverManager.getConnection("jdbc:wherobots://api.staging.wherobots.com", props)) {
             try (Statement stmt = conn.createStatement()) {