Search in sources :

Example 11 with FakeBatchTransactionId

use of com.google.cloud.spanner.FakeBatchTransactionId in project beam by apache.

the class SpannerIOReadTest method runReadUsingIndex.

@Test
public void runReadUsingIndex() throws Exception {
    Timestamp timestamp = Timestamp.ofTimeMicroseconds(12345);
    TimestampBound timestampBound = TimestampBound.ofReadTimestamp(timestamp);
    SpannerConfig spannerConfig = getSpannerConfig();
    PCollection<Struct> one = pipeline.apply("read q", SpannerIO.read().withTimestamp(Timestamp.now()).withSpannerConfig(spannerConfig).withTable("users").withColumns("id", "name").withIndex("theindex").withTimestampBound(timestampBound));
    FakeBatchTransactionId id = new FakeBatchTransactionId("runReadUsingIndexTest");
    when(mockBatchTx.getBatchTransactionId()).thenReturn(id);
    when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(timestampBound)).thenReturn(mockBatchTx);
    when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(any(BatchTransactionId.class))).thenReturn(mockBatchTx);
    Partition fakePartition = FakePartitionFactory.createFakeReadPartition(ByteString.copyFromUtf8("one"));
    when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(id)).thenReturn(mockBatchTx);
    when(mockBatchTx.partitionReadUsingIndex(any(PartitionOptions.class), eq("users"), eq("theindex"), eq(KeySet.all()), eq(Arrays.asList("id", "name")), any(ReadQueryUpdateTransactionOption.class))).thenReturn(Arrays.asList(fakePartition, fakePartition, fakePartition));
    when(mockBatchTx.execute(any(Partition.class))).thenReturn(ResultSets.forRows(FAKE_TYPE, FAKE_ROWS.subList(0, 2)), ResultSets.forRows(FAKE_TYPE, FAKE_ROWS.subList(2, 4)), ResultSets.forRows(FAKE_TYPE, FAKE_ROWS.subList(4, 6)));
    PAssert.that(one).containsInAnyOrder(FAKE_ROWS);
    pipeline.run();
}
Also used : Partition(com.google.cloud.spanner.Partition) PartitionOptions(com.google.cloud.spanner.PartitionOptions) FakeBatchTransactionId(com.google.cloud.spanner.FakeBatchTransactionId) BatchTransactionId(com.google.cloud.spanner.BatchTransactionId) ReadQueryUpdateTransactionOption(com.google.cloud.spanner.Options.ReadQueryUpdateTransactionOption) TimestampBound(com.google.cloud.spanner.TimestampBound) FakeBatchTransactionId(com.google.cloud.spanner.FakeBatchTransactionId) Timestamp(com.google.cloud.Timestamp) Struct(com.google.cloud.spanner.Struct) Test(org.junit.Test)

Aggregations

Timestamp (com.google.cloud.Timestamp)11 BatchTransactionId (com.google.cloud.spanner.BatchTransactionId)11 FakeBatchTransactionId (com.google.cloud.spanner.FakeBatchTransactionId)11 ReadQueryUpdateTransactionOption (com.google.cloud.spanner.Options.ReadQueryUpdateTransactionOption)11 Partition (com.google.cloud.spanner.Partition)11 PartitionOptions (com.google.cloud.spanner.PartitionOptions)11 TimestampBound (com.google.cloud.spanner.TimestampBound)11 Test (org.junit.Test)10 Struct (com.google.cloud.spanner.Struct)7 SpannerException (com.google.cloud.spanner.SpannerException)2 PipelineExecutionException (org.apache.beam.sdk.Pipeline.PipelineExecutionException)2 Read (org.apache.beam.sdk.io.gcp.spanner.SpannerIO.Read)2 BatchReadOnlyTransaction (com.google.cloud.spanner.BatchReadOnlyTransaction)1