com.google.cloud.spanner.BatchClient Java Examples

private ExposedSpannerAccessor(
    Spanner spanner,
    DatabaseClient databaseClient,
    DatabaseAdminClient databaseAdminClient,
    BatchClient batchClient) {
  this.spanner = spanner;
  this.databaseClient = databaseClient;
  this.databaseAdminClient = databaseAdminClient;
  this.batchClient = batchClient;
}
 

BatchClient getBatchClient() {
  // [START get_batch_client]
  SpannerOptions options = SpannerOptions.newBuilder().build();
  Spanner spanner = options.getService();
  final String project = "test-project";
  final String instance = "test-instance";
  final String database = "example-db";
  DatabaseId db = DatabaseId.of(project, instance, database);
  BatchClient batchClient = spanner.getBatchClient(db);
  // [END get_batch_client]

  return batchClient;
}
 

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

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

@VisibleForTesting
CloudSpannerConnection(DatabaseClient dbClient, BatchClient batchClient) {
  this.driver = null;
  this.database = null;
  this.url = null;
  this.suppliedProperties = null;
  int logLevel = CloudSpannerDriver.getLogLevel();
  synchronized (CloudSpannerConnection.class) {
    logger = new Logger(nextConnectionID++);
    logger.setLogLevel(logLevel);
  }
  this.dbClient = dbClient;
  this.transaction = new CloudSpannerTransaction(dbClient, batchClient, this);
  this.metaDataStore = new MetaDataStore(this);
}
 

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

@Test
public void testSupportsSavePoints() throws SQLException {
  assertTrue(testSubject.supportsSavepoints());
  // Check that the connection does not throw an Exception
  @SuppressWarnings("resource")
  CloudSpannerConnection connection =
      new CloudSpannerConnection(mock(DatabaseClient.class), mock(BatchClient.class));
  connection.setAutoCommit(false);
  connection.setSavepoint();
}
 

public FakeServiceFactory() {
  synchronized (lock) {
    index = count++;
    mockSpanners.add(mock(Spanner.class, withSettings().serializable()));
    mockDatabaseClients.add(mock(DatabaseClient.class, withSettings().serializable()));
    mockBatchClients.add(mock(BatchClient.class, withSettings().serializable()));
  }
  when(mockSpanner().getDatabaseClient(Matchers.any(DatabaseId.class)))
      .thenReturn(mockDatabaseClient());
  when(mockSpanner().getBatchClient(Matchers.any(DatabaseId.class)))
      .thenReturn(mockBatchClient());
}
 

private SpannerAccessor(
    Spanner spanner,
    DatabaseClient databaseClient,
    DatabaseAdminClient databaseAdminClient,
    BatchClient batchClient,
    SpannerConfig spannerConfig) {
  this.spanner = spanner;
  this.databaseClient = databaseClient;
  this.databaseAdminClient = databaseAdminClient;
  this.batchClient = batchClient;
  this.spannerConfig = spannerConfig;
}
 

public BatchClientSnippets(BatchClient batchClient) {
  this.batchClient = batchClient;
}
 

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

BatchClient mockBatchClient() {
  synchronized (lock) {
    return mockBatchClients.get(index);
  }
}
 

BatchClient getBatchClient() {
  return batchClient;
}
 

private static SpannerAccessor createAndConnect(SpannerConfig spannerConfig) {
  SpannerOptions.Builder builder = SpannerOptions.newBuilder();

  ValueProvider<Duration> commitDeadline = spannerConfig.getCommitDeadline();
  if (commitDeadline != null && commitDeadline.get().getMillis() > 0) {

    // Set the GRPC deadline on the Commit API call.
    UnaryCallSettings.Builder commitSettings =
        builder.getSpannerStubSettingsBuilder().commitSettings();
    RetrySettings.Builder commitRetrySettings = commitSettings.getRetrySettings().toBuilder();
    commitSettings.setRetrySettings(
        commitRetrySettings
            .setTotalTimeout(org.threeten.bp.Duration.ofMillis(commitDeadline.get().getMillis()))
            .setMaxRpcTimeout(org.threeten.bp.Duration.ofMillis(commitDeadline.get().getMillis()))
            .setInitialRpcTimeout(
                org.threeten.bp.Duration.ofMillis(commitDeadline.get().getMillis()))
            .build());
  }

  ValueProvider<String> projectId = spannerConfig.getProjectId();
  if (projectId != null) {
    builder.setProjectId(projectId.get());
  }
  ServiceFactory<Spanner, SpannerOptions> serviceFactory = spannerConfig.getServiceFactory();
  if (serviceFactory != null) {
    builder.setServiceFactory(serviceFactory);
  }
  ValueProvider<String> host = spannerConfig.getHost();
  if (host != null) {
    builder.setHost(host.get());
  }
  String userAgentString = USER_AGENT_PREFIX + "/" + ReleaseInfo.getReleaseInfo().getVersion();
  builder.setHeaderProvider(FixedHeaderProvider.create("user-agent", userAgentString));
  SpannerOptions options = builder.build();

  Spanner spanner = options.getService();
  String instanceId = spannerConfig.getInstanceId().get();
  String databaseId = spannerConfig.getDatabaseId().get();
  DatabaseClient databaseClient =
      spanner.getDatabaseClient(DatabaseId.of(options.getProjectId(), instanceId, databaseId));
  BatchClient batchClient =
      spanner.getBatchClient(DatabaseId.of(options.getProjectId(), instanceId, databaseId));
  DatabaseAdminClient databaseAdminClient = spanner.getDatabaseAdminClient();

  return new SpannerAccessor(
      spanner, databaseClient, databaseAdminClient, batchClient, spannerConfig);
}
 

@Before
public void setup() throws SQLException, NoSuchFieldException, SecurityException,
    IllegalArgumentException, IllegalAccessException {
  connection = new CloudSpannerConnection(mock(DatabaseClient.class), mock(BatchClient.class));
  connection.setAutoCommit(false);
}
 

public CloudSpannerTransaction(DatabaseClient dbClient, BatchClient batchClient,
    CloudSpannerConnection connection) {
  this.dbClient = dbClient;
  this.batchClient = batchClient;
  this.connection = connection;
}
 

public BatchClient getBatchClient(String dbName) {
  return client.getBatchClient(DatabaseId.of(projectId, instanceId, dbName));
}
 

public BatchClient getBatchClient() {
  return batchClient;
}
 

public static ExposedSpannerAccessor create(SpannerConfig spannerConfig) {
  SpannerOptions.Builder builder = SpannerOptions.newBuilder();

  ValueProvider<Duration> commitDeadline = spannerConfig.getCommitDeadline();
  if (commitDeadline != null && commitDeadline.get().getMillis() > 0) {

    // In Spanner API version 1.21 or above, we can set the deadline / total Timeout on an API
    // call using the following code:
    //
    // UnaryCallSettings.Builder commitSettings =
    // builder.getSpannerStubSettingsBuilder().commitSettings();
    // RetrySettings.Builder commitRetrySettings = commitSettings.getRetrySettings().toBuilder()
    // commitSettings.setRetrySettings(
    //     commitRetrySettings.setTotalTimeout(
    //         Duration.ofMillis(getCommitDeadlineMillis().get()))
    //     .build());
    //
    // However, at time of this commit, the Spanner API is at only at v1.6.0, where the only
    // method to set a deadline is with GRPC Interceptors, so we have to use that...
    SpannerInterceptorProvider interceptorProvider =
        SpannerInterceptorProvider.createDefault()
            .with(new CommitDeadlineSettingInterceptor(commitDeadline.get()));
    builder.setInterceptorProvider(interceptorProvider);
  }

  ValueProvider<String> projectId = spannerConfig.getProjectId();
  if (projectId != null) {
    builder.setProjectId(projectId.get());
  }
  ServiceFactory<Spanner, SpannerOptions> serviceFactory = spannerConfig.getServiceFactory();
  if (serviceFactory != null) {
    builder.setServiceFactory(serviceFactory);
  }
  ValueProvider<String> host = spannerConfig.getHost();
  if (host != null) {
    builder.setHost(host.get());
  }
  String userAgentString = USER_AGENT_PREFIX + "/" + ReleaseInfo.getReleaseInfo().getVersion();
  builder.setHeaderProvider(FixedHeaderProvider.create("user-agent", userAgentString));

  SessionPoolOptions.Builder sessionPoolOptions = SessionPoolOptions.newBuilder();
  sessionPoolOptions.setMinSessions(1);
  sessionPoolOptions.setMaxSessions(3);
  builder.setSessionPoolOption(sessionPoolOptions.build());

  SpannerOptions options = builder.build();

  Spanner spanner = options.getService();
  String instanceId = spannerConfig.getInstanceId().get();
  String databaseId = spannerConfig.getDatabaseId().get();
  DatabaseClient databaseClient =
      spanner.getDatabaseClient(DatabaseId.of(options.getProjectId(), instanceId, databaseId));
  BatchClient batchClient =
      spanner.getBatchClient(DatabaseId.of(options.getProjectId(), instanceId, databaseId));
  DatabaseAdminClient databaseAdminClient = spanner.getDatabaseAdminClient();

  return new ExposedSpannerAccessor(spanner, databaseClient, databaseAdminClient, batchClient);
}