Examples with BatchTransactionId com.google.cloud.spanner.BatchTransactionId used on opensource projects

Example 1 with BatchTransactionId

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

the class SpannerIOReadTest method readAllPipeline.

@Test
public void readAllPipeline() throws Exception {
    Timestamp timestamp = Timestamp.ofTimeMicroseconds(12345);
    TimestampBound timestampBound = TimestampBound.ofReadTimestamp(timestamp);
    SpannerConfig spannerConfig = getSpannerConfig();
    PCollectionView<Transaction> tx = pipeline.apply("tx", SpannerIO.createTransaction().withSpannerConfig(spannerConfig).withTimestampBound(timestampBound));
    PCollection<ReadOperation> reads = pipeline.apply(Create.of(ReadOperation.create().withQuery("SELECT * FROM users"), ReadOperation.create().withTable("users").withColumns("id", "name")));
    PCollection<Struct> one = reads.apply("read all", SpannerIO.readAll().withSpannerConfig(spannerConfig).withTransaction(tx));
    BatchTransactionId txId = new FakeBatchTransactionId("tx");
    when(mockBatchTx.getBatchTransactionId()).thenReturn(txId);
    when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(timestampBound)).thenReturn(mockBatchTx);
    when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(any(BatchTransactionId.class))).thenReturn(mockBatchTx);
    Partition fakePartition = FakePartitionFactory.createFakeReadPartition(ByteString.copyFromUtf8("partition"));
    when(mockBatchTx.partitionQuery(any(PartitionOptions.class), eq(Statement.of("SELECT * FROM users")), any(ReadQueryUpdateTransactionOption.class))).thenReturn(Arrays.asList(fakePartition, fakePartition));
    when(mockBatchTx.partitionRead(any(PartitionOptions.class), eq("users"), eq(KeySet.all()), eq(Arrays.asList("id", "name")), any(ReadQueryUpdateTransactionOption.class))).thenReturn(Arrays.asList(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();
}