feat: add support for Row Affinity app profiles (#2341)

djyau · web-flow · commit cb4d60e8ce20 · 2024-12-04T10:35:53.000-05:00
* Create of methods, add unit tests

Change-Id: I7cdd6c9ce85f9132c0fa3db3aa8c70abac26f12b

* Add unit tests for BigtableInstanceAdminClientTests

Change-Id: Id2cedefd90c52c249b2b15734eb1fbbcc5db2bf4

* Add tests for Create and Update App Profile

Change-Id: I01ef72617d391bdb956ae9f469f1fd5af208a2e1

* Add integration test

Change-Id: I9f18991cb61d99f30b811cea0988963d9e2577cd

* Rename ofWithRowAffinity to withRowAffinity

Change-Id: I9f88209e04cde4f628c878fbdb659c7344aa1c3c

* Run mvn com.coveo:fmt-maven-plugin:format

Change-Id: Ieda39f0d9825fae649755350b4f228fe0e8985e5

* Remove extra test for cluster IDs - one is enough

Change-Id: I6d8e7bd644ce2cd8f098c8ee37b4eec1a0d642e2

* Add unit test for row affinity with set of strings

Change-Id: I7b072214c8b38646058d0ec668bbb469bf16f23e

* Run formatter

Change-Id: Ieb04b2141c8843ff59cc4d756d736b4c00609ee2
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/models/AppProfile.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/models/AppProfile.java
@@ -18,7 +18,6 @@
 import com.google.api.core.InternalApi;
 import com.google.bigtable.admin.v2.AppProfile.DataBoostIsolationReadOnly;
 import com.google.bigtable.admin.v2.AppProfile.MultiClusterRoutingUseAny;
-import com.google.bigtable.admin.v2.AppProfile.Priority;
 import com.google.bigtable.admin.v2.AppProfile.StandardIsolation;
 import com.google.bigtable.admin.v2.AppProfileName;
 import com.google.common.base.Objects;
@@ -69,6 +68,10 @@ private AppProfile(@Nonnull com.google.bigtable.admin.v2.AppProfile proto) {
   @SuppressWarnings("WeakerAccess")
   public RoutingPolicy getPolicy() {
     if (proto.hasMultiClusterRoutingUseAny()) {
+      if (proto.getMultiClusterRoutingUseAny().hasRowAffinity()) {
+        return MultiClusterRoutingPolicy.withRowAffinity(
+            ImmutableSet.copyOf(proto.getMultiClusterRoutingUseAny().getClusterIdsList()));
+      }
       return MultiClusterRoutingPolicy.of(
           ImmutableSet.copyOf(proto.getMultiClusterRoutingUseAny().getClusterIdsList()));
     } else if (proto.hasSingleClusterRouting()) {
@@ -267,6 +270,34 @@ public static MultiClusterRoutingPolicy of(Set<String> clusterIds) {
           MultiClusterRoutingUseAny.newBuilder().addAllClusterIds(clusterIds).build());
     }
 
+    /** Creates a new instance of {@link MultiClusterRoutingPolicy}. */
+    public static MultiClusterRoutingPolicy withRowAffinity() {
+      return new MultiClusterRoutingPolicy(
+          MultiClusterRoutingUseAny.newBuilder()
+              .setRowAffinity(MultiClusterRoutingUseAny.RowAffinity.getDefaultInstance())
+              .build());
+    }
+
+    /**
+     * Creates a new instance of {@link MultiClusterRoutingPolicy} with row affinity enabled and
+     * specified cluster ids to route to.
+     */
+    public static MultiClusterRoutingPolicy withRowAffinity(String... clusterIds) {
+      return withRowAffinity(ImmutableSet.copyOf(clusterIds));
+    }
+
+    /**
+     * Creates a new instance of {@link MultiClusterRoutingPolicy} with specified cluster ids to
+     * route to.
+     */
+    public static MultiClusterRoutingPolicy withRowAffinity(Set<String> clusterIds) {
+      return new MultiClusterRoutingPolicy(
+          MultiClusterRoutingUseAny.newBuilder()
+              .addAllClusterIds(clusterIds)
+              .setRowAffinity(MultiClusterRoutingUseAny.RowAffinity.getDefaultInstance())
+              .build());
+    }
+
     /*
      * Returns the set of clusters to route to. The order is ignored; clusters will be
      * tried in order of distance. If empty, all clusters are eligible.
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/admin/v2/BigtableInstanceAdminClientTests.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/admin/v2/BigtableInstanceAdminClientTests.java
@@ -65,7 +65,9 @@
 import com.google.protobuf.FieldMask;
 import io.grpc.Status;
 import io.grpc.Status.Code;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import org.junit.Before;
@@ -1034,6 +1036,153 @@ public void testCreateAppProfileAddPriority() {
     assertThat(actualResult).isEqualTo(AppProfile.fromProto(expectedResponse));
   }
 
+  @Test
+  public void testCreateAppProfileAddRowAffinity() {
+    // Setup
+    Mockito.when(mockStub.createAppProfileCallable()).thenReturn(mockCreateAppProfileCallable);
+
+    com.google.bigtable.admin.v2.CreateAppProfileRequest expectedRequest =
+        com.google.bigtable.admin.v2.CreateAppProfileRequest.newBuilder()
+            .setParent(NameUtil.formatInstanceName(PROJECT_ID, INSTANCE_ID))
+            .setAppProfileId(APP_PROFILE_ID)
+            .setAppProfile(
+                com.google.bigtable.admin.v2.AppProfile.newBuilder()
+                    .setDescription("my description")
+                    .setMultiClusterRoutingUseAny(
+                        com.google.bigtable.admin.v2.AppProfile.MultiClusterRoutingUseAny
+                            .newBuilder()
+                            .setRowAffinity(
+                                com.google.bigtable.admin.v2.AppProfile.MultiClusterRoutingUseAny
+                                    .RowAffinity.getDefaultInstance())))
+            .build();
+
+    com.google.bigtable.admin.v2.AppProfile expectedResponse =
+        com.google.bigtable.admin.v2.AppProfile.newBuilder()
+            .setName(APP_PROFILE_NAME)
+            .setDescription("my description")
+            .setMultiClusterRoutingUseAny(
+                com.google.bigtable.admin.v2.AppProfile.MultiClusterRoutingUseAny.newBuilder()
+                    .setRowAffinity(
+                        com.google.bigtable.admin.v2.AppProfile.MultiClusterRoutingUseAny
+                            .RowAffinity.getDefaultInstance()))
+            .build();
+
+    Mockito.when(mockCreateAppProfileCallable.futureCall(expectedRequest))
+        .thenReturn(ApiFutures.immediateFuture(expectedResponse));
+
+    // Execute
+    AppProfile actualResult =
+        adminClient.createAppProfile(
+            CreateAppProfileRequest.of(INSTANCE_ID, APP_PROFILE_ID)
+                .setDescription("my description")
+                .setRoutingPolicy(MultiClusterRoutingPolicy.withRowAffinity()));
+
+    // Verify
+    assertThat(actualResult).isEqualTo(AppProfile.fromProto(expectedResponse));
+  }
+
+  @Test
+  public void testCreateAppProfileAddRowAffinityAddMultipleClusterIds() {
+    // Setup
+    Mockito.when(mockStub.createAppProfileCallable()).thenReturn(mockCreateAppProfileCallable);
+
+    com.google.bigtable.admin.v2.CreateAppProfileRequest expectedRequest =
+        com.google.bigtable.admin.v2.CreateAppProfileRequest.newBuilder()
+            .setParent(NameUtil.formatInstanceName(PROJECT_ID, INSTANCE_ID))
+            .setAppProfileId(APP_PROFILE_ID)
+            .setAppProfile(
+                com.google.bigtable.admin.v2.AppProfile.newBuilder()
+                    .setDescription("my description")
+                    .setMultiClusterRoutingUseAny(
+                        com.google.bigtable.admin.v2.AppProfile.MultiClusterRoutingUseAny
+                            .newBuilder()
+                            .addClusterIds("cluster-id-1")
+                            .addClusterIds("cluster-id-2")
+                            .setRowAffinity(
+                                com.google.bigtable.admin.v2.AppProfile.MultiClusterRoutingUseAny
+                                    .RowAffinity.getDefaultInstance())))
+            .build();
+
+    com.google.bigtable.admin.v2.AppProfile expectedResponse =
+        com.google.bigtable.admin.v2.AppProfile.newBuilder()
+            .setName(APP_PROFILE_NAME)
+            .setDescription("my description")
+            .setMultiClusterRoutingUseAny(
+                com.google.bigtable.admin.v2.AppProfile.MultiClusterRoutingUseAny.newBuilder()
+                    .addClusterIds("cluster-id-1")
+                    .addClusterIds("cluster-id-2")
+                    .setRowAffinity(
+                        com.google.bigtable.admin.v2.AppProfile.MultiClusterRoutingUseAny
+                            .RowAffinity.getDefaultInstance()))
+            .build();
+
+    Mockito.when(mockCreateAppProfileCallable.futureCall(expectedRequest))
+        .thenReturn(ApiFutures.immediateFuture(expectedResponse));
+
+    // Execute
+    AppProfile actualResult =
+        adminClient.createAppProfile(
+            CreateAppProfileRequest.of(INSTANCE_ID, APP_PROFILE_ID)
+                .setDescription("my description")
+                .setRoutingPolicy(
+                    MultiClusterRoutingPolicy.withRowAffinity("cluster-id-1", "cluster-id-2")));
+
+    // Verify
+    assertThat(actualResult).isEqualTo(AppProfile.fromProto(expectedResponse));
+  }
+
+  @Test
+  public void testCreateAppProfileAddRowAffinityAddSetOfClusterIds() {
+    // Setup
+    Mockito.when(mockStub.createAppProfileCallable()).thenReturn(mockCreateAppProfileCallable);
+
+    com.google.bigtable.admin.v2.CreateAppProfileRequest expectedRequest =
+        com.google.bigtable.admin.v2.CreateAppProfileRequest.newBuilder()
+            .setParent(NameUtil.formatInstanceName(PROJECT_ID, INSTANCE_ID))
+            .setAppProfileId(APP_PROFILE_ID)
+            .setAppProfile(
+                com.google.bigtable.admin.v2.AppProfile.newBuilder()
+                    .setDescription("my description")
+                    .setMultiClusterRoutingUseAny(
+                        com.google.bigtable.admin.v2.AppProfile.MultiClusterRoutingUseAny
+                            .newBuilder()
+                            .addClusterIds("cluster-id-1")
+                            .addClusterIds("cluster-id-2")
+                            .setRowAffinity(
+                                com.google.bigtable.admin.v2.AppProfile.MultiClusterRoutingUseAny
+                                    .RowAffinity.getDefaultInstance())))
+            .build();
+
+    com.google.bigtable.admin.v2.AppProfile expectedResponse =
+        com.google.bigtable.admin.v2.AppProfile.newBuilder()
+            .setName(APP_PROFILE_NAME)
+            .setDescription("my description")
+            .setMultiClusterRoutingUseAny(
+                com.google.bigtable.admin.v2.AppProfile.MultiClusterRoutingUseAny.newBuilder()
+                    .addClusterIds("cluster-id-1")
+                    .addClusterIds("cluster-id-2")
+                    .setRowAffinity(
+                        com.google.bigtable.admin.v2.AppProfile.MultiClusterRoutingUseAny
+                            .RowAffinity.getDefaultInstance()))
+            .build();
+
+    Mockito.when(mockCreateAppProfileCallable.futureCall(expectedRequest))
+        .thenReturn(ApiFutures.immediateFuture(expectedResponse));
+
+    // Execute
+    Set<String> clusterIds = new HashSet<String>();
+    clusterIds.add("cluster-id-1");
+    clusterIds.add("cluster-id-2");
+    AppProfile actualResult =
+        adminClient.createAppProfile(
+            CreateAppProfileRequest.of(INSTANCE_ID, APP_PROFILE_ID)
+                .setDescription("my description")
+                .setRoutingPolicy(MultiClusterRoutingPolicy.withRowAffinity(clusterIds)));
+
+    // Verify
+    assertThat(actualResult).isEqualTo(AppProfile.fromProto(expectedResponse));
+  }
+
   @Test
   public void testGetAppProfile() {
     // Setup
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/admin/v2/it/BigtableInstanceAdminClientIT.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/admin/v2/it/BigtableInstanceAdminClientIT.java
@@ -285,6 +285,43 @@ public void appProfileTestDataBoost() {
     }
   }
 
+  @Test
+  public void appProfileTestRowAffinity() {
+    String newInstanceId = prefixGenerator.newPrefix();
+    String newClusterId = newInstanceId + "-c1";
+    String newClusterId2 = newInstanceId + "-c2";
+
+    client.createInstance(
+        CreateInstanceRequest.of(newInstanceId)
+            .addCluster(newClusterId, testEnvRule.env().getPrimaryZone(), 1, StorageType.SSD)
+            .addCluster(newClusterId2, testEnvRule.env().getSecondaryZone(), 1, StorageType.SSD)
+            .setDisplayName("Row-Affinity-Instance-Test")
+            .addLabel("state", "readytodelete")
+            .setType(Type.PRODUCTION));
+
+    try {
+      assertThat(client.exists(newInstanceId)).isTrue();
+
+      String testAppProfile = prefixGenerator.newPrefix();
+
+      CreateAppProfileRequest request =
+          CreateAppProfileRequest.of(newInstanceId, testAppProfile)
+              .setRoutingPolicy(
+                  AppProfile.MultiClusterRoutingPolicy.withRowAffinity(newClusterId, newClusterId2))
+              .setDescription("row affinity app profile");
+
+      AppProfile newlyCreateAppProfile = client.createAppProfile(request);
+      AppProfile.RoutingPolicy routingPolicy = newlyCreateAppProfile.getPolicy();
+      assertThat(routingPolicy)
+          .isEqualTo(
+              AppProfile.MultiClusterRoutingPolicy.withRowAffinity(newClusterId, newClusterId2));
+    } finally {
+      if (client.exists(newInstanceId)) {
+        client.deleteInstance(newInstanceId);
+      }
+    }
+  }
+
   @Test
   public void iamUpdateTest() {
     Policy policy = client.getIamPolicy(instanceId);
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/admin/v2/models/AppProfileTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/admin/v2/models/AppProfileTest.java
@@ -291,4 +291,54 @@ public void testFromProtoWithDataBoostIsolation() {
             AppProfile.DataBoostIsolationReadOnlyPolicy.of(
                 AppProfile.ComputeBillingOwner.UNSPECIFIED));
   }
+
+  @Test
+  public void testFromProtoWithRowAffinityNoClusterGroup() {
+    AppProfile profile =
+        AppProfile.fromProto(
+            com.google.bigtable.admin.v2.AppProfile.newBuilder()
+                .setName(AppProfileName.of("my-project", "my-instance", "my-profile").toString())
+                .setDescription("my description")
+                .setMultiClusterRoutingUseAny(
+                    com.google.bigtable.admin.v2.AppProfile.MultiClusterRoutingUseAny.newBuilder()
+                        .setRowAffinity(
+                            com.google.bigtable.admin.v2.AppProfile.MultiClusterRoutingUseAny
+                                .RowAffinity.getDefaultInstance())
+                        .build())
+                .setEtag("my-etag")
+                .build());
+
+    assertThat(profile.getInstanceId()).isEqualTo("my-instance");
+    assertThat(profile.getId()).isEqualTo("my-profile");
+    assertThat(profile.getDescription()).isEqualTo("my description");
+    System.out.println(profile.getPolicy());
+    System.out.println(AppProfile.MultiClusterRoutingPolicy.withRowAffinity());
+    assertThat(profile.getPolicy())
+        .isEqualTo(AppProfile.MultiClusterRoutingPolicy.withRowAffinity());
+  }
+
+  @Test
+  public void testFromProtoWithRowAffinityClusterGroup() {
+    AppProfile profile =
+        AppProfile.fromProto(
+            com.google.bigtable.admin.v2.AppProfile.newBuilder()
+                .setName(AppProfileName.of("my-project", "my-instance", "my-profile").toString())
+                .setDescription("my description")
+                .setMultiClusterRoutingUseAny(
+                    com.google.bigtable.admin.v2.AppProfile.MultiClusterRoutingUseAny.newBuilder()
+                        .addAllClusterIds(ImmutableList.of("cluster-id-1", "cluster-id-2"))
+                        .setRowAffinity(
+                            com.google.bigtable.admin.v2.AppProfile.MultiClusterRoutingUseAny
+                                .RowAffinity.getDefaultInstance())
+                        .build())
+                .setEtag("my-etag")
+                .build());
+
+    assertThat(profile.getInstanceId()).isEqualTo("my-instance");
+    assertThat(profile.getId()).isEqualTo("my-profile");
+    assertThat(profile.getDescription()).isEqualTo("my description");
+    assertThat(profile.getPolicy())
+        .isEqualTo(
+            AppProfile.MultiClusterRoutingPolicy.withRowAffinity("cluster-id-1", "cluster-id-2"));
+  }
 }
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/admin/v2/models/CreateAppProfileRequestTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/admin/v2/models/CreateAppProfileRequestTest.java
@@ -101,4 +101,17 @@ public void testDataBoostIsolationReadOnly() {
                 .setComputeBillingOwner(DataBoostIsolationReadOnly.ComputeBillingOwner.HOST_PAYS)
                 .build());
   }
+
+  @Test
+  public void testRowAffinity() {
+    CreateAppProfileRequest wrapper =
+        CreateAppProfileRequest.of("my-instance", "my-profile")
+            .setRoutingPolicy(MultiClusterRoutingPolicy.withRowAffinity());
+
+    assertThat(wrapper.toProto("my-project").getAppProfile().getMultiClusterRoutingUseAny())
+        .isEqualTo(
+            MultiClusterRoutingUseAny.newBuilder()
+                .setRowAffinity(MultiClusterRoutingUseAny.RowAffinity.getDefaultInstance())
+                .build());
+  }
 }
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/admin/v2/models/UpdateAppProfileRequestTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/admin/v2/models/UpdateAppProfileRequestTest.java
@@ -146,4 +146,34 @@ public void testUpdateExistingDataBoostIsolationReadOnly() {
                 .setUpdateMask(FieldMask.newBuilder().addPaths("data_boost_isolation_read_only"))
                 .build());
   }
+
+  @Test
+  public void testUpdateRowAffinity() {
+    com.google.bigtable.admin.v2.AppProfile existingProto =
+        com.google.bigtable.admin.v2.AppProfile.newBuilder()
+            .setName("projects/my-project/instances/my-instance/appProfiles/my-profile")
+            .setEtag("my-etag")
+            .setDescription("description")
+            .setMultiClusterRoutingUseAny(MultiClusterRoutingUseAny.getDefaultInstance())
+            .build();
+
+    AppProfile existingWrapper = AppProfile.fromProto(existingProto);
+
+    UpdateAppProfileRequest updateWrapper =
+        UpdateAppProfileRequest.of(existingWrapper)
+            .setRoutingPolicy(AppProfile.MultiClusterRoutingPolicy.withRowAffinity());
+
+    assertThat(updateWrapper.toProto("my-project"))
+        .isEqualTo(
+            com.google.bigtable.admin.v2.UpdateAppProfileRequest.newBuilder()
+                .setAppProfile(
+                    existingProto
+                        .toBuilder()
+                        .setMultiClusterRoutingUseAny(
+                            MultiClusterRoutingUseAny.newBuilder()
+                                .setRowAffinity(
+                                    MultiClusterRoutingUseAny.RowAffinity.getDefaultInstance())))
+                .setUpdateMask(FieldMask.newBuilder().addPaths("multi_cluster_routing_use_any"))
+                .build());
+  }
 }
