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 index bd7a534640..2507ef4dd3 100644 --- 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 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 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 index d8522db71a..388631d93a 100644 --- 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 clusterIds = new HashSet(); + 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 index 76413165bd..6e9f82b5b4 100644 --- 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 @@ -279,6 +279,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 index 8215e5f8fc..d6e6e410e8 100644 --- 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 index 088dc2bcfe..7e9cc81541 100644 --- 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 index 04cf3f0813..603943c533 100644 --- 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()); + } }