com.google.cloud.spanner.Partition 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());
}
 

@Override
public PCollection<Struct> expand(PCollection<ReadOperation> input) {
  PCollectionView<Transaction> txView = getTxView();
  if (txView == null) {
    Pipeline begin = input.getPipeline();
    SpannerIO.CreateTransaction createTx =
        SpannerIO.createTransaction()
            .withSpannerConfig(getSpannerConfig())
            .withTimestampBound(getTimestampBound());
    txView = begin.apply(createTx);
  }
  return input
      .apply(
          "Generate Partitions",
          ParDo.of(new GeneratePartitionsFn(getSpannerConfig(), txView)).withSideInputs(txView))
      .apply("Shuffle partitions", Reshuffle.<Partition>viaRandomKey())
      .apply(
          "Read from Partitions",
          ParDo.of(new ReadFromPartitionFn(getSpannerConfig(), txView)).withSideInputs(txView));
}
 

@Override
public PCollection<Struct> expand(PCollection<ReadOperation> input) {
  PCollectionView<Transaction> txView = getTxView();
  if (txView == null) {
    Pipeline begin = input.getPipeline();
    LocalSpannerIO.CreateTransaction createTx =
        LocalSpannerIO.createTransaction()
            .withSpannerConfig(getSpannerConfig())
            .withTimestampBound(getTimestampBound());
    txView = begin.apply(createTx);
  }
  return input
      .apply(
          "Generate Partitions",
          ParDo.of(new GeneratePartitionsFn(getSpannerConfig(), txView)).withSideInputs(txView))
      .apply("Shuffle partitions", Reshuffle.<Partition>viaRandomKey())
      .apply(
          "Read from Partitions",
          ParDo.of(new ReadFromPartitionFn(getSpannerConfig(), txView)).withSideInputs(txView));
}
 

public static CloudSpannerConnection createConnection() throws SQLException {
  CloudSpannerConnection connection = Mockito.mock(CloudSpannerConnection.class);
  Mockito.doCallRealMethod().when(connection).setAutoCommit(Mockito.anyBoolean());
  Mockito.when(connection.getAutoCommit()).thenCallRealMethod();
  connection.setAutoCommit(false);
  Mockito.doCallRealMethod().when(connection).setBatchReadOnly(Mockito.anyBoolean());
  Mockito.when(connection.isBatchReadOnly()).thenCallRealMethod();
  Mockito.doCallRealMethod().when(connection).setReadOnly(Mockito.anyBoolean());
  Mockito.when(connection.isReadOnly()).thenCallRealMethod();

  Mockito.when(connection.isAllowExtendedMode()).thenAnswer(new Returns(true));
  Mockito.when(connection.createArrayOf(Mockito.anyString(), Mockito.any())).thenCallRealMethod();
  CloudSpannerDatabaseMetaData metadata = createMetaData();
  Mockito.when(connection.getMetaData()).thenReturn(metadata);
  CloudSpannerTransaction transaction = Mockito.mock(CloudSpannerTransaction.class);
  Mockito.when(transaction.executeQuery(Mockito.any()))
      .thenReturn(Mockito.mock(com.google.cloud.spanner.ResultSet.class));
  Mockito.when(transaction.partitionQuery(Mockito.any(), Mockito.any())).thenReturn(
      Arrays.asList(mock(Partition.class), mock(Partition.class), mock(Partition.class)));
  Mockito.when(connection.getTransaction()).thenReturn(transaction);

  TableKeyMetaData tableFoo = Mockito.mock(TableKeyMetaData.class);
  Mockito.when(tableFoo.getKeyColumns()).thenAnswer(new Returns(Arrays.asList("ID")));
  Mockito
      .when(connection.getTable(Mockito
          .matches(Pattern.compile("FOO", Pattern.CASE_INSENSITIVE | Pattern.UNICODE_CASE))))
      .thenAnswer(new Returns(tableFoo));

  TableKeyMetaData tableBar = Mockito.mock(TableKeyMetaData.class);
  Mockito.when(tableBar.getKeyColumns()).thenAnswer(new Returns(Arrays.asList("ID1", "ID2")));
  Mockito
      .when(connection.getTable(Mockito
          .matches(Pattern.compile("BAR", Pattern.CASE_INSENSITIVE | Pattern.UNICODE_CASE))))
      .thenAnswer(new Returns(tableBar));
  Mockito.when(connection.getLogger()).thenAnswer(new Returns(new Logger()));

  mockXAMethods(connection);

  return connection;
}
 

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

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

@Override
public ResultSet execute(Partition partition) throws SpannerException {
  checkTransaction();
  if (batchReadOnlyTransaction != null) {
    return batchReadOnlyTransaction.execute(partition);
  }
  throw SpannerExceptionFactory.newSpannerException(ErrorCode.FAILED_PRECONDITION,
      METHOD_ONLY_IN_BATCH_READONLY);
}
 

@Override
public List<Partition> partitionQuery(PartitionOptions partitionOptions, Statement statement,
    QueryOption... options) throws SpannerException {
  checkTransaction();
  if (batchReadOnlyTransaction != null) {
    return batchReadOnlyTransaction.partitionQuery(partitionOptions, statement, options);
  }
  throw SpannerExceptionFactory.newSpannerException(ErrorCode.FAILED_PRECONDITION,
      METHOD_ONLY_IN_BATCH_READONLY);
}
 

@Override
public List<Partition> partitionReadUsingIndex(PartitionOptions partitionOptions, String table,
    String index, KeySet keys, Iterable<String> columns, ReadOption... options)
    throws SpannerException {
  throw SpannerExceptionFactory.newSpannerException(ErrorCode.UNIMPLEMENTED,
      METHOD_NOT_IMPLEMENTED);
}
 

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

@Override
public boolean execute() throws SQLException {
  CustomDriverStatement custom = getCustomDriverStatement(sqlTokens);
  if (custom != null)
    return custom.execute(sqlTokens);
  Statement statement = null;
  boolean ddl = isDDLStatement();
  if (!ddl) {
    try {
      statement = CCJSqlParserUtil.parse(sanitizeSQL(sql));
    } catch (JSQLParserException | TokenMgrException e) {
      throw new CloudSpannerSQLException(PARSE_ERROR + sql + ": " + e.getLocalizedMessage(),
          Code.INVALID_ARGUMENT, e);
    }
  }
  if (!ddl && statement instanceof Select) {
    determineForceSingleUseReadContext((Select) statement);
    com.google.cloud.spanner.Statement.Builder builder = createSelectBuilder(statement, sql);
    if (!isForceSingleUseReadContext() && getConnection().isBatchReadOnly()) {
      List<Partition> partitions = partitionQuery(builder.build());
      currentResultSets = new ArrayList<>(partitions.size());
      for (Partition p : partitions) {
        currentResultSets
            .add(new CloudSpannerPartitionResultSet(this, getBatchReadOnlyTransaction(), p, sql));
      }
    } else {
      try (ReadContext context = getReadContext()) {
        com.google.cloud.spanner.ResultSet rs = context.executeQuery(builder.build());
        currentResultSets = Arrays.asList(new CloudSpannerResultSet(this, rs, sql));
        currentResultSetIndex = 0;
        lastUpdateCount = -1;
      }
    }
    return true;
  } else {
    lastUpdateCount = executeUpdate();
    currentResultSets = null;
    currentResultSetIndex = 0;
    return false;
  }
}
 

@Override
public boolean execute(String sql) throws SQLException {
  String[] sqlTokens = getTokens(sql);
  CustomDriverStatement custom = getCustomDriverStatement(sqlTokens);
  if (custom != null)
    return custom.execute(sqlTokens);
  Statement statement = null;
  boolean ddl = isDDLStatement(sqlTokens);
  if (!ddl) {
    try {
      statement = CCJSqlParserUtil.parse(sanitizeSQL(sql));
    } catch (JSQLParserException | TokenMgrException e) {
      throw new CloudSpannerSQLException(
          "Error while parsing sql statement " + sql + ": " + e.getLocalizedMessage(),
          Code.INVALID_ARGUMENT, e);
    }
  }
  if (!ddl && statement instanceof Select) {
    determineForceSingleUseReadContext((Select) statement);
    if (!isForceSingleUseReadContext() && getConnection().isBatchReadOnly()) {
      List<Partition> partitions = partitionQuery(com.google.cloud.spanner.Statement.of(sql));
      currentResultSets = new ArrayList<>(partitions.size());
      for (Partition p : partitions) {
        currentResultSets
            .add(new CloudSpannerPartitionResultSet(this, getBatchReadOnlyTransaction(), p, sql));
      }
    } else {
      try (ReadContext context = getReadContext()) {
        com.google.cloud.spanner.ResultSet rs =
            context.executeQuery(com.google.cloud.spanner.Statement.of(sql));
        currentResultSets = Arrays.asList(new CloudSpannerResultSet(this, rs, sql));
        currentResultSetIndex = 0;
        lastUpdateCount = -1;
      }
    }
    return true;
  } else {
    lastUpdateCount = executeUpdate(sql);
    currentResultSetIndex = 0;
    currentResultSets = null;
    return false;
  }
}
 

protected List<Partition> partitionQuery(com.google.cloud.spanner.Statement statement) {
  PartitionOptions po = PartitionOptions.getDefaultInstance();
  return connection.getTransaction().partitionQuery(po, statement);
}
 

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

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

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

@Test
public void readPipeline() 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 txId = new FakeBatchTransactionId("readPipelineTest");
  when(mockBatchTx.getBatchTransactionId()).thenReturn(txId);

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

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

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

  PCollectionView<Transaction> tx =
      pipeline.apply(
          "tx",
          SpannerIO.createTransaction()
              .withSpannerConfig(spannerConfig)
              .withTimestampBound(timestampBound));

  PCollection<ReadOperation> reads =
      pipeline.apply(
          Create.of(
              ReadOperation.create().withQuery("SELECT * FROM users"),
              ReadOperation.create().withTable("users").withColumns("id", "name")));

  PCollection<Struct> one =
      reads.apply(
          "read all", SpannerIO.readAll().withSpannerConfig(spannerConfig).withTransaction(tx));

  BatchTransactionId txId = new FakeBatchTransactionId("tx");
  when(mockBatchTx.getBatchTransactionId()).thenReturn(txId);

  when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(timestampBound))
      .thenReturn(mockBatchTx);

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

  Partition fakePartition =
      FakePartitionFactory.createFakeReadPartition(ByteString.copyFromUtf8("partition"));
  when(mockBatchTx.partitionQuery(
          any(PartitionOptions.class), eq(Statement.of("SELECT * FROM users"))))
      .thenReturn(Arrays.asList(fakePartition, fakePartition));
  when(mockBatchTx.partitionRead(
          any(PartitionOptions.class),
          eq("users"),
          eq(KeySet.all()),
          eq(Arrays.asList("id", "name"))))
      .thenReturn(Arrays.asList(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();
}
 

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

@Override
public List<Partition> partitionRead(PartitionOptions partitionOptions, String table, KeySet keys,
    Iterable<String> columns, ReadOption... options) throws SpannerException {
  throw SpannerExceptionFactory.newSpannerException(ErrorCode.UNIMPLEMENTED,
      METHOD_NOT_IMPLEMENTED);
}