Skip to content

Commit

Permalink
feat: add support for Row Affinity app profiles (#2341)
Browse files Browse the repository at this point in the history 
* 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
  • Loading branch information
@djyau
djyau authored Dec 4, 2024
1 parent 7bd3a30 commit cb4d60e
Show file tree
Hide file tree
Showing 6 changed files with 311 additions and 1 deletion.
33 changes: 32 additions & 1 deletion ...le-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/models/AppProfile.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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()) {
Expand Down Expand Up @@ -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.
Expand Down
149 changes: 149 additions & 0 deletions ...le/src/test/java/com/google/cloud/bigtable/admin/v2/BigtableInstanceAdminClientTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand Down
37 changes: 37 additions & 0 deletions ...le/src/test/java/com/google/cloud/bigtable/admin/v2/it/BigtableInstanceAdminClientIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
50 changes: 50 additions & 0 deletions ...loud-bigtable/src/test/java/com/google/cloud/bigtable/admin/v2/models/AppProfileTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"));
}
}
13 changes: 13 additions & 0 deletions .../src/test/java/com/google/cloud/bigtable/admin/v2/models/CreateAppProfileRequestTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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());
}
}
30 changes: 30 additions & 0 deletions .../src/test/java/com/google/cloud/bigtable/admin/v2/models/UpdateAppProfileRequestTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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());
}
}

0 comments on commit cb4d60e

Please sign in to comment.