.thenReturn(fakePartitions); List<Partition> result = fnTester.processBundle(read.getReadOperation()); assertThat(result, Matchers.containsInAnyOrder(fakePartitions.toArray()));
@Test public void runQuery() throws Exception { SpannerIO.Read read = SpannerIO.read() .withProjectId("test") .withInstanceId("123") .withDatabaseId("aaa") .withQuery("SELECT * FROM users") .withServiceFactory(serviceFactory); List<Partition> fakePartitions = Arrays.asList(mock(Partition.class), mock(Partition.class), mock(Partition.class)); BatchTransactionId id = mock(BatchTransactionId.class); Transaction tx = Transaction.create(id); PCollectionView<Transaction> txView = pipeline.apply(Create.of(tx)).apply(View.<Transaction>asSingleton()); BatchSpannerRead.GeneratePartitionsFn fn = new BatchSpannerRead.GeneratePartitionsFn(read.getSpannerConfig(), txView); DoFnTester<ReadOperation, Partition> fnTester = DoFnTester.of(fn); fnTester.setSideInput(txView, GlobalWindow.INSTANCE, tx); when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(id)).thenReturn(mockBatchTx); when(mockBatchTx.partitionQuery(any(PartitionOptions.class), any(Statement.class))) .thenReturn(fakePartitions); List<Partition> result = fnTester.processBundle(read.getReadOperation()); assertThat(result, Matchers.containsInAnyOrder(fakePartitions.toArray())); verify(serviceFactory.mockBatchClient()).batchReadOnlyTransaction(id); verify(mockBatchTx) .partitionQuery(any(PartitionOptions.class), eq(Statement.of("SELECT * " + "FROM users"))); }
.thenReturn(fakePartitions); List<Partition> result = fnTester.processBundle(read.getReadOperation()); assertThat(result, Matchers.containsInAnyOrder(fakePartitions.toArray()));
+ "with withTimestampBound or withTimestamp method"); if (getReadOperation().getQuery() != null) { } else if (getReadOperation().getTable() != null) { getReadOperation().getColumns(), "For a read operation SpannerIO.read() requires a list of " + "columns to set with withColumns method"); checkArgument( !getReadOperation().getColumns().isEmpty(), "For a read operation SpannerIO.read() requires a" + " list of columns to set with withColumns method"); .withBatching(getBatching()) .withTransaction(getTransaction()); return input.apply(Create.of(getReadOperation())).apply("Execute query", readAll);
public Read withKeySet(KeySet keySet) { return withReadOperation(getReadOperation().withKeySet(keySet)); }
public Read withTable(String table) { return withReadOperation(getReadOperation().withTable(table)); }
public Read withPartitionOptions(PartitionOptions partitionOptions) { return withReadOperation(getReadOperation().withPartitionOptions(partitionOptions)); }
public Read withQuery(Statement statement) { return withReadOperation(getReadOperation().withQuery(statement)); }
public Read withIndex(String index) { return withReadOperation(getReadOperation().withIndex(index)); }
public Read withColumns(List<String> columns) { return withReadOperation(getReadOperation().withColumns(columns)); }