Might as well check all the getters (#94)

looker-open-source · Jul 2, 2021 · ea48fc8 · ea48fc8
1 parent d75fb42
commit ea48fc8
Showing 1 changed file with 17 additions and 5 deletions.
diff --git a/src/main/java/net/starschema/clouddb/jdbc/BQStatementRoot.java b/src/main/java/net/starschema/clouddb/jdbc/BQStatementRoot.java
@@ -27,6 +27,8 @@
 
 package net.starschema.clouddb.jdbc;
 
+import com.google.api.services.bigquery.model.TableSchema;
+import com.google.api.services.bigquery.model.TableRow;
 import com.google.api.services.bigquery.model.Job;
 import com.google.api.services.bigquery.model.QueryResponse;
 import org.slf4j.Logger;
@@ -35,6 +37,8 @@
 import java.io.IOException;
 import java.math.BigInteger;
 import java.sql.*;
+import java.util.ArrayList;
+import java.util.List;
 
 /**
  * This class partially implements java.sql.Statement, and
@@ -267,9 +271,9 @@ private int executeDML(String sql) throws SQLException {
                     this.connection.getLabels()
             );
 
-            if (qr.getJobComplete()) {
+            if (defaultValueIfNull(qr.getJobComplete(), false)) {
                 // I hope they don't insert more than 2^32-1 :)
-                return Math.toIntExact(qr.getNumDmlAffectedRows());
+                return Math.toIntExact(defaultValueIfNull(qr.getNumDmlAffectedRows(), 0L));
             }
 
             referencedJob = this.connection.getBigquery().jobs().get(projectId, qr.getJobReference().getJobId()).execute();
@@ -324,9 +328,13 @@ public ResultSet executeQuery(String querySql, boolean unlimitedBillingBytes) th
                     (long) getMaxRows(),
                     this.connection.getLabels()
             );
-            if (qr.getJobComplete()) {
-                if (qr.getTotalRows().equals(BigInteger.valueOf(qr.getRows().size()))) {
-                    return new BQScrollableResultSet(qr.getRows(), this, qr.getSchema(), qr.getTotalBytesProcessed(), qr.getCacheHit());
+            if (defaultValueIfNull(qr.getJobComplete(), false)) {
+                List<TableRow> rows = defaultValueIfNull(qr.getRows(), new ArrayList<TableRow>());
+                if (BigInteger.valueOf(rows.size()).equals(qr.getTotalRows())) {
+                    Boolean cacheHit = defaultValueIfNull(qr.getCacheHit(), false);
+                    Long totalBytesProcessed = defaultValueIfNull(qr.getTotalBytesProcessed(), 0L);
+                    TableSchema schema = defaultValueIfNull(qr.getSchema(), new TableSchema());
+                    return new BQScrollableResultSet(rows, this, schema, totalBytesProcessed, cacheHit);
                 }
                 jobAlreadyCompleted = true;
             }
@@ -374,6 +382,10 @@ public ResultSet executeQuery(String querySql, boolean unlimitedBillingBytes) th
                 "Query run took more than the specified timeout");
     }
 
+    private static <T> T defaultValueIfNull(T value, T defaultValue) {
+        return value == null ? defaultValue : value;
+    }
+
     /**
      * Only supported for standard SQL (non-legacy)
      *