@VisibleForTesting Read withServiceFactory(ServiceFactory<Spanner, SpannerOptions> serviceFactory) { SpannerConfig config = getSpannerConfig(); return withSpannerConfig(config.withServiceFactory(serviceFactory)); }
new BatchSpannerRead.GeneratePartitionsFn(read.getSpannerConfig(), txView); DoFnTester<ReadOperation, Partition> fnTester = DoFnTester.of(fn); fnTester.setSideInput(txView, GlobalWindow.INSTANCE, tx);
@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"))); }
new BatchSpannerRead.GeneratePartitionsFn(read.getSpannerConfig(), txView); DoFnTester<ReadOperation, Partition> fnTester = DoFnTester.of(fn); fnTester.setSideInput(txView, GlobalWindow.INSTANCE, tx);
@Override public PCollection<Struct> expand(PBegin input) { getSpannerConfig().validate(); checkArgument( getTimestampBound() != null, .withSpannerConfig(getSpannerConfig()) .withTimestampBound(getTimestampBound()) .withBatching(getBatching())
/** Specifies the Cloud Spanner instance. */ public Read withInstanceId(ValueProvider<String> instanceId) { SpannerConfig config = getSpannerConfig(); return withSpannerConfig(config.withInstanceId(instanceId)); }
/** Specifies the Cloud Spanner host. */ public Read withHost(ValueProvider<String> host) { SpannerConfig config = getSpannerConfig(); return withSpannerConfig(config.withHost(host)); }
/** Specifies the Cloud Spanner database. */ public Read withDatabaseId(ValueProvider<String> databaseId) { SpannerConfig config = getSpannerConfig(); return withSpannerConfig(config.withDatabaseId(databaseId)); }
/** Specifies the Cloud Spanner project. */ public Read withProjectId(ValueProvider<String> projectId) { SpannerConfig config = getSpannerConfig(); return withSpannerConfig(config.withProjectId(projectId)); }