com.google.datastore.v1.Mutation 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());
}
 

@ProcessElement
public void processElement(ProcessContext c) throws Exception {
  Mutation write = c.element();
  int size = write.getSerializedSize();
  if (mutations.size() > 0
      && mutationsSize + size >= DatastoreV1.DATASTORE_BATCH_UPDATE_BYTES_LIMIT) {
    flushBatch();
  }
  mutations.add(c.element());
  mutationsSize += size;
  if (mutations.size() >= writeBatcher.nextBatchSize(System.currentTimeMillis())) {
    flushBatch();
  }
}
 

@Override
public Mutation apply(Entity entity) {
  // Verify that the entity to write has a complete key.
  checkArgument(
      isValidKey(entity.getKey()),
      "Entities to be written to the Cloud Datastore must have complete keys:\n%s",
      entity);

  return makeUpsert(entity).build();
}
 

@Override
public Mutation apply(Entity entity) {
  // Verify that the entity to delete has a complete key.
  checkArgument(
      isValidKey(entity.getKey()),
      "Entities to be deleted from the Cloud Datastore must have complete keys:\n%s",
      entity);

  return makeDelete(entity.getKey()).build();
}
 

@Override
public Mutation apply(Key key) {
  // Verify that the entity to delete has a complete key.
  checkArgument(
      isValidKey(key),
      "Keys to be deleted from the Cloud Datastore must be complete:\n%s",
      key);

  return makeDelete(key).build();
}
 

@Test
/** Test that entities with valid keys are transformed to upsert mutations. */
public void testAddEntities() throws Exception {
  Key key = makeKey("bird", "finch").build();
  Entity entity = Entity.newBuilder().setKey(key).build();
  UpsertFn upsertFn = new UpsertFn();

  Mutation expectedMutation = makeUpsert(entity).build();
  assertEquals(expectedMutation, upsertFn.apply(entity));
}
 

/** Test that entities with valid keys are transformed to delete mutations. */
@Test
public void testDeleteEntities() throws Exception {
  Key key = makeKey("bird", "finch").build();
  Entity entity = Entity.newBuilder().setKey(key).build();
  DeleteEntityFn deleteEntityFn = new DeleteEntityFn();

  Mutation expectedMutation = makeDelete(entity.getKey()).build();
  assertEquals(expectedMutation, deleteEntityFn.apply(entity));
}
 

/** Test that valid keys are transformed to delete mutations. */
@Test
public void testDeleteKeys() {
  Key key = makeKey("bird", "finch").build();
  DeleteKeyFn deleteKeyFn = new DeleteKeyFn();

  Mutation expectedMutation = makeDelete(key).build();
  assertEquals(expectedMutation, deleteKeyFn.apply(key));
}
 

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

/**
 * Note that {@code projectId} is only {@code @Nullable} as a matter of build order, but if it
 * is {@code null} at instantiation time, an error will be thrown.
 */
Mutate(
    @Nullable ValueProvider<String> projectId,
    @Nullable String localhost,
    SimpleFunction<T, Mutation> mutationFn) {
  this.projectId = projectId;
  this.localhost = localhost;
  this.mutationFn = checkNotNull(mutationFn);
}
 

@Override
public ListenableFuture<MutationResult> executeAsync(final MutationStatement statement, final ListenableFuture<TransactionResult> txn) {
  final List<Mutation> mutations = Optional
    .ofNullable(statement)
    .flatMap(s -> Optional.of(ImmutableList.of(s.getPb(config.getNamespace()))))
    .orElse(ImmutableList.of());

  return executeAsyncMutations(mutations, txn);
}
 

@Override
public Mutation getPb(final String namespace) {
  final com.google.datastore.v1.Mutation.Builder mutation =
    com.google.datastore.v1.Mutation.newBuilder();

  return mutation.setInsert(entity.build().getPb(namespace)).build();
}
 

@Override
public Mutation getPb(final String namespace) {
  final com.google.datastore.v1.Mutation.Builder mutation =
    com.google.datastore.v1.Mutation.newBuilder();
  if (upsert) {
    mutation.setUpsert(entity.build().getPb(namespace));
  } else {
    mutation.setUpdate(entity.build().getPb(namespace));
  }
  return mutation.build();
}
 

@Override
public Mutation.Builder apply(Entity entity) {
  return makeDelete(entity.getKey());
}
 

public List<Mutation> getPb(final String namespace) {
  return statements
      .stream()
      .map(m -> m.getPb(namespace))
      .collect(Collectors.toList());
}
 

@Override
public Mutation getPb(final String namespace) {
  final com.google.datastore.v1.Key mutationKey = getKey().getPb(namespace);
  return com.google.datastore.v1.Mutation.newBuilder().setDelete(mutationKey).build();
}
 

@Override
public Mutation.Builder apply(Entity entity) {
  return makeUpsert(entity);
}
 

/**
 * @param entity the entity to insert
 * @return a mutation that will insert an entity
 */
public static Mutation.Builder makeInsert(Entity entity) {
  return Mutation.newBuilder().setInsert(entity);
}
 

/**
 * @param entity the entity to update
 * @return a mutation that will update an entity
 */
public static Mutation.Builder makeUpdate(Entity entity) {
  return Mutation.newBuilder().setUpdate(entity);
}
 

/**
 * @param entity the entity to upsert
 * @return a mutation that will upsert an entity
 */
public static Mutation.Builder makeUpsert(Entity entity) {
  return Mutation.newBuilder().setUpsert(entity);
}
 

/**
 * @param key the key of the entity to delete
 * @return a mutation that will delete an entity
 */
public static Mutation.Builder makeDelete(Key key) {
  return Mutation.newBuilder().setDelete(key);
}
 

Mutation getPb(final String namespace); 

Mutation.Builder apply(Entity entity);