use of com.google.cloud.spanner.Mutation in project beam by apache.
the class SpannerIOWriteTest method testBatchableMutationFilterFn_size.
@Test
public void testBatchableMutationFilterFn_size() {
Mutation all = Mutation.delete("test", KeySet.all());
Mutation prefix = Mutation.delete("test", KeySet.prefixRange(Key.of(1L)));
Mutation range = Mutation.delete("test", KeySet.range(KeyRange.openOpen(Key.of(1L), Key.newBuilder().build())));
MutationGroup[] mutationGroups = new MutationGroup[] { g(m(1L)), g(m(2L), m(3L)), // not batchable - too big.
g(m(1L), m(3L), m(4L), m(5L)), g(del(1L)), // not point delete.
g(del(5L, 6L)), g(all), g(prefix), g(range) };
long mutationSize = MutationSizeEstimator.sizeOf(m(1L));
BatchableMutationFilterFn testFn = new BatchableMutationFilterFn(null, null, mutationSize * 3, 1000, 1000);
DoFn<MutationGroup, MutationGroup>.ProcessContext mockProcessContext = Mockito.mock(ProcessContext.class);
when(mockProcessContext.sideInput(any())).thenReturn(getSchema());
// Capture the outputs.
doNothing().when(mockProcessContext).output(mutationGroupCaptor.capture());
doNothing().when(mockProcessContext).output(any(), mutationGroupListCaptor.capture());
// Process all elements.
for (MutationGroup m : mutationGroups) {
when(mockProcessContext.element()).thenReturn(m);
testFn.processElement(mockProcessContext);
}
// Verify captured batchable elements.
assertThat(mutationGroupCaptor.getAllValues(), containsInAnyOrder(g(m(1L)), g(m(2L), m(3L)), g(del(1L))));
// Verify captured unbatchable mutations
Iterable<MutationGroup> unbatchableMutations = Iterables.concat(mutationGroupListCaptor.getAllValues());
assertThat(unbatchableMutations, containsInAnyOrder(// not batchable - too big.
g(m(1L), m(3L), m(4L), m(5L)), // not point delete.
g(del(5L, 6L)), g(all), g(prefix), g(range)));
}
use of com.google.cloud.spanner.Mutation in project beam by apache.
the class SpannerIOWriteTest method streamingWritesWithGroupingWithPriority.
@Test
public void streamingWritesWithGroupingWithPriority() throws Exception {
// verify that grouping/sorting occurs when set.
TestStream<Mutation> testStream = TestStream.create(SerializableCoder.of(Mutation.class)).addElements(m(1L), m(5L), m(2L), m(4L), m(3L), m(6L)).advanceWatermarkToInfinity();
Write write = SpannerIO.write().withProjectId("test-project").withInstanceId("test-instance").withDatabaseId("test-database").withServiceFactory(serviceFactory).withGroupingFactor(40).withMaxNumRows(2).withLowPriority();
pipeline.apply(testStream).apply(write);
pipeline.run();
assertEquals(RpcPriority.LOW, write.getSpannerConfig().getRpcPriority().get());
// Output should be batches of sorted mutations.
verifyBatches(batch(m(1L), m(2L)), batch(m(3L), m(4L)), batch(m(5L), m(6L)));
}
use of com.google.cloud.spanner.Mutation in project beam by apache.
the class SpannerIOWriteTest method testBatchableMutationFilterFn_rows.
@Test
public void testBatchableMutationFilterFn_rows() {
Mutation all = Mutation.delete("test", KeySet.all());
Mutation prefix = Mutation.delete("test", KeySet.prefixRange(Key.of(1L)));
Mutation range = Mutation.delete("test", KeySet.range(KeyRange.openOpen(Key.of(1L), Key.newBuilder().build())));
MutationGroup[] mutationGroups = new MutationGroup[] { g(m(1L)), g(m(2L), m(3L)), // not batchable - too many rows.
g(m(1L), m(3L), m(4L), m(5L)), g(del(1L)), // not point delete.
g(del(5L, 6L)), g(all), g(prefix), g(range) };
BatchableMutationFilterFn testFn = new BatchableMutationFilterFn(null, null, 1000, 1000, 3);
BatchableMutationFilterFn.ProcessContext mockProcessContext = Mockito.mock(ProcessContext.class);
when(mockProcessContext.sideInput(any())).thenReturn(getSchema());
// Capture the outputs.
doNothing().when(mockProcessContext).output(mutationGroupCaptor.capture());
doNothing().when(mockProcessContext).output(any(), mutationGroupListCaptor.capture());
// Process all elements.
for (MutationGroup m : mutationGroups) {
when(mockProcessContext.element()).thenReturn(m);
testFn.processElement(mockProcessContext);
}
// Verify captured batchable elements.
assertThat(mutationGroupCaptor.getAllValues(), containsInAnyOrder(g(m(1L)), g(m(2L), m(3L)), g(del(1L))));
// Verify captured unbatchable mutations
Iterable<MutationGroup> unbatchableMutations = Iterables.concat(mutationGroupListCaptor.getAllValues());
assertThat(unbatchableMutations, containsInAnyOrder(// not batchable - too many rows.
g(m(1L), m(3L), m(4L), m(5L)), // not point delete.
g(del(5L, 6L)), g(all), g(prefix), g(range)));
}
use of com.google.cloud.spanner.Mutation in project beam by apache.
the class SpannerIOWriteTest method retryOnAbortedAndDeadlineExceeded.
@Test
public void retryOnAbortedAndDeadlineExceeded() throws InterruptedException {
List<Mutation> mutationList = Arrays.asList(m((long) 1));
String errString = "Transaction aborted. " + "Database schema probably changed during transaction, retry may succeed.";
// mock sleeper so that it does not actually sleep.
WriteToSpannerFn.sleeper = Mockito.mock(Sleeper.class);
// Respond with (1) Aborted transaction a couple of times (2) deadline exceeded
// (3) Aborted transaction 3 times (4) deadline exceeded and finally return success.
when(serviceFactory.mockDatabaseClient().writeAtLeastOnceWithOptions(any(), any(ReadQueryUpdateTransactionOption.class))).thenThrow(SpannerExceptionFactory.newSpannerException(ErrorCode.ABORTED, errString)).thenThrow(SpannerExceptionFactory.newSpannerException(ErrorCode.ABORTED, errString)).thenThrow(SpannerExceptionFactory.newSpannerException(ErrorCode.DEADLINE_EXCEEDED, "simulated Timeout 1")).thenThrow(SpannerExceptionFactory.newSpannerException(ErrorCode.ABORTED, errString)).thenThrow(SpannerExceptionFactory.newSpannerException(ErrorCode.ABORTED, errString)).thenThrow(SpannerExceptionFactory.newSpannerException(ErrorCode.ABORTED, errString)).thenThrow(SpannerExceptionFactory.newSpannerException(ErrorCode.DEADLINE_EXCEEDED, "simulated Timeout 2")).thenReturn(new CommitResponse(Timestamp.now()));
SpannerWriteResult result = pipeline.apply(Create.of(mutationList)).apply(SpannerIO.write().withProjectId("test-project").withInstanceId("test-instance").withDatabaseId("test-database").withServiceFactory(serviceFactory).withBatchSizeBytes(0).withFailureMode(FailureMode.FAIL_FAST));
// Zero error
PAssert.that(result.getFailedMutations()).satisfies(m -> {
assertEquals(0, Iterables.size(m));
return null;
});
pipeline.run().waitUntilFinish();
// 2 calls to sleeper
verify(WriteToSpannerFn.sleeper, times(2)).sleep(anyLong());
// 8 write attempts for the single mutationGroup.
verify(serviceFactory.mockDatabaseClient(), times(8)).writeAtLeastOnceWithOptions(any(), any(ReadQueryUpdateTransactionOption.class));
}
use of com.google.cloud.spanner.Mutation in project beam by apache.
the class MutationKeyEncoderTest method verifyEncodedOrdering.
private void verifyEncodedOrdering(SpannerSchema schema, List<Mutation> expectedMutations) {
MutationKeyEncoder encoder = new MutationKeyEncoder(schema);
Assert.assertEquals(5, expectedMutations.size());
// mix them up.
List<Mutation> unsortedMutations = Arrays.asList(expectedMutations.get(3), expectedMutations.get(4), expectedMutations.get(1), expectedMutations.get(2), expectedMutations.get(0));
// Use a map to sort the list by encoded table/key, then by Mutation contents to give a defined
// order when the same key is given, or if it is an unknown table.
TreeMultimap<byte[], Mutation> mutationsByEncoding = TreeMultimap.create(UnsignedBytes.lexicographicalComparator(), Comparator.comparing(Mutation::toString));
for (Mutation m : unsortedMutations) {
mutationsByEncoding.put(encoder.encodeTableNameAndKey(m), m);
}
Assert.assertEquals(expectedMutations, new ArrayList<>(mutationsByEncoding.values()));
}
Aggregations