Add numDmlAffectedRows to QueryResponse (#2320)

saicheems · web-flow · commit d3095c7df003 · 2017-08-08T16:53:08.000-07:00
diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java
@@ -571,6 +571,7 @@ public com.google.api.services.bigquery.model.QueryResponse call() {
       QueryResponse.Builder builder = QueryResponse.newBuilder();
       JobId completeJobId = JobId.fromPb(results.getJobReference());
       builder.setJobId(completeJobId);
+      builder.setNumDmlAffectedRows(results.getNumDmlAffectedRows());
       builder.setJobCompleted(results.getJobComplete());
       List<TableRow> rowsPb = results.getRows();
       if (results.getJobComplete()) {
@@ -617,6 +618,7 @@ public GetQueryResultsResponse call() {
           }, serviceOptions.getRetrySettings(), EXCEPTION_HANDLER, serviceOptions.getClock());
       QueryResponse.Builder builder = QueryResponse.newBuilder();
       builder.setJobId(JobId.fromPb(results.getJobReference()));
+      builder.setNumDmlAffectedRows(results.getNumDmlAffectedRows());
       builder.setEtag(results.getEtag());
       builder.setJobCompleted(results.getJobComplete());
       List<TableRow> rowsPb = results.getRows();
diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryResponse.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryResponse.java
@@ -56,6 +56,7 @@ public class QueryResponse implements Serializable {
   private final QueryResult result;
   private final String etag;
   private final JobId jobId;
+  private final Long numDmlAffectedRows;
   private final boolean jobCompleted;
   private final List<BigQueryError> executionErrors;
 
@@ -64,6 +65,7 @@ static final class Builder {
     private QueryResult result;
     private String etag;
     private JobId jobId;
+    private Long numDmlAffectedRows;
     private boolean jobCompleted;
     private List<BigQueryError> executionErrors;
 
@@ -84,6 +86,11 @@ Builder setJobId(JobId jobId) {
       return this;
     }
 
+    Builder setNumDmlAffectedRows(Long numDmlAffectedRows) {
+      this.numDmlAffectedRows = numDmlAffectedRows;
+      return this;
+    }
+
     Builder setJobCompleted(boolean jobCompleted) {
       this.jobCompleted = jobCompleted;
       return this;
@@ -103,6 +110,7 @@ private QueryResponse(Builder builder) {
     this.result = builder.result;
     this.etag = builder.etag;
     this.jobId = builder.jobId;
+    this.numDmlAffectedRows = builder.numDmlAffectedRows;
     this.jobCompleted = builder.jobCompleted;
     this.executionErrors = builder.executionErrors != null ? builder.executionErrors
       : ImmutableList.<BigQueryError>of();
@@ -134,6 +142,12 @@ public JobId getJobId() {
     return jobId;
   }
 
+  /**
+   * Returns the number of rows affected by a DML statement. Present only for DML statements INSERT,
+   * UPDATE or DELETE.
+   */
+  public Long getNumDmlAffectedRows() { return numDmlAffectedRows; }
+
   /**
    * Returns whether the job running the query has completed or not. If {@link #getResult()} is not
    * {@code null}, this method will always return {@code true}. If this method returns {@code false}
@@ -168,14 +182,15 @@ public String toString() {
         .add("result", result)
         .add("etag", etag)
         .add("jobId", jobId)
+        .add("numDmlAffectedRows", numDmlAffectedRows)
         .add("jobCompleted", jobCompleted)
         .add("executionErrors", executionErrors)
         .toString();
   }
 
   @Override
   public final int hashCode() {
-    return Objects.hash(jobId);
+    return Objects.hash(result, etag, jobId, numDmlAffectedRows, jobCompleted, executionErrors);
   }
 
   @Override
@@ -191,6 +206,7 @@ public final boolean equals(Object obj) {
         && Objects.equals(etag, response.etag)
         && Objects.equals(result, response.result)
         && Objects.equals(jobId, response.jobId)
+        && Objects.equals(numDmlAffectedRows, response.numDmlAffectedRows)
         && Objects.equals(executionErrors, response.executionErrors);
   }
 
diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryResponseTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryResponseTest.java
@@ -38,6 +38,7 @@ public class QueryResponseTest {
 
   private static final Schema SCHEMA = Schema.of(FIELD_SCHEMA1);
   private static final JobId JOB_ID = JobId.of("project", "job");
+  private static final Long NUM_DML_AFFECTED_ROWS = 24L;
   private static final Long TOTAL_ROWS = 42L;
   private static final QueryResult.QueryResultsPageFetcher FETCHER =
       new QueryResult.QueryResultsPageFetcher() {
@@ -66,6 +67,7 @@ public QueryResult getNextPage() {
   private static final QueryResponse QUERY_RESPONSE = QueryResponse.newBuilder()
       .setEtag(ETAG)
       .setJobId(JOB_ID)
+      .setNumDmlAffectedRows(NUM_DML_AFFECTED_ROWS)
       .setJobCompleted(JOB_COMPLETE)
       .setExecutionErrors(ERRORS)
       .setResult(QUERY_RESULT)
@@ -76,6 +78,7 @@ public void testBuilder() {
     assertEquals(ETAG, QUERY_RESPONSE.getEtag());
     assertEquals(QUERY_RESULT, QUERY_RESPONSE.getResult());
     assertEquals(JOB_ID, QUERY_RESPONSE.getJobId());
+    assertEquals(NUM_DML_AFFECTED_ROWS, QUERY_RESPONSE.getNumDmlAffectedRows());
     assertEquals(JOB_COMPLETE, QUERY_RESPONSE.jobCompleted());
     assertEquals(ERRORS, QUERY_RESPONSE.getExecutionErrors());
     assertTrue(QUERY_RESPONSE.hasErrors());
@@ -88,6 +91,7 @@ public void testBuilderIncomplete() {
     assertNull(queryResponse.getEtag());
     assertNull(queryResponse.getResult());
     assertNull(queryResponse.getJobId());
+    assertNull(queryResponse.getNumDmlAffectedRows());
     assertFalse(queryResponse.jobCompleted());
     assertEquals(ImmutableList.<BigQueryError>of(), queryResponse.getExecutionErrors());
     assertFalse(queryResponse.hasErrors());
@@ -103,6 +107,7 @@ private void compareQueryResponse(QueryResponse expected, QueryResponse value) {
     assertEquals(expected.getEtag(), value.getEtag());
     assertEquals(expected.getResult(), value.getResult());
     assertEquals(expected.getJobId(), value.getJobId());
+    assertEquals(expected.getNumDmlAffectedRows(), value.getNumDmlAffectedRows());
     assertEquals(expected.jobCompleted(), value.jobCompleted());
     assertEquals(expected.getExecutionErrors(), value.getExecutionErrors());
     assertEquals(expected.hasErrors(), value.hasErrors());
diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/SerializationTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/SerializationTest.java
@@ -160,6 +160,7 @@ public class SerializationTest extends BaseSerializationTest {
   private static final JobStatus JOB_STATUS = new JobStatus(JobStatus.State.DONE, BIGQUERY_ERROR,
       ImmutableList.of(BIGQUERY_ERROR));
   private static final JobId JOB_ID = JobId.of("project", "job");
+  private static final Long NUM_DL_AFFECTED_ROWS = 24L;
   private static final CopyJobConfiguration COPY_JOB_CONFIGURATION =
       CopyJobConfiguration.of(TABLE_ID, TABLE_ID);
   private static final ExtractJobConfiguration EXTRACT_JOB_CONFIGURATION =
@@ -211,6 +212,7 @@ public class SerializationTest extends BaseSerializationTest {
   private static final QueryResponse QUERY_RESPONSE = QueryResponse.newBuilder()
       .setEtag(ETAG)
       .setJobId(JOB_ID)
+      .setNumDmlAffectedRows(NUM_DL_AFFECTED_ROWS)
       .setJobCompleted(true)
       .setResult(QUERY_RESULT)
       .build();
