From ef80cbd8ff5dea116371c90a742a1fa4a57dc0e7 Mon Sep 17 00:00:00 2001
From: Ryan Culbertson <ryan@spotify.com>
Date: Sun, 23 Jun 2019 13:07:54 -0400
Subject: [PATCH] Allow Bigtable Query to be copied

Added a new create method to the Bigtable Query object that behaves like a
copy constructor. It takes an existing Query object and performs a deep copy.

This is to address #5550, but this could be useful in other situations as well.
---
 .../cloud/bigtable/data/v2/models/Query.java     | 16 +++++++++++-----
 .../cloud/bigtable/data/v2/models/QueryTest.java |  7 +++++++
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/google-cloud-clients/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java b/google-cloud-clients/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java
index 5288ee0109a2..64b1f9c4f51a 100644
--- a/google-cloud-clients/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java
+++ b/google-cloud-clients/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java
@@ -43,7 +43,7 @@ public final class Query implements Serializable {
   private static final long serialVersionUID = -316972783499434755L;
 
   private final String tableId;
-  private transient ReadRowsRequest.Builder builder = ReadRowsRequest.newBuilder();
+  private transient ReadRowsRequest.Builder builder;
 
   /**
    * Constructs a new Query object for the specified table id. The table id will be combined with
@@ -51,11 +51,17 @@ public final class Query implements Serializable {
    * com.google.cloud.bigtable.data.v2.BigtableDataSettings}.
    */
   public static Query create(String tableId) {
-    return new Query(tableId);
+    return new Query(tableId, ReadRowsRequest.newBuilder());
   }
 
-  private Query(String tableId) {
+  /** Constructs a new Query object copying the state of the specified Query. */
+  public static Query create(Query query) {
+    return new Query(query.tableId, query.builder);
+  }
+
+  private Query(String tableId, ReadRowsRequest.Builder builder) {
     this.tableId = tableId;
+    this.builder = ReadRowsRequest.newBuilder(builder.buildPartial());
   }
 
   private void readObject(ObjectInputStream input) throws IOException, ClassNotFoundException {
@@ -229,7 +235,7 @@ public List<Query> shard(SortedSet<ByteString> splitPoints) {
     List<Query> shards = Lists.newArrayListWithCapacity(shardedRowSets.size());
 
     for (RowSet rowSet : shardedRowSets) {
-      Query queryShard = new Query(tableId);
+      Query queryShard = Query.create(tableId);
       queryShard.builder.mergeFrom(this.builder.build());
       queryShard.builder.setRows(rowSet);
       shards.add(queryShard);
@@ -268,7 +274,7 @@ public ReadRowsRequest toProto(RequestContext requestContext) {
   public static Query fromProto(@Nonnull ReadRowsRequest request) {
     Preconditions.checkArgument(request != null, "ReadRowsRequest must not be null");
 
-    Query query = new Query(NameUtil.extractTableIdFromTableName(request.getTableName()));
+    Query query = Query.create(NameUtil.extractTableIdFromTableName(request.getTableName()));
     query.builder = request.toBuilder();
 
     return query;
diff --git a/google-cloud-clients/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java b/google-cloud-clients/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java
index 857a90916d14..f6b3d7b52c70 100644
--- a/google-cloud-clients/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java
+++ b/google-cloud-clients/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java
@@ -67,6 +67,13 @@ public void requestContextTest() {
     assertThat(proto).isEqualTo(expectedProtoBuilder().build());
   }
 
+  @Test
+  public void copyTest() {
+    Query original = Query.create(TABLE_ID).limit(1).range("a", "z");
+    Query copy = Query.create(original);
+    assertThat(copy).isEqualTo(original);
+  }
+
   @Test
   public void rowKeysTest() {
     Query query =