Fix result set NPE for procedure creation during async query (#96)

looker-open-source · Jul 12, 2021 · 0f81d63 · 0f81d63
1 parent e6c5953
commit 0f81d63
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 3 deletions.
diff --git a/src/main/java/net/starschema/clouddb/jdbc/BQForwardOnlyResultSet.java b/src/main/java/net/starschema/clouddb/jdbc/BQForwardOnlyResultSet.java
@@ -158,7 +158,7 @@ public BQForwardOnlyResultSet(Bigquery bigquery, String projectId,
                     fetchPos = fetchPos.add(BigInteger.valueOf(this.rowsofResult.size()));
                 }
                 this.totalBytesProcessed = result.getTotalBytesProcessed();
-                this.cacheHit = result.getCacheHit();
+                this.cacheHit = Boolean.TRUE.equals(result.getCacheHit()); // coerce Boolean nullable object to boolean primitive
             }
         }
     }

diff --git a/src/main/java/net/starschema/clouddb/jdbc/BQStatement.java b/src/main/java/net/starschema/clouddb/jdbc/BQStatement.java
@@ -110,6 +110,10 @@ public BQStatement(String projectId, BQConnection bqConnection,
 
     }
 
+    protected long getSyncTimeoutMillis() {
+        return SYNC_TIMEOUT_MILLIS;
+    }
+
     public void setLabels(Map<String, String> statementLabels) {
         this.statementLabels = ImmutableMap.copyOf(statementLabels);
     }
@@ -297,7 +301,7 @@ protected QueryResponse runSyncQuery(String querySql, boolean unlimitedBillingBy
                         connection.getDataSet(),
                         this.connection.getUseLegacySql(),
                         !unlimitedBillingBytes ? this.connection.getMaxBillingBytes() : null,
-                        SYNC_TIMEOUT_MILLIS, // we need this to respond fast enough to avoid any socket timeouts
+                        getSyncTimeoutMillis(), // we need this to respond fast enough to avoid any socket timeouts
                         (long) getMaxRows(),
                         allLabels);
                 syncResponseFromCurrentQuery.set(resp);
@@ -347,7 +351,7 @@ public void cancel() throws SQLException {
         } else if (currentlyRunningSyncThread != null) {
             // The sync part of the query has not completed yet: wait for it so we can find the job to cancel
             try {
-                currentlyRunningSyncThread.join(SYNC_TIMEOUT_MILLIS);
+                currentlyRunningSyncThread.join(getSyncTimeoutMillis());
                 QueryResponse resp = syncResponseFromCurrentQuery.get();
                 if (resp != null && !resp.getJobComplete()) { // Don't bother cancel if the job is complete
                     jobRefToCancel = resp.getJobReference();

diff --git a/src/test/java/net/starschema/clouddb/jdbc/BQForwardOnlyResultSetFunctionTest.java b/src/test/java/net/starschema/clouddb/jdbc/BQForwardOnlyResultSetFunctionTest.java
@@ -26,6 +26,7 @@
 
 import com.google.common.reflect.TypeToken;
 import com.google.gson.Gson;
+import java.io.IOException;
 import junit.framework.Assert;
 import org.junit.Before;
 import org.junit.Test;
@@ -55,6 +56,13 @@ public class BQForwardOnlyResultSetFunctionTest {
     Logger logger = LoggerFactory.getLogger(BQForwardOnlyResultSetFunctionTest.class);
     private Integer maxRows = null;
 
+    private BQConnection conn()  throws SQLException, IOException {
+        String url = BQSupportFuncts.constructUrlFromPropertiesFile(BQSupportFuncts
+            .readFromPropFile(getClass().getResource("/installedaccount.properties").getFile()), true, null);
+        url += "&useLegacySql=false";
+        return new BQConnection(url, new Properties());
+    }
+
     @Test
     public void ChainedCursorFunctionTest() {
         this.QueryLoad();
@@ -615,4 +623,34 @@ public void testResultSetProcedures() throws SQLException, ParseException {
 
 		System.out.println(result.toString());
 	}
+
+    @Test
+    public void testResultSetProceduresAsync() throws SQLException {
+        final String sql = "CREATE PROCEDURE looker_test.long_procedure(target_id INT64)\n"
+            + "BEGIN\n"
+            + "END;";
+        this.NewConnection(false);
+
+        try {
+            BQConnection bq = conn();
+            BQStatement stmt = new BQStatement(bq.getProjectId(), bq, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY) {
+                @Override
+                protected long getSyncTimeoutMillis() {
+                    return 0; // force async condition
+                }
+            };
+
+            stmt.setQueryTimeout(500);
+            stmt.executeQuery(sql);
+        } catch (SQLException | IOException e) {
+            this.logger.error("SQLexception" + e.toString());
+            Assert.fail("SQLException" + e.toString());
+        } finally {
+            String cleanupSql = "DROP PROCEDURE looker_test.long_procedure;\n";
+            Statement stmt = BQForwardOnlyResultSetFunctionTest.con
+                .createStatement(ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY);
+            stmt.setQueryTimeout(500);
+            stmt.executeQuery(cleanupSql);
+        }
+    }
 }