@Test @SuppressWarnings("all") public void mutateRowsTest() throws Exception { MutateRowsResponse expectedResponse = MutateRowsResponse.newBuilder().build(); mockBigtable.addResponse(expectedResponse); TableName tableName = TableName.of("[PROJECT]", "[INSTANCE]", "[TABLE]"); List<MutateRowsRequest.Entry> entries = new ArrayList<>(); MutateRowsRequest request = MutateRowsRequest.newBuilder() .setTableName(tableName.toString()) .addAllEntries(entries) .build(); MockStreamObserver<MutateRowsResponse> responseObserver = new MockStreamObserver<>(); ServerStreamingCallable<MutateRowsRequest, MutateRowsResponse> callable = client.mutateRowsCallable(); callable.serverStreamingCall(request, responseObserver); List<MutateRowsResponse> actualResponses = responseObserver.future().get(); Assert.assertEquals(1, actualResponses.size()); Assert.assertEquals(expectedResponse, actualResponses.get(0)); }
private static MutateRowsRequest createRequest(int count) { MutateRowsRequest.Builder request = MutateRowsRequest.newBuilder() .setTableName(NameUtil.formatTableName(PROJECT_ID, INSTANCE_ID, TABLE_ID)); for (int i = 0; i < count; i++) { Builder entry = Entry.newBuilder() .addMutations( Mutation.newBuilder() .setSetCell( SetCell.newBuilder() .setFamilyName("family") .setColumnQualifier(ByteString.copyFromUtf8("col" + i)) .setTimestampMicros(1000) .setValue(ByteString.copyFromUtf8("value" + i)))); request.addEntries(entry); } return request.build(); } }
@Test @SuppressWarnings("all") public void mutateRowsExceptionTest() throws Exception { StatusRuntimeException exception = new StatusRuntimeException(Status.INVALID_ARGUMENT); mockBigtable.addException(exception); TableName tableName = TableName.of("[PROJECT]", "[INSTANCE]", "[TABLE]"); List<MutateRowsRequest.Entry> entries = new ArrayList<>(); MutateRowsRequest request = MutateRowsRequest.newBuilder() .setTableName(tableName.toString()) .addAllEntries(entries) .build(); MockStreamObserver<MutateRowsResponse> responseObserver = new MockStreamObserver<>(); ServerStreamingCallable<MutateRowsRequest, MutateRowsResponse> callable = client.mutateRowsCallable(); callable.serverStreamingCall(request, responseObserver); try { List<MutateRowsResponse> actualResponses = responseObserver.future().get(); Assert.fail("No exception thrown"); } catch (ExecutionException e) { Assert.assertTrue(e.getCause() instanceof InvalidArgumentException); InvalidArgumentException apiException = (InvalidArgumentException) e.getCause(); Assert.assertEquals(StatusCode.Code.INVALID_ARGUMENT, apiException.getStatusCode().getCode()); } }
/** * Creates a new {@link MutateRowsRequest} that's a subset of the original request that * corresponds to a set of indices. * * @param indiciesToRetry * @return the new {@link MutateRowsRequest}. */ private MutateRowsRequest createRetryRequest(List<Integer> indiciesToRetry) { MutateRowsRequest.Builder updatedRequest = MutateRowsRequest.newBuilder() .setTableName(originalRequest.getTableName()); mapToOriginalIndex = new int[indiciesToRetry.size()]; for (int i = 0; i < indiciesToRetry.size(); i++) { mapToOriginalIndex[i] = indiciesToRetry.get(i); updatedRequest.addEntries(originalRequest.getEntries(indiciesToRetry.get(i))); } return updatedRequest.build(); }
/** * Creates a single entry bulk {@link com.google.bigtable.v2.MutateRowsRequest}, which will be * merged by the batching logic in the callable chain. */ @InternalApi public MutateRowsRequest toBulkProto(RequestContext requestContext) { String tableName = NameUtil.formatTableName( requestContext.getProjectId(), requestContext.getInstanceId(), tableId); return MutateRowsRequest.newBuilder() .setAppProfileId(requestContext.getAppProfileId()) .setTableName(tableName) .addEntries( Entry.newBuilder().setRowKey(key).addAllMutations(mutation.getMutations()).build()) .build(); } }
@InternalApi public MutateRowsRequest toProto(RequestContext requestContext) { String tableName = NameUtil.formatTableName( requestContext.getProjectId(), requestContext.getInstanceId(), tableId); return builder .setTableName(tableName) .setAppProfileId(requestContext.getAppProfileId()) .build(); } }
/** * Creates a single entry bulk {@link com.google.bigtable.v2.MutateRowsRequest}, which will be * merged by the batching logic in the callable chain. */ @InternalApi public MutateRowsRequest toBulkProto(RequestContext requestContext) { String tableName = NameUtil.formatTableName( requestContext.getProjectId(), requestContext.getInstanceId(), tableId); return MutateRowsRequest.newBuilder() .setAppProfileId(requestContext.getAppProfileId()) .setTableName(tableName) .addEntries( Entry.newBuilder().setRowKey(key).addAllMutations(mutation.getMutations()).build()) .build(); } }
@Test public void requestBuilderTest() { RequestBuilder<MutateRowsRequest> builder = descriptor.getRequestBuilder(); MutateRowsRequest expected = createRequest(5); for (Entry entry : expected.getEntriesList()) { MutateRowsRequest singleReq = MutateRowsRequest.newBuilder() .setTableName(NameUtil.formatTableName(PROJECT_ID, INSTANCE_ID, TABLE_ID)) .addEntries(entry) .build(); builder.appendRequest(singleReq); } MutateRowsRequest actual = builder.build(); assertThat(actual).isEqualTo(expected); }
@Test public void partitionKeyTest() { String myTableName = NameUtil.formatTableName("my-project", "my-instance", "my-table"); MutateRowsRequest request = createRequest(2).toBuilder().setTableName(myTableName).build(); PartitionKey actual = descriptor.getBatchPartitionKey(request); assertThat(actual).isEqualTo(new PartitionKey(myTableName)); }
@InternalApi public MutateRowsRequest toProto(RequestContext requestContext) { String tableName = NameUtil.formatTableName( requestContext.getProjectId(), requestContext.getInstanceId(), tableId); return builder .setTableName(tableName) .setAppProfileId(requestContext.getAppProfileId()) .build(); } }