Support default workspace for queries

rockset · Apr 3, 2024 · 377d8a6 · 377d8a6
1 parent 00f48d9
commit 377d8a6
Show file tree

Hide file tree

Showing 3 changed files with 492 additions and 392 deletions.
diff --git a/src/main/java/com/rockset/jdbc/RocksetStatement.java b/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) {

diff --git a/src/test/java/com/rockset/client/TestWorkspace.java b/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;
@@ -60,7 +62,8 @@ public void testDeleteWorkspace() throws Exception {
 
     // wait for collection to go away
     Awaitility.await("Waiting for collection to be cleaned up ")
-        .atMost(60, TimeUnit.SECONDS)
+        .atMost(120, TimeUnit.SECONDS)
+        .pollInterval(1, TimeUnit.SECONDS)
         .until(
             (Callable<Boolean>)
                 () -> {
@@ -70,7 +73,6 @@ public void testDeleteWorkspace() throws Exception {
                   } catch (Exception e) {
                     return true; // collection deleted
                   }
-                  Thread.sleep(1000);
                   return false;
                 });