Support default workspace for queries (#78)

src/main/java/com/rockset/jdbc/RocksetStatement.java

@@ -185,19 +185,24 @@ private static String getQueryIdFromQueryResponse(QueryResponse response) {
   protected boolean executeWithParams(String sql, List<QueryParameter> params) throws SQLException {
     clearCurrentResults();
     checkOpen();
+    final String schema = this.connection.get().getSchema();
+
+    final String sqlWithWorkspace = schema.equals(RocksetConnection.DEFAULT_SCHEMA)
+            ? sql
+            : String.format("OPTION(default_workspace='%s')\n %s", schema, sql);
 
     ResultSet resultSet = null;
     try {
       // Make query to rockset service. We do not use queryTimeoutSeconds
       // because rockset queries do not yet have a client-side timeout.
       QueryResponse resp =
           connection()
-              .startQuery(sql, this.fetchSize.get(), params, getStatementSessionProperties());
+              .startQuery(sqlWithWorkspace, this.fetchSize.get(), params, getStatementSessionProperties());
 
       // store resuts in memory
       resultSet =
           new RocksetResultSet(
-              sql,
+              sqlWithWorkspace,
               resp,
               this.maxRows.get(),
               RocksetResultSetPaginationParams.builder()
@@ -210,7 +215,7 @@ protected boolean executeWithParams(String sql, List<QueryParameter> params) thr
       this.currentResult.set(resultSet);
       return true;
     } catch (RuntimeException e) {
-      String msg = "Error executing query '" + sql + "'" + " error =  " + e.getMessage();
+      String msg = "Error executing query '" + sqlWithWorkspace + "'" + " error =  " + e.getMessage();
       RocksetDriver.log(msg);
       throw new SQLException(msg, e);
     } catch (Exception e) {

src/test/java/com/rockset/client/TestWorkspace.java

@@ -1,6 +1,8 @@
 package com.rockset.client;
 
 import com.rockset.client.model.*;
+
+import java.sql.Time;
 import java.util.concurrent.Callable;
 import java.util.concurrent.TimeUnit;
 import org.apache.commons.lang3.RandomStringUtils;
@@ -61,7 +63,7 @@ public void testDeleteWorkspace() throws Exception {
     // wait for collection to go away
     Awaitility.await("Waiting for collection to be cleaned up ")
         .atMost(3, TimeUnit.MINUTES)
-            .pollInterval(1, TimeUnit.SECONDS)
+        .pollInterval(1, TimeUnit.SECONDS)
         .until(
             (Callable<Boolean>)
                 () -> {

src/test/java/com/rockset/client/util/RetryHelper.java

@@ -0,0 +1,25 @@
+package com.rockset.client.util;
+
+import com.rockset.client.ApiException;
+import org.awaitility.Awaitility;
+import org.hamcrest.Matchers;
+
+import java.util.concurrent.Callable;
+import java.util.concurrent.TimeUnit;
+
+public class RetryHelper {
+    public static <T> T retryOnApiException(Callable<T> request) {
+        return Awaitility
+                .await()
+                .atMost(5, TimeUnit.MINUTES)
+                .pollInterval(1, TimeUnit.SECONDS)
+                .until(() -> {
+                    try {
+                        return request.call();
+                    } catch (ApiException e) {
+                        System.out.println("Encountered error, retrying request: " + e.getMessage());
+                        return null;
+                    }
+                }, Matchers.notNullValue());
+    }
+}