Expose the schema of dry run query jobs. (#3159)

googleapis · Apr 17, 2018 · 167ef27 · 167ef27
1 parent 239c456
commit 167ef27
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 2 deletions.
diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java
@@ -323,6 +323,7 @@ public static class QueryStatistics extends JobStatistics {
     private final Long totalBytesBilled;
     private final Long totalBytesProcessed;
     private final List<QueryStage> queryPlan;
+    private final Schema schema;
 
     static final class Builder extends JobStatistics.Builder<QueryStatistics, Builder> {
 
@@ -331,6 +332,7 @@ static final class Builder extends JobStatistics.Builder<QueryStatistics, Builde
       private Long totalBytesBilled;
       private Long totalBytesProcessed;
       private List<QueryStage> queryPlan;
+      private Schema schema;
 
       private Builder() {}
 
@@ -346,6 +348,9 @@ private Builder(com.google.api.services.bigquery.model.JobStatistics statisticsP
                 Lists.transform(
                     statisticsPb.getQuery().getQueryPlan(), QueryStage.FROM_PB_FUNCTION);
           }
+          if (statisticsPb.getQuery().getSchema() != null) {
+            this.schema = Schema.fromPb(statisticsPb.getQuery().getSchema());
+          }
         }
       }
 
@@ -374,6 +379,11 @@ Builder setQueryPlan(List<QueryStage> queryPlan) {
         return self();
       }
 
+      Builder setSchema(Schema schema) {
+        this.schema = schema;
+        return self();
+      }
+
       @Override
       QueryStatistics build() {
         return new QueryStatistics(this);
@@ -387,6 +397,7 @@ private QueryStatistics(Builder builder) {
       this.totalBytesBilled = builder.totalBytesBilled;
       this.totalBytesProcessed = builder.totalBytesProcessed;
       this.queryPlan = builder.queryPlan;
+      this.schema = builder.schema;
     }
 
 
@@ -437,14 +448,24 @@ public List<QueryStage> getQueryPlan() {
       return queryPlan;
     }
 
+
+    /**
+     * Returns the schema for the query result. Present only for successful dry run of
+     * non-legacy SQL queries.
+     */
+    public Schema getSchema() {
+      return schema;
+    }
+
     @Override
     ToStringHelper toStringHelper() {
       return super.toStringHelper()
           .add("billingTier", billingTier)
           .add("cacheHit", cacheHit)
           .add("totalBytesBilled", totalBytesBilled)
           .add("totalBytesProcessed", totalBytesProcessed)
-          .add("queryPlan", queryPlan);
+          .add("queryPlan", queryPlan)
+          .add("schema", schema);
     }
 
     @Override
@@ -458,7 +479,7 @@ public final boolean equals(Object obj) {
     @Override
     public final int hashCode() {
       return Objects.hash(baseHashCode(), billingTier, cacheHit, totalBytesBilled,
-          totalBytesProcessed, queryPlan);
+          totalBytesProcessed, queryPlan, schema);
     }
 
     @Override
@@ -471,6 +492,9 @@ com.google.api.services.bigquery.model.JobStatistics toPb() {
       if (queryPlan != null) {
         queryStatisticsPb.setQueryPlan(Lists.transform(queryPlan, QueryStage.TO_PB_FUNCTION));
       }
+      if (schema != null) {
+        queryStatisticsPb.setSchema(schema.toPb());
+      }
       return super.toPb().setQuery(queryStatisticsPb);
     }
 

diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobStatisticsTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobStatisticsTest.java
@@ -89,6 +89,7 @@ public class JobStatisticsTest {
       .setWriteRatioMax(10.10)
       .build();
   private static final List<QueryStage> QUERY_PLAN = ImmutableList.of(QUERY_STAGE);
+  private static final Schema SCHEMA = Schema.of(Field.of("column", LegacySQLTypeName.DATETIME));
   private static final QueryStatistics QUERY_STATISTICS = QueryStatistics.newBuilder()
       .setCreationTimestamp(CREATION_TIME)
       .setEndTime(END_TIME)
@@ -98,6 +99,7 @@ public class JobStatisticsTest {
       .setTotalBytesBilled(TOTAL_BYTES_BILLED)
       .setTotalBytesProcessed(TOTAL_BYTES_PROCESSED)
       .setQueryPlan(QUERY_PLAN)
+      .setSchema(SCHEMA)
       .build();
   private static final QueryStatistics QUERY_STATISTICS_INCOMPLETE = QueryStatistics.newBuilder()
       .setCreationTimestamp(CREATION_TIME)
@@ -218,6 +220,7 @@ private void compareQueryStatistics(QueryStatistics expected, QueryStatistics va
     assertEquals(expected.getTotalBytesBilled(), value.getTotalBytesBilled());
     assertEquals(expected.getTotalBytesProcessed(), value.getTotalBytesProcessed());
     assertEquals(expected.getQueryPlan(), value.getQueryPlan());
+    assertEquals(expected.getSchema(), value.getSchema());
   }
 
   private void compareStatistics(JobStatistics expected, JobStatistics value) {