com.google.cloud.spanner.TimestampBound Java Examples

void partitionReadUsingIndex() {
  // [START partition_read_using_index]
  final BatchReadOnlyTransaction txn =
      batchClient.batchReadOnlyTransaction(TimestampBound.strong());
  List<Partition> partitions =
      txn.partitionReadUsingIndex(
          PartitionOptions.getDefaultInstance(),
          "Singers",
          "SingerId",
          KeySet.all(),
          Arrays.asList("SingerId", "FirstName", "LastName"));

  for (Partition p : partitions) {
    try (ResultSet results = txn.execute(p)) {
      while (results.next()) {
        long singerId = results.getLong(0);
        String firstName = results.getString(1);
        String lastName = results.getString(2);
        System.out.println("[" + singerId + "] " + firstName + " " + lastName);
      }
    }
  }
  // [END partition_read_using_index]
}
 

@Test
public void readOnlyTransactionTest() {

	ReadOnlyTransaction readOnlyTransaction = mock(ReadOnlyTransaction.class);
	when(this.databaseClient.readOnlyTransaction(
			eq(TimestampBound.ofMinReadTimestamp(Timestamp.ofTimeMicroseconds(333)))))
					.thenReturn(readOnlyTransaction);

	String finalResult = this.spannerTemplate
			.performReadOnlyTransaction((spannerOperations) -> {
				List<TestEntity> items = spannerOperations.readAll(TestEntity.class);
				TestEntity item = spannerOperations.read(TestEntity.class,
						Key.of("key"));
				return "all done";
			}, new SpannerReadOptions()
					.setTimestampBound(TimestampBound.ofMinReadTimestamp(Timestamp.ofTimeMicroseconds(333L))));

	assertThat(finalResult).isEqualTo("all done");
	verify(readOnlyTransaction, times(2)).read(eq("custom_test_table"), any(), any());
}
 

public void begin() throws SQLException {
  if (connection.isBatchReadOnly()) {
    if (batchReadOnlyTransaction == null) {
      batchReadOnlyTransaction = batchClient.batchReadOnlyTransaction(TimestampBound.strong());
    }
  } else if (connection.isReadOnly()) {
    if (readOnlyTransaction == null) {
      readOnlyTransaction = dbClient.readOnlyTransaction();
    }
  } else {
    if (transactionThread == null) {
      transactionThread = new TransactionThread(dbClient, connection.getLogger());
      transactionThread.start();
    }
  }
}
 

@Test
public void readOnlyTransactionDmlTest() {

	this.expectedException.expectMessage("A read-only transaction template cannot execute DML.");

	ReadOnlyTransaction readOnlyTransaction = mock(ReadOnlyTransaction.class);
	when(this.databaseClient.readOnlyTransaction(
			// exact staleness is expected.
			eq(TimestampBound.ofReadTimestamp(Timestamp.ofTimeMicroseconds(333)))))
					.thenReturn(readOnlyTransaction);

	this.spannerTemplate
			.performReadOnlyTransaction((spannerOperations) -> {
				spannerOperations.executeDmlStatement(Statement.of("fail"));
				return null;
			}, new SpannerReadOptions()
					.setTimestamp(Timestamp.ofTimeMicroseconds(333)));
}
 

@Test
public void readOnlyTransactionPartitionedDmlTest() {

	this.expectedException.expectMessage("A read-only transaction template cannot execute partitioned" +
			" DML.");

	ReadOnlyTransaction readOnlyTransaction = mock(ReadOnlyTransaction.class);
	when(this.databaseClient.readOnlyTransaction(
			eq(TimestampBound.ofReadTimestamp(Timestamp.ofTimeMicroseconds(333)))))
					.thenReturn(readOnlyTransaction);

	this.spannerTemplate
			.performReadOnlyTransaction((spannerOperations) -> {
				spannerOperations.executePartitionedDmlStatement(Statement.of("fail"));
				return null;
			}, new SpannerReadOptions()
					.setTimestamp(Timestamp.ofTimeMicroseconds(333)));
}
 

@Test
public void findMultipleKeysTest() {
	ResultSet results = mock(ResultSet.class);
	ReadOption readOption = mock(ReadOption.class);
	SpannerReadOptions options = new SpannerReadOptions().addReadOption(readOption)
			.setTimestampBound(TimestampBound.ofMinReadTimestamp(Timestamp.ofTimeMicroseconds(333L)));
	KeySet keySet = KeySet.singleKey(Key.of("key"));
	when(this.readContext.read(any(), any(), any(), any())).thenReturn(results);
	when(this.databaseClient.singleUse(eq(TimestampBound.ofMinReadTimestamp(Timestamp.ofTimeMicroseconds(333L)))))
			.thenReturn(this.readContext);

	verifyAfterEvents(new AfterReadEvent(Collections.emptyList(), keySet, options),
			() -> this.spannerTemplate.read(TestEntity.class, keySet, options), x -> {
				verify(this.objectMapper, times(1)).mapToList(same(results),
						eq(TestEntity.class), isNull(), eq(false));
				verify(this.readContext, times(1)).read(eq("custom_test_table"), same(keySet),
						any(), same(readOption));
			});
	verify(this.databaseClient, times(1))
			.singleUse(TimestampBound.ofMinReadTimestamp(Timestamp.ofTimeMicroseconds(333L)));
}
 

@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();
}
 

@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();
}
 

/** Example of read only transaction with timestamp bound. */
// [TARGET readOnlyTransaction(TimestampBound)]
// [VARIABLE my_singer_id]
// [VARIABLE my_album_id]
public String readOnlyTransactionTimestamp(long singerId, long albumId) {
  // [START readOnlyTransactionTimestamp]
  String singerColumn = "FirstName";
  String albumColumn = "AlbumTitle";
  String albumTitle = null;
  // ReadOnlyTransaction should be closed to prevent resource leak.
  try (ReadOnlyTransaction txn =
      dbClient.readOnlyTransaction(TimestampBound.ofExactStaleness(10, TimeUnit.SECONDS))) {
    Struct singerRow =
        txn.readRow("Singers", Key.of(singerId), Collections.singleton(singerColumn));
    Struct albumRow =
        txn.readRow("Albums", Key.of(singerId, albumId), Collections.singleton(albumColumn));
    singerRow.getString(singerColumn);
    albumTitle = albumRow.getString(albumColumn);
  }
  // [END readOnlyTransactionTimestamp]
  return albumTitle;
}
 

void partitionRead() {
  // [START partition_read]
  final BatchReadOnlyTransaction txn =
      batchClient.batchReadOnlyTransaction(TimestampBound.strong());
  List<Partition> partitions =
      txn.partitionRead(
          PartitionOptions.getDefaultInstance(),
          "Singers",
          KeySet.all(),
          Arrays.asList("SingerId", "FirstName", "LastName"));
  for (final Partition p : partitions) {
    try (ResultSet results = txn.execute(p)) {
      while (results.next()) {
        long singerId = results.getLong(0);
        String firstName = results.getString(1);
        String lastName = results.getString(2);
        System.out.println("[" + singerId + "] " + firstName + " " + lastName);
      }
    }
  }
  // [END partition_read]
}
 

void partitionQuery() {
  // [START partition_query]
  final BatchReadOnlyTransaction txn =
      batchClient.batchReadOnlyTransaction(TimestampBound.strong());
  List<Partition> partitions =
      txn.partitionQuery(
          PartitionOptions.getDefaultInstance(),
          Statement.of("SELECT SingerId, FirstName, LastName FROM Singers"));

  for (final Partition p : partitions) {
    try (ResultSet results = txn.execute(p)) {
      while (results.next()) {
        long singerId = results.getLong(0);
        String firstName = results.getString(1);
        String lastName = results.getString(2);
        System.out.println("[" + singerId + "] " + firstName + " " + lastName);
      }
    }
  }
  // [END partition_query]

}
 

@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();
}
 

/**
 * Returns a transform that creates a batch transaction. By default, {@link
 * TimestampBound#strong()} transaction is created, to override this use {@link
 * CreateTransaction#withTimestampBound(TimestampBound)}.
 */
@Experimental
public static CreateTransaction createTransaction() {
  return new AutoValue_LocalSpannerIO_CreateTransaction.Builder()
      .setSpannerConfig(SpannerConfig.create())
      .setTimestampBound(TimestampBound.strong())
      .build();
}
 

@Test
public void testExecuteBatchReadOnly() throws SQLException, NoSuchFieldException,
    SecurityException, IllegalArgumentException, IllegalAccessException {
  for (int testRun = 0; testRun < 2; testRun++) {
    final int numberOfPartitions = 6;
    BatchClient batchClient = mock(BatchClient.class);
    BatchReadOnlyTransaction tx = mock(BatchReadOnlyTransaction.class);
    List<Partition> partitions = new ArrayList<>(numberOfPartitions);
    for (int i = 0; i < numberOfPartitions; i++)
      partitions.add(mock(Partition.class));
    when(tx.partitionQuery(any(), any())).then(new Returns(partitions));
    when(batchClient.batchReadOnlyTransaction(TimestampBound.strong())).then(new Returns(tx));
    Field field = CloudSpannerTransaction.class.getDeclaredField("batchClient");
    field.setAccessible(true);
    field.set(connection.getTransaction(), batchClient);
    connection.setBatchReadOnly(true);
    Statement statement;
    if (testRun % 2 == 0) {
      statement = connection.createStatement();
      assertTrue(statement.execute(SELECT_ALL_FROM_FOO));
    } else {
      PreparedStatement ps = connection.prepareStatement(SELECT_ALL_FROM_FOO);
      assertTrue(ps.execute());
      statement = ps;
    }
    List<ResultSet> resultSets = new ArrayList<>();
    do {
      resultSets.add(statement.getResultSet());
    } while (statement.getMoreResults());
    assertEquals(numberOfPartitions, resultSets.size());
  }
}
 

private Ddl getDatabaseDdl() {
  BatchClient batchClient = spannerServer.getBatchClient(dbId);
  BatchReadOnlyTransaction batchTx =
      batchClient.batchReadOnlyTransaction(TimestampBound.strong());
  InformationSchemaScanner scanner = new InformationSchemaScanner(batchTx);
  return scanner.scan();
}
 

@ProcessElement
public void processElement(ProcessContext c) throws Exception {
  String timestamp = ExportTransform.this.snapshotTime.get();

  TimestampBound tsb;
  if ("".equals(timestamp)) {
    /* If no timestamp is specified, read latest data */
    tsb = TimestampBound.strong();
  } else {
    /* Else try to read data in the timestamp specified. */
    com.google.cloud.Timestamp tsVal;
    try {
      tsVal = com.google.cloud.Timestamp.parseTimestamp(timestamp);
    } catch (Exception e) {
      throw new IllegalStateException("Invalid timestamp specified " + timestamp);
    }

    /*
     * If timestamp specified is in the future, spanner read will wait
     * till the time has passed. Abort the job and complain early.
     */
    if (tsVal.compareTo(com.google.cloud.Timestamp.now()) > 0) {
      throw new IllegalStateException("Timestamp specified is in future " + timestamp);
    }

    /*
     * Export jobs with Timestamps which are older than 
     * maximum staleness time (one hour) fail with the FAILED_PRECONDITION
     * error - https://cloud.google.com/spanner/docs/timestamp-bounds
     * Hence we do not handle the case.
     */

    tsb = TimestampBound.ofReadTimestamp(tsVal);
  }
  BatchReadOnlyTransaction tx =
    spannerAccessor.getBatchClient().batchReadOnlyTransaction(tsb);
  c.output(Transaction.create(tx.getBatchTransactionId()));
}
 

/**
 * Creates an uninitialized instance of {@link Read}. Before use, the {@link Read} must be
 * configured with a {@link Read#withInstanceId} and {@link Read#withDatabaseId} that identify the
 * Cloud Spanner database.
 */
public static Read read() {
  return new AutoValue_LocalSpannerIO_Read.Builder()
      .setSpannerConfig(SpannerConfig.create())
      .setTimestampBound(TimestampBound.strong())
      .setReadOperation(ReadOperation.create())
      .setBatching(true)
      .build();
}
 

/**
 * A {@link PTransform} that works like {@link #read}, but executes read operations coming from a
 * {@link PCollection}.
 */
public static ReadAll readAll() {
  return new AutoValue_LocalSpannerIO_ReadAll.Builder()
      .setSpannerConfig(SpannerConfig.create())
      .setTimestampBound(TimestampBound.strong())
      .setBatching(true)
      .build();
}
 

/** Example to do a batch strong read. */
BatchReadOnlyTransaction readStrong() {
  // [START batch_client_strong_read]
  BatchReadOnlyTransaction txn = batchClient.batchReadOnlyTransaction(TimestampBound.strong());
  // [END batch_client_strong_read]
  return txn;
}
 

/**
 * Creates an uninitialized instance of {@link Read}. Before use, the {@link Read} must be
 * configured with a {@link Read#withInstanceId} and {@link Read#withDatabaseId} that identify the
 * Cloud Spanner database.
 */
public static Read read() {
  return new AutoValue_SpannerIO_Read.Builder()
      .setSpannerConfig(SpannerConfig.create())
      .setTimestampBound(TimestampBound.strong())
      .setReadOperation(ReadOperation.create())
      .setBatching(true)
      .build();
}
 

/**
 * A {@link PTransform} that works like {@link #read}, but executes read operations coming from a
 * {@link PCollection}.
 */
public static ReadAll readAll() {
  return new AutoValue_SpannerIO_ReadAll.Builder()
      .setSpannerConfig(SpannerConfig.create())
      .setTimestampBound(TimestampBound.strong())
      .setBatching(true)
      .build();
}
 

/**
 * Returns a transform that creates a batch transaction. By default, {@link
 * TimestampBound#strong()} transaction is created, to override this use {@link
 * CreateTransaction#withTimestampBound(TimestampBound)}.
 */
@Experimental
public static CreateTransaction createTransaction() {
  return new AutoValue_SpannerIO_CreateTransaction.Builder()
      .setSpannerConfig(SpannerConfig.create())
      .setTimestampBound(TimestampBound.strong())
      .build();
}
 

/** Example of single use read only transaction with timestamp bound. */
// [TARGET singleUseReadOnlyTransaction(TimestampBound)]
// [VARIABLE my_singer_id]
public Timestamp singleUseReadOnlyTransactionTimestamp(long singerId) {
  // [START singleUseReadOnlyTransactionTimestamp]
  String column = "FirstName";
  ReadOnlyTransaction txn =
      dbClient.singleUseReadOnlyTransaction(TimestampBound.ofMaxStaleness(10, TimeUnit.SECONDS));
  Struct row = txn.readRow("Singers", Key.of(singerId), Collections.singleton(column));
  row.getString(column);
  Timestamp timestamp = txn.getReadTimestamp();
  // [END singleUseReadOnlyTransactionTimestamp]
  return timestamp;
}
 

/** Example of single use with timestamp bound. */
// [TARGET singleUse(TimestampBound)]
// [VARIABLE my_singer_id]
public String singleUseStale(long singerId) {
  // [START singleUseStale]
  String column = "FirstName";
  Struct row =
      dbClient
          .singleUse(TimestampBound.ofMaxStaleness(10, TimeUnit.SECONDS))
          .readRow("Singers", Key.of(singerId), Collections.singleton(column));
  String firstName = row.getString(column);
  // [END singleUseStale]
  return firstName;
}
 

static void readStaleData(DatabaseClient dbClient) {
  try (ResultSet resultSet =
      dbClient
          .singleUse(TimestampBound.ofExactStaleness(15, TimeUnit.SECONDS))
          .read(
              "Albums", KeySet.all(), Arrays.asList("SingerId", "AlbumId", "MarketingBudget"))) {
    while (resultSet.next()) {
      System.out.printf(
          "%d %d %s\n",
          resultSet.getLong(0),
          resultSet.getLong(1),
          resultSet.isNull(2) ? "NULL" : resultSet.getLong("MarketingBudget"));
    }
  }
}
 

@Test
public void runQuery() throws Exception {
  Timestamp timestamp = Timestamp.ofTimeMicroseconds(12345);
  TimestampBound timestampBound = TimestampBound.ofReadTimestamp(timestamp);

  SpannerConfig spannerConfig =
      SpannerConfig.create()
          .withProjectId("test")
          .withInstanceId("123")
          .withDatabaseId("aaa")
          .withServiceFactory(serviceFactory);

  PCollection<Struct> one =
      pipeline.apply(
          "read q",
          SpannerIO.read()
              .withSpannerConfig(spannerConfig)
              .withQuery("SELECT * FROM users")
              .withTimestampBound(timestampBound));

  FakeBatchTransactionId id = new FakeBatchTransactionId("runQueryTest");
  when(mockBatchTx.getBatchTransactionId()).thenReturn(id);

  when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(timestampBound))
      .thenReturn(mockBatchTx);
  when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(any(BatchTransactionId.class)))
      .thenReturn(mockBatchTx);

  Partition fakePartition =
      FakePartitionFactory.createFakeQueryPartition(ByteString.copyFromUtf8("one"));

  when(mockBatchTx.partitionQuery(
          any(PartitionOptions.class), eq(Statement.of("SELECT * FROM users"))))
      .thenReturn(Arrays.asList(fakePartition, 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, 6)));

  PAssert.that(one).containsInAnyOrder(FAKE_ROWS);

  pipeline.run();
}
 

public CreateTransaction withTimestampBound(TimestampBound timestampBound) {
  return toBuilder().setTimestampBound(timestampBound).build();
}
 

@Nullable
abstract TimestampBound getTimestampBound();
 

@Test
public void runRead() throws Exception {
  Timestamp timestamp = Timestamp.ofTimeMicroseconds(12345);
  TimestampBound timestampBound = TimestampBound.ofReadTimestamp(timestamp);

  SpannerConfig spannerConfig =
      SpannerConfig.create()
          .withProjectId("test")
          .withInstanceId("123")
          .withDatabaseId("aaa")
          .withServiceFactory(serviceFactory);

  PCollection<Struct> one =
      pipeline.apply(
          "read q",
          SpannerIO.read()
              .withSpannerConfig(spannerConfig)
              .withTable("users")
              .withColumns("id", "name")
              .withTimestampBound(timestampBound));

  FakeBatchTransactionId id = new FakeBatchTransactionId("runReadTest");
  when(mockBatchTx.getBatchTransactionId()).thenReturn(id);

  when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(timestampBound))
      .thenReturn(mockBatchTx);
  when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(any(BatchTransactionId.class)))
      .thenReturn(mockBatchTx);

  Partition fakePartition =
      FakePartitionFactory.createFakeReadPartition(ByteString.copyFromUtf8("one"));

  when(mockBatchTx.partitionRead(
          any(PartitionOptions.class),
          eq("users"),
          eq(KeySet.all()),
          eq(Arrays.asList("id", "name"))))
      .thenReturn(Arrays.asList(fakePartition, fakePartition, 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();
}
 

@Test
public void runReadUsingIndex() throws Exception {
  Timestamp timestamp = Timestamp.ofTimeMicroseconds(12345);
  TimestampBound timestampBound = TimestampBound.ofReadTimestamp(timestamp);

  SpannerConfig spannerConfig =
      SpannerConfig.create()
          .withProjectId("test")
          .withInstanceId("123")
          .withDatabaseId("aaa")
          .withServiceFactory(serviceFactory);

  PCollection<Struct> one =
      pipeline.apply(
          "read q",
          SpannerIO.read()
              .withTimestamp(Timestamp.now())
              .withSpannerConfig(spannerConfig)
              .withTable("users")
              .withColumns("id", "name")
              .withIndex("theindex")
              .withTimestampBound(timestampBound));

  FakeBatchTransactionId id = new FakeBatchTransactionId("runReadUsingIndexTest");
  when(mockBatchTx.getBatchTransactionId()).thenReturn(id);

  when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(timestampBound))
      .thenReturn(mockBatchTx);
  when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(any(BatchTransactionId.class)))
      .thenReturn(mockBatchTx);

  Partition fakePartition =
      FakePartitionFactory.createFakeReadPartition(ByteString.copyFromUtf8("one"));

  when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(id)).thenReturn(mockBatchTx);
  when(mockBatchTx.partitionReadUsingIndex(
          any(PartitionOptions.class),
          eq("users"),
          eq("theindex"),
          eq(KeySet.all()),
          eq(Arrays.asList("id", "name"))))
      .thenReturn(Arrays.asList(fakePartition, fakePartition, 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();
}