SpannerIO.Read read = SpannerIO.read() .withProjectId("test") .withInstanceId("123") .withDatabaseId("aaa") .withTable("users") .withColumns("id", "name") .withServiceFactory(serviceFactory); new BatchSpannerRead.GeneratePartitionsFn(read.getSpannerConfig(), txView); DoFnTester<ReadOperation, Partition> fnTester = DoFnTester.of(fn); fnTester.setSideInput(txView, GlobalWindow.INSTANCE, tx); .thenReturn(fakePartitions); List<Partition> result = fnTester.processBundle(read.getReadOperation()); assertThat(result, Matchers.containsInAnyOrder(fakePartitions.toArray()));
"read q", SpannerIO.read() .withSpannerConfig(spannerConfig) .withQuery("SELECT * FROM users") .withTimestampBound(timestampBound)); FakeBatchTransactionId txId = new FakeBatchTransactionId("readPipelineTest"); when(mockBatchTx.getBatchTransactionId()).thenReturn(txId);
SpannerIO.Read read = SpannerIO.read() .withProjectId("test") .withInstanceId("123") .withDatabaseId("aaa") .withTimestamp(Timestamp.now()) .withTable("users") .withColumns("id", "name") .withIndex("theindex") .withServiceFactory(serviceFactory); new BatchSpannerRead.GeneratePartitionsFn(read.getSpannerConfig(), txView); DoFnTester<ReadOperation, Partition> fnTester = DoFnTester.of(fn); fnTester.setSideInput(txView, GlobalWindow.INSTANCE, tx); .thenReturn(fakePartitions); List<Partition> result = fnTester.processBundle(read.getReadOperation()); assertThat(result, Matchers.containsInAnyOrder(fakePartitions.toArray()));
@Override public PCollection<Struct> expand(PBegin input) { getSpannerConfig().validate(); checkArgument( getTimestampBound() != null, "SpannerIO.read() runs in a read only transaction and requires timestamp to be set " + "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"); .withSpannerConfig(getSpannerConfig()) .withTimestampBound(getTimestampBound()) .withBatching(getBatching()) .withTransaction(getTransaction()); return input.apply(Create.of(getReadOperation())).apply("Execute query", readAll);
@Test public void testReadAllRecordsInDb() throws Exception { SpannerConfig spannerConfig = createSpannerConfig(); PCollectionView<Transaction> tx = p.apply( SpannerIO.createTransaction() .withSpannerConfig(spannerConfig) .withTimestampBound(TimestampBound.strong())); PCollection<Struct> allRecords = p.apply( SpannerIO.read() .withSpannerConfig(spannerConfig) .withBatching(false) .withQuery( "SELECT t.table_name FROM information_schema.tables AS t WHERE t" + ".table_catalog = '' AND t.table_schema = ''")) .apply( MapElements.into(TypeDescriptor.of(ReadOperation.class)) .via( (SerializableFunction<Struct, ReadOperation>) input -> { String tableName = input.getString(0); return ReadOperation.create().withQuery("SELECT * FROM " + tableName); })) .apply(SpannerIO.readAll().withTransaction(tx).withSpannerConfig(spannerConfig)); PAssert.thatSingleton(allRecords.apply("Count rows", Count.globally())).isEqualTo(5L); p.run(); }
@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"))); }
@Test public void testQuery() throws Exception { SpannerConfig spannerConfig = createSpannerConfig(); PCollectionView<Transaction> tx = p.apply( SpannerIO.createTransaction() .withSpannerConfig(spannerConfig) .withTimestampBound(TimestampBound.strong())); PCollection<Struct> output = p.apply( SpannerIO.read() .withSpannerConfig(spannerConfig) .withQuery("SELECT * FROM " + options.getTable()) .withTransaction(tx)); PAssert.thatSingleton(output.apply("Count rows", Count.globally())).isEqualTo(5L); p.run(); }
@Test public void testRead() throws Exception { SpannerConfig spannerConfig = createSpannerConfig(); PCollectionView<Transaction> tx = p.apply( SpannerIO.createTransaction() .withSpannerConfig(spannerConfig) .withTimestampBound(TimestampBound.strong())); PCollection<Struct> output = p.apply( SpannerIO.read() .withSpannerConfig(spannerConfig) .withTable(options.getTable()) .withColumns("Key", "Value") .withTransaction(tx)); PAssert.thatSingleton(output.apply("Count rows", Count.<Struct>globally())).isEqualTo(5L); p.run(); }
public Read withQuery(String sql) { return withQuery(Statement.of(sql)); }
public Read withPartitionOptions(PartitionOptions partitionOptions) { return withReadOperation(getReadOperation().withPartitionOptions(partitionOptions)); }
public Read withKeySet(KeySet keySet) { return withReadOperation(getReadOperation().withKeySet(keySet)); }
public Read withTable(String table) { return withReadOperation(getReadOperation().withTable(table)); }
public Read withTimestamp(Timestamp timestamp) { return withTimestampBound(TimestampBound.ofReadTimestamp(timestamp)); }
public Read withQuery(Statement statement) { return withReadOperation(getReadOperation().withQuery(statement)); }
public Read withColumns(String... columns) { return withColumns(Arrays.asList(columns)); }
public Read withIndex(String index) { return withReadOperation(getReadOperation().withIndex(index)); }
public Read withColumns(List<String> columns) { return withReadOperation(getReadOperation().withColumns(columns)); }
/** Specifies the Cloud Spanner database. */ public Read withDatabaseId(String databaseId) { return withDatabaseId(ValueProvider.StaticValueProvider.of(databaseId)); }
/** Specifies the Cloud Spanner instance. */ public Read withInstanceId(String instanceId) { return withInstanceId(ValueProvider.StaticValueProvider.of(instanceId)); }
/** Specifies the Cloud Spanner project. */ public Read withProjectId(String projectId) { return withProjectId(ValueProvider.StaticValueProvider.of(projectId)); }