Search in sources :

Example 21 with Mutation

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)));
}
Also used : DoFn(org.apache.beam.sdk.transforms.DoFn) Mutation(com.google.cloud.spanner.Mutation) BatchableMutationFilterFn(org.apache.beam.sdk.io.gcp.spanner.SpannerIO.BatchableMutationFilterFn) Test(org.junit.Test)

Example 22 with Mutation

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)));
}
Also used : Write(org.apache.beam.sdk.io.gcp.spanner.SpannerIO.Write) Mutation(com.google.cloud.spanner.Mutation) Test(org.junit.Test)

Example 23 with Mutation

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)));
}
Also used : Mutation(com.google.cloud.spanner.Mutation) BatchableMutationFilterFn(org.apache.beam.sdk.io.gcp.spanner.SpannerIO.BatchableMutationFilterFn) Test(org.junit.Test)

Example 24 with Mutation

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));
}
Also used : ReadQueryUpdateTransactionOption(com.google.cloud.spanner.Options.ReadQueryUpdateTransactionOption) CommitResponse(com.google.cloud.spanner.CommitResponse) Mutation(com.google.cloud.spanner.Mutation) Sleeper(org.apache.beam.sdk.util.Sleeper) Test(org.junit.Test)

Example 25 with Mutation

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()));
}
Also used : Mutation(com.google.cloud.spanner.Mutation)

Aggregations

Mutation (com.google.cloud.spanner.Mutation)70 Test (org.junit.Test)53 ArrayList (java.util.ArrayList)11 CommitResponse (com.google.cloud.spanner.CommitResponse)4 KeySet (com.google.cloud.spanner.KeySet)4 ReadQueryUpdateTransactionOption (com.google.cloud.spanner.Options.ReadQueryUpdateTransactionOption)4 Timestamp (com.google.cloud.Timestamp)3 DatabaseClient (com.google.cloud.spanner.DatabaseClient)3 ResultSet (com.google.cloud.spanner.ResultSet)3 TransactionContext (com.google.cloud.spanner.TransactionContext)3 BatchableMutationFilterFn (org.apache.beam.sdk.io.gcp.spanner.SpannerIO.BatchableMutationFilterFn)3 Sleeper (org.apache.beam.sdk.util.Sleeper)3 Database (com.google.cloud.spanner.Database)2 DatabaseAdminClient (com.google.cloud.spanner.DatabaseAdminClient)2 Key (com.google.cloud.spanner.Key)2 WriteBuilder (com.google.cloud.spanner.Mutation.WriteBuilder)2 SpannerException (com.google.cloud.spanner.SpannerException)2 TransactionRunner (com.google.cloud.spanner.TransactionRunner)2 CreateDatabaseMetadata (com.google.spanner.admin.database.v1.CreateDatabaseMetadata)2 BigDecimal (java.math.BigDecimal)2