Added Query#fromProto to convert protobuf ReadRowsRequest (#4291)

rahulKQL · JesseLovelace · commit 3d1165982108 · 2019-01-09T13:49:19.000-08:00
* Adding Query.fromProto

* Adding Query.fromProto

* adding changes as per feedback

* updated javadoc as per feedback

* reformatted Query &amp; QueryTest.java
diff --git a/google-cloud-clients/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/NameUtil.java b/google-cloud-clients/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/NameUtil.java
@@ -16,6 +16,8 @@
 package com.google.cloud.bigtable.data.v2.internal;
 
 import com.google.api.core.InternalApi;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 import javax.annotation.Nonnull;
 
 /**
@@ -26,6 +28,9 @@
  */
 @InternalApi
 public class NameUtil {
+  private static final Pattern TABLE_PATTERN =
+      Pattern.compile("projects/([^/]+)/instances/([^/]+)/tables/([^/]+)");
+
   public static String formatInstanceName(@Nonnull String projectId, @Nonnull String instanceId) {
     return "projects/" + projectId + "/instances/" + instanceId;
   }
@@ -34,4 +39,12 @@ public static String formatTableName(
       @Nonnull String projectId, @Nonnull String instanceId, @Nonnull String tableId) {
     return formatInstanceName(projectId, instanceId) + "/tables/" + tableId;
   }
+
+  public static String extractTableIdFromTableName(@Nonnull String fullTableName) {
+    Matcher matcher = TABLE_PATTERN.matcher(fullTableName);
+    if (!matcher.matches()) {
+      throw new IllegalArgumentException("Invalid table name: " + fullTableName);
+    }
+    return matcher.group(3);
+  }
 }
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
@@ -36,6 +36,7 @@
 import java.io.Serializable;
 import java.util.List;
 import java.util.SortedSet;
+import javax.annotation.Nonnull;
 
 /** A simple wrapper to construct a query for the ReadRows RPC. */
 public final class Query implements Serializable {
@@ -258,6 +259,21 @@ public ReadRowsRequest toProto(RequestContext requestContext) {
         .build();
   }
 
+  /**
+   * Wraps the protobuf {@link ReadRowsRequest}.
+   *
+   * <p>WARNING: Please note that the project id & instance id in the table name will be overwritten
+   * by the configuration in the BigtableDataClient.
+   */
+  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.builder = request.toBuilder();
+
+    return query;
+  }
+
   private static ByteString wrapKey(String key) {
     if (key == null) {
       return null;
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
@@ -38,7 +38,9 @@
 import java.util.List;
 import java.util.SortedSet;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.ExpectedException;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
@@ -50,6 +52,8 @@ public class QueryTest {
   private static final String APP_PROFILE_ID = "fake-profile-id";
   private RequestContext requestContext;
 
+  @Rule public ExpectedException expect = ExpectedException.none();
+
   @Before
   public void setUp() {
     requestContext = RequestContext.create(PROJECT_ID, INSTANCE_ID, APP_PROFILE_ID);
@@ -226,4 +230,32 @@ private static ReadRowsRequest.Builder expectedProtoBuilder() {
         .setTableName(NameUtil.formatTableName(PROJECT_ID, INSTANCE_ID, TABLE_ID))
         .setAppProfileId(APP_PROFILE_ID);
   }
+
+  @Test
+  public void testFromProto() {
+    ReadRowsRequest request =
+        ReadRowsRequest.newBuilder()
+            .setTableName(NameUtil.formatTableName(PROJECT_ID, INSTANCE_ID, TABLE_ID))
+            .setAppProfileId(APP_PROFILE_ID)
+            .setFilter(RowFilter.newBuilder().setRowKeyRegexFilter(ByteString.copyFromUtf8(".*")))
+            .setRows(
+                RowSet.newBuilder()
+                    .addRowKeys(ByteString.copyFromUtf8("row-key"))
+                    .addRowRanges(
+                        RowRange.newBuilder()
+                            .setStartKeyClosed(ByteString.copyFromUtf8("j"))
+                            .setEndKeyClosed(ByteString.copyFromUtf8("z"))))
+            .build();
+    Query query = Query.fromProto(request);
+
+    assertThat(query.toProto(requestContext)).isEqualTo(request);
+  }
+
+  @Test(expected = IllegalArgumentException.class)
+  public void testFromProtoWithEmptyTableId() {
+    Query.fromProto(ReadRowsRequest.getDefaultInstance());
+
+    expect.expect(IllegalArgumentException.class);
+    expect.expectMessage("Invalid table name:");
+  }
 }
