com.google.cloud.spanner.BatchReadOnlyTransaction 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]
}
 

private List<Partition> execute(ReadOperation op, BatchReadOnlyTransaction tx) {
  // Query was selected.
  if (op.getQuery() != null) {
    return tx.partitionQuery(op.getPartitionOptions(), op.getQuery());
  }
  // Read with index was selected.
  if (op.getIndex() != null) {
    return tx.partitionReadUsingIndex(
        op.getPartitionOptions(),
        op.getTable(),
        op.getIndex(),
        op.getKeySet(),
        op.getColumns());
  }
  // Read from table was selected.
  return tx.partitionRead(
      op.getPartitionOptions(), op.getTable(), op.getKeySet(), op.getColumns());
}
 

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]

}
 

private List<Partition> execute(ReadOperation op, BatchReadOnlyTransaction tx) {
  // Query was selected.
  if (op.getQuery() != null) {
    return tx.partitionQuery(op.getPartitionOptions(), op.getQuery());
  }
  // Read with index was selected.
  if (op.getIndex() != null) {
    return tx.partitionReadUsingIndex(
        op.getPartitionOptions(),
        op.getTable(),
        op.getIndex(),
        op.getKeySet(),
        op.getColumns());
  }
  // Read from table was selected.
  return tx.partitionRead(
      op.getPartitionOptions(), op.getTable(), op.getKeySet(), op.getColumns());
}
 

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

/** 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;
}
 

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

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

private ResultSet execute(ReadOperation op, BatchReadOnlyTransaction readOnlyTransaction) {
  if (op.getQuery() != null) {
    return readOnlyTransaction.executeQuery(op.getQuery());
  }
  if (op.getIndex() != null) {
    return readOnlyTransaction.readUsingIndex(
        op.getTable(), op.getIndex(), op.getKeySet(), op.getColumns());
  }
  return readOnlyTransaction.read(op.getTable(), op.getKeySet(), op.getColumns());
}
 

@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 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 CloudSpannerPartitionResultSet createSubject() {
  Partition partition = mock(Partition.class);
  BatchReadOnlyTransaction transaction = mock(BatchReadOnlyTransaction.class);
  ResultSet rs = CloudSpannerResultSetTest.getMockResultSet();
  when(transaction.execute(partition)).thenReturn(rs);
  return new CloudSpannerPartitionResultSet(mock(CloudSpannerStatement.class), transaction,
      partition, "SELECT * FROM FOO");
}
 

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) {
  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);
}
 

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

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

@ProcessElement
public void processElement(ProcessContext c) throws Exception {
  BatchReadOnlyTransaction tx =
      spannerAccessor.getBatchClient().batchReadOnlyTransaction(config.getTimestampBound());
  c.output(Transaction.create(tx.getBatchTransactionId()));
}
 

protected BatchReadOnlyTransaction getBatchReadOnlyTransaction() {
  return connection.getTransaction().getBatchReadOnlyTransaction();
}
 

public CloudSpannerPartitionResultSet(CloudSpannerStatement statement,
    BatchReadOnlyTransaction transaction, Partition partition, String sql) {
  super(statement, sql);
  this.transaction = transaction;
  this.partition = partition;
}
 

public BatchReadOnlyTransaction getBatchReadOnlyTransaction() {
  return batchReadOnlyTransaction;
}
 

@Before
@SuppressWarnings("unchecked")
public void setUp() throws Exception {
  serviceFactory = new FakeServiceFactory();
  mockBatchTx = Mockito.mock(BatchReadOnlyTransaction.class);
}
 

/**
 * This example showcases how to create a batch client, partition a query, and concurrently read
 * from multiple partitions.
 */
public static void main(String[] args) throws InterruptedException {
  if (args.length != 2) {
    System.err.println("Usage: BatchSample <instance_id> <database_id>");
    return;
  }

  /*
   * CREATE TABLE Singers (
   *   SingerId   INT64 NOT NULL,
   *   FirstName  STRING(1024),
   *   LastName   STRING(1024),
   *   SingerInfo BYTES(MAX),
   * ) PRIMARY KEY (SingerId);
   */

  String instanceId = args[0];
  String databaseId = args[1];

  SpannerOptions options = SpannerOptions.newBuilder().build();
  Spanner spanner = options.getService();

  // [START spanner_batch_client]
  int numThreads = Runtime.getRuntime().availableProcessors();
  ExecutorService executor = Executors.newFixedThreadPool(numThreads);

  // Statistics
  int totalPartitions;
  AtomicInteger totalRecords = new AtomicInteger(0);

  try {
    BatchClient batchClient =
        spanner.getBatchClient(DatabaseId.of(options.getProjectId(), instanceId, databaseId));

    final BatchReadOnlyTransaction txn =
        batchClient.batchReadOnlyTransaction(TimestampBound.strong());

    // A Partition object is serializable and can be used from a different process.
    List<Partition> partitions =
        txn.partitionQuery(
            PartitionOptions.getDefaultInstance(),
            Statement.of("SELECT SingerId, FirstName, LastName FROM Singers"));

    totalPartitions = partitions.size();

    for (final Partition p : partitions) {
      executor.execute(
          () -> {
            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);
                totalRecords.getAndIncrement();
              }
            }
          });
    }
  } finally {
    executor.shutdown();
    executor.awaitTermination(1, TimeUnit.HOURS);
    spanner.close();
  }

  double avgRecordsPerPartition = 0.0;
  if (totalPartitions != 0) {
    avgRecordsPerPartition = (double) totalRecords.get() / totalPartitions;
  }
  System.out.println("totalPartitions=" + totalPartitions);
  System.out.println("totalRecords=" + totalRecords);
  System.out.println("avgRecordsPerPartition=" + avgRecordsPerPartition);
  // [END spanner_batch_client]
}
 

@ProcessElement
public void processElement(ProcessContext c) throws Exception {
  BatchReadOnlyTransaction tx =
      spannerAccessor.getBatchClient().batchReadOnlyTransaction(config.getTimestampBound());
  c.output(Transaction.create(tx.getBatchTransactionId()));
}