@ProcessElement public void processElement(ProcessContext c) throws Exception { BatchReadOnlyTransaction tx = spannerAccessor.getBatchClient().batchReadOnlyTransaction(config.getTimestampBound()); c.output(Transaction.create(tx.getBatchTransactionId())); } }
@ProcessElement public void processElement(ProcessContext c) throws Exception { Transaction tx = c.sideInput(txView); BatchReadOnlyTransaction context = spannerAccessor.getBatchClient().batchReadOnlyTransaction(tx.transactionId()); for (Partition p : execute(c.element(), context)) { c.output(p); } }
@ProcessElement public void processElement(ProcessContext c) { Transaction transaction = c.sideInput(tx); BatchTransactionId transactionId = transaction.transactionId(); BatchClient batchClient = spannerAccessor.getBatchClient(); BatchReadOnlyTransaction context = batchClient.batchReadOnlyTransaction(transactionId); InformationSchemaScanner scanner = new InformationSchemaScanner(context); Ddl ddl = scanner.scan(); c.output(ddl); } }
Transaction tx = Transaction.create(id); PCollectionView<Transaction> txView = pipeline.apply(Create.of(tx)).apply(View.<Transaction>asSingleton());
@ProcessElement public void processElement(ProcessContext c) throws Exception { Transaction tx = c.sideInput(txView); ReadOperation op = c.element(); BatchReadOnlyTransaction context = spannerAccessor.getBatchClient().batchReadOnlyTransaction(tx.transactionId()); try (ResultSet resultSet = execute(op, context)) { while (resultSet.next()) { c.output(resultSet.getCurrentRowAsStruct()); } } }
@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"))); }
@ProcessElement public void processElement(ProcessContext c) throws Exception { Transaction tx = c.sideInput(txView); BatchReadOnlyTransaction batchTx = spannerAccessor.getBatchClient().batchReadOnlyTransaction(tx.transactionId()); Partition p = c.element(); try (ResultSet resultSet = batchTx.execute(p)) { while (resultSet.next()) { Struct s = resultSet.getCurrentRowAsStruct(); c.output(s); } } } }
Transaction tx = Transaction.create(id); PCollectionView<Transaction> txView = pipeline.apply(Create.of(tx)).apply(View.<Transaction>asSingleton());