From 59a7329e5db9a651851189dfac53f3d2a7428dbb Mon Sep 17 00:00:00 2001
From: Mike Kaplinskiy <mike@ladderlife.com>
Date: Sun, 15 Apr 2018 14:55:56 -0700
Subject: [PATCH] Expose the schema of dry run query jobs.

---
 .../google/cloud/bigquery/JobStatistics.java  | 28 +++++++++++++++++--
 .../cloud/bigquery/JobStatisticsTest.java     |  3 ++
 2 files changed, 29 insertions(+), 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
index 49ce6302592e..05e0856e5235 100644
--- 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,6 +448,15 @@ 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()
@@ -444,7 +464,8 @@ ToStringHelper toStringHelper() {
           .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
index cf4e9c424b27..16537fc17b81 100644
--- 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) {