com.google.datastore.v1.CommitRequest Java Examples

private void datastoreWriterFnTest(int numMutations) throws Exception {
  // Create the requested number of mutations.
  List<Mutation> mutations = new ArrayList<>(numMutations);
  for (int i = 0; i < numMutations; ++i) {
    mutations.add(
        makeUpsert(Entity.newBuilder().setKey(makeKey("key" + i, i + 1)).build()).build());
  }

  DatastoreWriterFn datastoreWriter =
      new DatastoreWriterFn(
          StaticValueProvider.of(PROJECT_ID), null, mockDatastoreFactory, new FakeWriteBatcher());
  DoFnTester<Mutation, Void> doFnTester = DoFnTester.of(datastoreWriter);
  doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
  doFnTester.processBundle(mutations);

  int start = 0;
  while (start < numMutations) {
    int end = Math.min(numMutations, start + DatastoreV1.DATASTORE_BATCH_UPDATE_ENTITIES_START);
    CommitRequest.Builder commitRequest = CommitRequest.newBuilder();
    commitRequest.setMode(CommitRequest.Mode.NON_TRANSACTIONAL);
    commitRequest.addAllMutations(mutations.subList(start, end));
    // Verify all the batch requests were made with the expected mutations.
    verify(mockDatastore, times(1)).commit(commitRequest.build());
    start = end;
  }
}
 

/** Tests {@link DatastoreWriterFn} with a failed request which is retried. */
@Test
public void testDatatoreWriterFnRetriesErrors() throws Exception {
  List<Mutation> mutations = new ArrayList<>();
  int numRpcs = 2;
  for (int i = 0; i < DatastoreV1.DATASTORE_BATCH_UPDATE_ENTITIES_START * numRpcs; ++i) {
    mutations.add(
        makeUpsert(Entity.newBuilder().setKey(makeKey("key" + i, i + 1)).build()).build());
  }

  CommitResponse successfulCommit = CommitResponse.getDefaultInstance();
  when(mockDatastore.commit(any(CommitRequest.class)))
      .thenReturn(successfulCommit)
      .thenThrow(new DatastoreException("commit", Code.DEADLINE_EXCEEDED, "", null))
      .thenReturn(successfulCommit);

  DatastoreWriterFn datastoreWriter =
      new DatastoreWriterFn(
          StaticValueProvider.of(PROJECT_ID), null, mockDatastoreFactory, new FakeWriteBatcher());
  DoFnTester<Mutation, Void> doFnTester = DoFnTester.of(datastoreWriter);
  doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
  doFnTester.processBundle(mutations);
}
 

private ListenableFuture<MutationResult> executeAsyncMutations(final List<Mutation> mutations, final ListenableFuture<TransactionResult> txn) {
  final ListenableFuture<Response> httpResponse = Futures.transformAsync(txn, result -> {
    final CommitRequest.Builder request = CommitRequest.newBuilder();
    if (mutations != null) {
      request.addAllMutations(mutations);
    }

    final ByteString transaction = result.getTransaction();
    if (transaction != null) {
      request.setTransaction(transaction);
    } else {
      request.setMode(CommitRequest.Mode.NON_TRANSACTIONAL);
    }
    final ProtoHttpContent payload = new ProtoHttpContent(request.build());
    return ListenableFutureAdapter.asGuavaFuture(prepareRequest("commit", payload).execute());
  }, MoreExecutors.directExecutor());
  return Futures.transformAsync(httpResponse, response -> {
    if (!isSuccessful(response.getStatusCode())) {
      throw new DatastoreException(response.getStatusCode(), response.getResponseBody());
    }
    final CommitResponse commit = CommitResponse.parseFrom(streamResponse(response));
    return Futures.immediateFuture(MutationResult.build(commit));
  }, MoreExecutors.directExecutor());
}
 

private void flushBatch() throws DatastoreException, IOException, InterruptedException {
  LOG.info("Writing batch of {} entities", entities.size());
  Sleeper sleeper = Sleeper.DEFAULT;
  BackOff backoff =
      FluentBackoff.DEFAULT
          .withMaxRetries(MAX_RETRIES)
          .withInitialBackoff(INITIAL_BACKOFF)
          .backoff();

  while (true) {
    // Batch mutate entities.
    try {
      CommitRequest.Builder commitRequest = CommitRequest.newBuilder();
      for (Entity entity : entities) {
        commitRequest.addMutations(mutationBuilder.apply(entity));
      }
      commitRequest.setMode(CommitRequest.Mode.NON_TRANSACTIONAL);
      datastore.commit(commitRequest.build());
      // Break if the commit threw no exception.
      break;
    } catch (DatastoreException exception) {
      LOG.error(
          "Error writing to the Datastore ({}): {}",
          exception.getCode(),
          exception.getMessage());
      if (!BackOffUtils.next(sleeper, backoff)) {
        LOG.error("Aborting after {} retries.", MAX_RETRIES);
        throw exception;
      }
    }
  }
  LOG.info("Successfully wrote {} entities", entities.size());
  entities.clear();
}
 

/**
  * Insert an entity into the datastore.
  *
  * The entity must have no ids.
  *
  * @return The key for the inserted entity.
  * @throws DatastoreException on error
  */
 private Key insert(Entity entity) throws DatastoreException {
   CommitRequest request = CommitRequest.newBuilder()
.addMutations(Mutation.newBuilder()
    .setInsert(entity))
       .setMode(CommitRequest.Mode.NON_TRANSACTIONAL)
.build();
   return datastore.commit(request).getMutationResults(0).getKey();
 }
 

public CommitResponse commit(CommitRequest request) throws DatastoreException {
  try (InputStream is = remoteRpc.call("commit", request)) {
    return CommitResponse.parseFrom(is);
  } catch (IOException exception) {
    throw invalidResponseException("commit", exception);
  }
}
 

@Test
public void commit() throws Exception {
  CommitRequest.Builder request = CommitRequest.newBuilder();
  request.setTransaction(ByteString.copyFromUtf8("project-id"));
  CommitResponse.Builder response = CommitResponse.newBuilder();
  expectRpc("commit", request.build(), response.build());
}
 

/**
 * Writes a batch of mutations to Cloud Datastore.
 *
 * <p>If a commit fails, it will be retried up to {@link #MAX_RETRIES} times. All mutations in
 * the batch will be committed again, even if the commit was partially successful. If the retry
 * limit is exceeded, the last exception from Cloud Datastore will be thrown.
 *
 * @throws DatastoreException if the commit fails or IOException or InterruptedException if
 *     backing off between retries fails.
 */
private void flushBatch() throws DatastoreException, IOException, InterruptedException {
  LOG.debug("Writing batch of {} mutations", mutations.size());
  Sleeper sleeper = Sleeper.DEFAULT;
  BackOff backoff = BUNDLE_WRITE_BACKOFF.backoff();

  while (true) {
    // Batch upsert entities.
    CommitRequest.Builder commitRequest = CommitRequest.newBuilder();
    commitRequest.addAllMutations(mutations);
    commitRequest.setMode(CommitRequest.Mode.NON_TRANSACTIONAL);
    long startTime = System.currentTimeMillis(), endTime;

    if (throttler.throttleRequest(startTime)) {
      LOG.info("Delaying request due to previous failures");
      throttledSeconds.inc(WriteBatcherImpl.DATASTORE_BATCH_TARGET_LATENCY_MS / 1000);
      sleeper.sleep(WriteBatcherImpl.DATASTORE_BATCH_TARGET_LATENCY_MS);
      continue;
    }

    try {
      datastore.commit(commitRequest.build());
      endTime = System.currentTimeMillis();

      writeBatcher.addRequestLatency(endTime, endTime - startTime, mutations.size());
      throttler.successfulRequest(startTime);
      rpcSuccesses.inc();

      // Break if the commit threw no exception.
      break;
    } catch (DatastoreException exception) {
      if (exception.getCode() == Code.DEADLINE_EXCEEDED) {
        /* Most errors are not related to request size, and should not change our expectation of
         * the latency of successful requests. DEADLINE_EXCEEDED can be taken into
         * consideration, though. */
        endTime = System.currentTimeMillis();
        writeBatcher.addRequestLatency(endTime, endTime - startTime, mutations.size());
      }
      // Only log the code and message for potentially-transient errors. The entire exception
      // will be propagated upon the last retry.
      LOG.error(
          "Error writing batch of {} mutations to Datastore ({}): {}",
          mutations.size(),
          exception.getCode(),
          exception.getMessage());
      rpcErrors.inc();

      if (NON_RETRYABLE_ERRORS.contains(exception.getCode())) {
        throw exception;
      }
      if (!BackOffUtils.next(sleeper, backoff)) {
        LOG.error("Aborting after {} retries.", MAX_RETRIES);
        throw exception;
      }
    }
  }
  LOG.debug("Successfully wrote {} mutations", mutations.size());
  mutations.clear();
  mutationsSize = 0;
}