com.google.cloud.datastore.Transaction Java Examples

@Override
public void run(Transaction tx, Key userKey, String content) {
  Entity user = tx.get(userKey);
  if (user == null) {
    System.out.println("Adding a new user.");
    user = Entity.newBuilder(userKey).set("count", 1).build();
    tx.add(user);
  } else {
    user = Entity.newBuilder(user).set("count", user.getLong("count") + 1L).build();
    tx.update(user);
  }
  IncompleteKey commentKey = IncompleteKey.newBuilder(userKey, COMMENT_KIND).build();
  FullEntity<IncompleteKey> comment =
      FullEntity.newBuilder(commentKey)
          .set("content", content)
          .set("timestamp", Timestamp.now())
          .build();
  tx.addWithDeferredIdAllocation(comment);
  System.out.printf("Adding a comment to user '%s'.%n", userKey.getName());
}
 

/**
 * Delete a value stored in the project's datastore.
 *
 * @param sessionId Request from which the session is extracted.
 */
protected void deleteSessionVariables(String sessionId, String... varNames) {
  if (sessionId.equals("")) {
    return;
  }
  Key key = keyFactory.newKey(sessionId);
  Transaction transaction = datastore.newTransaction();
  try {
    Entity stateEntity = transaction.get(key);
    if (stateEntity != null) {
      Entity.Builder builder = Entity.newBuilder(stateEntity);
      StringBuilder delNames = new StringBuilder();
      for (String varName : varNames) {
        delNames.append(varName + " ");
        builder = builder.remove(varName);
      }
      datastore.update(builder.build());
    }
  } finally {
    if (transaction.isActive()) {
      transaction.rollback();
    }
  }
}
 

protected void deleteSessionWithValue(String varName, String varValue) {
  Transaction transaction = datastore.newTransaction();
  try {
    Query<Entity> query = Query.newEntityQueryBuilder()
        .setKind("SessionVariable")
        .setFilter(PropertyFilter.eq(varName, varValue))
        .build();
    QueryResults<Entity> resultList = transaction.run(query);
    while (resultList.hasNext()) {
      Entity stateEntity = resultList.next();
      transaction.delete(stateEntity.getKey());
    }
    transaction.commit();
  } finally {
    if (transaction.isActive()) {
      transaction.rollback();
    }
  }
}
 

protected void deleteSessionWithValue(String varName, String varValue) {
  Transaction transaction = datastore.newTransaction();
  try {
    Query<Entity> query = Query.newEntityQueryBuilder()
        .setKind("SessionVariable")
        .setFilter(PropertyFilter.eq(varName, varValue))
        .build();
    QueryResults<Entity> resultList = transaction.run(query);
    while (resultList.hasNext()) {
      Entity stateEntity = resultList.next();
      transaction.delete(stateEntity.getKey());
    }
    transaction.commit();
  } finally {
    if (transaction.isActive()) {
      transaction.rollback();
    }
  }
}
 

/**
 * Delete a value stored in the project's datastore.
 *
 * @param sessionId Request from which the session is extracted.
 */
protected void deleteSessionVariables(String sessionId, String... varNames) {
  if (sessionId.equals("")) {
    return;
  }
  Key key = keyFactory.newKey(sessionId);
  Transaction transaction = datastore.newTransaction();
  try {
    Entity stateEntity = transaction.get(key);
    if (stateEntity != null) {
      Entity.Builder builder = Entity.newBuilder(stateEntity);
      StringBuilder delNames = new StringBuilder();
      for (String varName : varNames) {
        delNames.append(varName + " ");
        builder = builder.remove(varName);
      }
      datastore.update(builder.build());
    }
  } finally {
    if (transaction.isActive()) {
      transaction.rollback();
    }
  }
}
 

private Long getCount(String fieldName) {
  Transaction txn = datastore.newTransaction(
      TransactionOptions.newBuilder()
          .setReadOnly(ReadOnly.newBuilder().build())
          .build()
  );

  QueryResults<ProjectionEntity> results;
  try {
    Query<ProjectionEntity> countQuery = Query.newProjectionEntityQueryBuilder()
        .setKind(SHARD)
        .setNamespace(namespace)
        .setProjection(fieldName)
        .build();

    results = datastore.run(countQuery);
    return StreamSupport.stream(Spliterators.spliteratorUnknownSize(results, Spliterator.NONNULL), false)
        .map(entity -> Long.valueOf(entity.getLong(fieldName)))
        .reduce(0L, (valueA, valueB) -> valueA + valueB);
  } finally {
    if (txn.isActive()) {
      txn.rollback();
    }
  }
}
 

protected void deleteSessionWithValue(String varName, String varValue) {
  Transaction transaction = datastore.newTransaction();
  try {
    Query<Entity> query = Query.newEntityQueryBuilder()
        .setKind("SessionVariable")
        .setFilter(PropertyFilter.eq(varName, varValue))
        .build();
    QueryResults<Entity> resultList = transaction.run(query);
    while (resultList.hasNext()) {
      Entity stateEntity = resultList.next();
      transaction.delete(stateEntity.getKey());
    }
    transaction.commit();
  } finally {
    if (transaction.isActive()) {
      transaction.rollback();
    }
  }
}
 

/**
 * Delete a value stored in the project's datastore.
 *
 * @param sessionId Request from which the session is extracted.
 */
protected void deleteSessionVariables(String sessionId, String... varNames) {
  if (sessionId.equals("")) {
    return;
  }
  Key key = keyFactory.newKey(sessionId);
  Transaction transaction = datastore.newTransaction();
  try {
    Entity stateEntity = transaction.get(key);
    if (stateEntity != null) {
      Entity.Builder builder = Entity.newBuilder(stateEntity);
      StringBuilder delNames = new StringBuilder();
      for (String varName : varNames) {
        delNames.append(varName + " ");
        builder = builder.remove(varName);
      }
      datastore.update(builder.build());
    }
  } finally {
    if (transaction.isActive()) {
      transaction.rollback();
    }
  }
}
 

/**
 * Marks a task entity as done.
 *
 * @param id The ID of the task entity as given by {@link Key#id()}
 * @return true if the task was found, false if not
 * @throws DatastoreException if the transaction fails
 */
boolean markDone(long id) {
  Transaction transaction = datastore.newTransaction();
  try {
    Entity task = transaction.get(keyFactory.newKey(id));
    if (task != null) {
      transaction.put(Entity.newBuilder(task).set("done", true).build());
    }
    transaction.commit();
    return task != null;
  } finally {
    if (transaction.isActive()) {
      transaction.rollback();
    }
  }
}
 

protected void deleteSessionWithValue(String varName, String varValue) {
  Transaction transaction = datastore.newTransaction();
  try {
    Query<Entity> query = Query.newEntityQueryBuilder()
        .setKind("SessionVariable")
        .setFilter(PropertyFilter.eq(varName, varValue))
        .build();
    QueryResults<Entity> resultList = transaction.run(query);
    while (resultList.hasNext()) {
      Entity stateEntity = resultList.next();
      transaction.delete(stateEntity.getKey());
    }
    transaction.commit();
  } finally {
    if (transaction.isActive()) {
      transaction.rollback();
    }
  }
}
 

/**
 * Delete a value stored in the project's datastore.
 * @param sessionId Request from which the session is extracted.
 */
protected void deleteSessionVariables(String sessionId, String... varNames) {
  if (sessionId.equals("")) {
    return;
  }
  Key key = keyFactory.newKey(sessionId);
  Transaction transaction = datastore.newTransaction();
  try {
    Entity stateEntity = transaction.get(key);
    if (stateEntity != null) {
      Entity.Builder builder = Entity.newBuilder(stateEntity);
      StringBuilder delNames = new StringBuilder();
      for (String varName : varNames) {
        delNames.append(varName + " ");
        builder = builder.remove(varName);
      }
      datastore.update(builder.build());
    }
  } finally {
    if (transaction.isActive()) {
      transaction.rollback();
    }
  }
}
 

void transferFunds(Key fromKey, Key toKey, long amount) {
  Transaction txn = datastore.newTransaction();
  try {
    List<Entity> entities = txn.fetch(fromKey, toKey);
    Entity from = entities.get(0);
    Entity updatedFrom =
        Entity.newBuilder(from).set("balance", from.getLong("balance") - amount).build();
    Entity to = entities.get(1);
    Entity updatedTo = Entity.newBuilder(to).set("balance", to.getLong("balance") + amount)
        .build();
    txn.put(updatedFrom, updatedTo);
    txn.commit();
  } finally {
    if (txn.isActive()) {
      txn.rollback();
    }
  }
}
 

@Test
public void testTransactionalGetOrCreate() {
  // [START datastore_transactional_get_or_create]
  Entity task;
  Transaction txn = datastore.newTransaction();
  try {
    task = txn.get(taskKey);
    if (task == null) {
      task = Entity.newBuilder(taskKey).build();
      txn.put(task);
      txn.commit();
    }
  } finally {
    if (txn.isActive()) {
      txn.rollback();
    }
  }
  // [END datastore_transactional_get_or_create]
  assertEquals(task, datastore.get(taskKey));
}
 

@Override
public void run(Transaction tx, Key userKey, Void arg) {
  Entity user = tx.get(userKey);
  if (user == null) {
    System.out.println("Nothing to delete, user does not exist.");
    return;
  }
  Query<Key> query =
      Query.newKeyQueryBuilder()
          .setNamespace(NAMESPACE)
          .setKind(COMMENT_KIND)
          .setFilter(PropertyFilter.hasAncestor(userKey))
          .build();
  QueryResults<Key> comments = tx.run(query);
  int count = 0;
  while (comments.hasNext()) {
    tx.delete(comments.next());
    count++;
  }
  tx.delete(userKey);
  System.out.printf("Deleting user '%s' and %d comment[s].%n", userKey.getName(), count);
}
 

static InstrumentedTransaction of(Stats stats, Transaction transaction) {
  Objects.requireNonNull(stats, "stats");
  Objects.requireNonNull(transaction, "transaction");
  return new InstrumentedTransaction() {
    @Override
    public Transaction transaction() {
      return transaction;
    }

    @Override
    public Stats stats() {
      return stats;
    }

    @Override
    public DatastoreBatchWriter batchWriter() {
      return transaction;
    }
  };
}
 

@Override
public void run(Transaction tx, Key userKey, Contact contact) {
  Entity user = tx.get(userKey);
  if (user == null) {
    System.out.println("Adding a new user.");
    user = Entity.newBuilder(userKey).set("count", 0L).build();
    tx.add(user);
  }
  FullEntity<IncompleteKey> contactEntity =
      FullEntity.newBuilder()
          .set("email", contact.email())
          .set("phone", contact.phone())
          .build();
  tx.update(Entity.newBuilder(user).set("contact", contactEntity).build());
  System.out.printf("Setting contact for user '%s'.%n", userKey.getName());
}
 

/**
 * Inserts a new {@link PortabilityJob} keyed by {@code jobId} in Datastore.
 *
 * <p>To update an existing {@link PortabilityJob} instead, use {@link JobStore#update}.
 *
 * @throws IOException if a job already exists for {@code jobId}, or if there was a different
 *     problem inserting the job.
 */
@Override
public void createJob(UUID jobId, PortabilityJob job) throws IOException {
  Preconditions.checkNotNull(jobId);
  Transaction transaction = datastore.newTransaction();
  Entity shouldNotExist = transaction.get(getJobKey(jobId));
  if (shouldNotExist != null) {
    transaction.rollback();
    throw new IOException(
        "Record already exists for jobID: " + jobId + ". Record: " + shouldNotExist);
  }
  Entity entity = createNewEntity(jobId, job.toMap());
  try {
    transaction.put(entity);
  } catch (DatastoreException e) {
    transaction.rollback();
    throw new IOException(
        "Could not create initial record for jobID: " + jobId + ". Record: " + entity, e);
  }
  transaction.commit();
}
 

@Override
public <T extends DataModel> void update(UUID jobId, String key, T model) {
  Transaction transaction = datastore.newTransaction();
  Key entityKey = getDataKey(jobId, key);

  try {
    Entity previousEntity = transaction.get(entityKey);
    if (previousEntity == null) {
      throw new IOException("Could not find record for data key: " + entityKey.getName());
    }

    String serialized = objectMapper.writeValueAsString(model);
    Entity entity =
        Entity.newBuilder(entityKey)
            .set(CREATED_FIELD, Timestamp.now())
            .set(model.getClass().getName(), serialized)
            .build();

    transaction.put(entity);
    transaction.commit();
  } catch (IOException t) {
    transaction.rollback();
    throw new RuntimeException("Failed atomic update of key: " + key, t);
  }
}
 

@Test
public void testPutGetMultiple() {
  Transaction transaction = datastore.newTransaction();
  TransactionSnippets transactionSnippets = new TransactionSnippets(transaction);
  transactionSnippets.multipleAddEntities(
      registerKey("add_get_multiple_key_1"), registerKey("put_get_multiple_key_2"));

  transaction = datastore.newTransaction();
  transactionSnippets = new TransactionSnippets(transaction);
  List<Entity> entities =
      transactionSnippets.getMultiple("add_get_multiple_key_1", "put_get_multiple_key_2");
  assertEquals(2, entities.size());
  Set<String> values =
      ImmutableSet.of(
          entities.get(0).getString("propertyName"), entities.get(1).getString("propertyName"));
  assertTrue(values.contains("value1"));
  assertTrue(values.contains("value2"));
}
 

@Override
public void addCounts(UUID jobId, Map<String, Integer> newCounts) throws IOException {
  if (newCounts == null) {
    return;
  }

  Transaction transaction = datastore.newTransaction();

  for (String dataType : newCounts.keySet()) {
    Key key = getCountsKey(jobId, dataType);
    Entity current = datastore.get(key);
    Integer oldCount = 0;

    if (current != null && current.getNames().contains(COUNTS_FIELD)) {
      // Datastore only allows Long properties, but we only ever write Integers through this
      // interface so the conversion is OK
      oldCount = Math.toIntExact(current.getLong(COUNTS_FIELD));
    }
    transaction.put(
        GoogleCloudUtils.createEntityBuilder(
                key, ImmutableMap.of(COUNTS_FIELD, oldCount + newCounts.get(dataType)))
            .build());
  }
  transaction.commit();
}
 

@Test
public void testRun() {
  Key key1 = datastore.newKeyFactory().setKind("ParentKind").newKey("run_key_1");
  Entity entity1 = Entity.newBuilder(key1).set("description", "run1").build();
  datastore.put(entity1);
  Key key2 =
      datastore
          .newKeyFactory()
          .setKind("MyKind")
          .addAncestor(PathElement.of("ParentKind", "run_key_1"))
          .newKey("run_key_2");
  registerKey(key1);
  registerKey(key2);
  Entity entity2 = Entity.newBuilder(key2).set("description", "run2").build();
  datastore.put(entity2);

  Transaction transaction = datastore.newTransaction();
  TransactionSnippets transactionSnippets = new TransactionSnippets(transaction);
  List<Entity> entities = transactionSnippets.run("run_key_1");
  assertEquals(1, entities.size());
  assertEquals(entity2, entities.get(0));
}
 

@Test
public void testIsActive() {
  Transaction transaction = datastore.newTransaction();
  TransactionSnippets transactionSnippets = new TransactionSnippets(transaction);
  Key key = transactionSnippets.isActive();
  Entity result = datastore.get(key);
  assertNotNull(result);
  datastore.delete(key);
}
 

@Test
public void testEntityAddGet() {
  Transaction transaction = datastore.newTransaction();
  TransactionSnippets transactionSnippets = new TransactionSnippets(transaction);
  transactionSnippets.addSingleEntity(registerKey("add_get_key"));

  transaction = datastore.newTransaction();
  transactionSnippets = new TransactionSnippets(transaction);
  Entity entity = transactionSnippets.get("add_get_key");
  assertEquals("value", entity.getString("propertyName"));
}
 

@Test
public void testActive() {
  Transaction transaction = datastore.newTransaction();
  TransactionSnippets transactionSnippets = new TransactionSnippets(transaction);
  Key key = transactionSnippets.active();
  Entity result = datastore.get(key);
  assertNotNull(result);
  datastore.delete(key);
}
 

@Test
public void testRollback() {
  Transaction transaction = datastore.newTransaction();
  TransactionSnippets transactionSnippets = new TransactionSnippets(transaction);
  Key key = transactionSnippets.rollback();
  Entity result = datastore.get(key);
  assertNull(result);
}
 

/**
 * Stores the state value in each key-value pair in the project's datastore.
 *
 * @param sessionId Request from which to extract session.
 * @param varName   the name of the desired session variable
 * @param varValue  the value of the desired session variable
 */
protected void setSessionVariables(String sessionId, Map<String, String> setMap) {
  if (sessionId.equals("")) {
    return;
  }
  Key key = keyFactory.newKey(sessionId);
  Transaction transaction = datastore.newTransaction();
  DateTime dt = DateTime.now(DateTimeZone.UTC);
  dt.toString(dtf);
  try {
    Entity stateEntity = transaction.get(key);
    Entity.Builder seBuilder;
    if (stateEntity == null) {
      seBuilder = Entity.newBuilder(key);
    } else {
      seBuilder = Entity.newBuilder(stateEntity);
    }
    for (String varName : setMap.keySet()) {
      seBuilder.set(varName, setMap.get(varName));
    }
    transaction.put(seBuilder.set("lastModified", dt.toString(dtf)).build());
    transaction.commit();
  } finally {
    if (transaction.isActive()) {
      transaction.rollback();
    }
  }
}
 

@Test
public void testCommit() {
  Transaction transaction = datastore.newTransaction();
  TransactionSnippets transactionSnippets = new TransactionSnippets(transaction);
  Key key = transactionSnippets.commit();
  Entity result = datastore.get(key);
  assertNotNull(result);
  datastore.delete(key);
}
 

/**
 * Take an HttpServletRequest, and copy all of the current session variables over to it
 *
 * @param req Request from which to extract session.
 * @return a map of strings containing all the session variables loaded or an empty map.
 */
protected Map<String, String> loadSessionVariables(HttpServletRequest req)
    throws ServletException {
  Map<String, String> datastoreMap = new HashMap<>();
  String sessionId = getCookieValue(req, "bookshelfSessionId");
  if (sessionId.equals("")) {
    return datastoreMap;
  }
  Key key = keyFactory.newKey(sessionId);
  Transaction transaction = datastore.newTransaction();
  try {
    Entity stateEntity = transaction.get(key);
    StringBuilder logNames = new StringBuilder();
    if (stateEntity != null) {
      for (String varName : stateEntity.getNames()) {
        req.getSession().setAttribute(varName, stateEntity.getString(varName));
        datastoreMap.put(varName, stateEntity.getString(varName));
        logNames.append(varName + " ");
      }
    }
  } finally {
    if (transaction.isActive()) {
      transaction.rollback();
    }
  }
  return datastoreMap;
}
 

@Test
public void testFetchDeleteEntitiesWithKeys() {
  Key key1 = datastore.newKeyFactory().setKind("MyKind").newKey("fetch_key_1");
  Key key2 = datastore.newKeyFactory().setKind("MyKind").newKey("fetch_key_2");
  Entity entity1 = Entity.newBuilder(key1).set("description", "fetch1").build();
  Entity entity2 = Entity.newBuilder(key2).set("description", "fetch2").build();
  datastore.put(entity1, entity2);
  registerKey("fetch_key_1");
  registerKey("fetch_key_2");

  Transaction transaction = datastore.newTransaction();
  TransactionSnippets transactionSnippets = new TransactionSnippets(transaction);
  Set<Entity> entities =
      Sets.newHashSet(transactionSnippets.fetchEntitiesWithKeys("fetch_key_1", "fetch_key_2"));
  assertEquals(2, entities.size());
  assertTrue(entities.contains(entity1));
  assertTrue(entities.contains(entity2));

  transaction = datastore.newTransaction();
  transactionSnippets = new TransactionSnippets(transaction);
  transactionSnippets.multipleDeleteEntities("fetch_key_1", "fetch_key_2");

  transaction = datastore.newTransaction();
  transactionSnippets = new TransactionSnippets(transaction);
  List<Entity> deletedEntities =
      transactionSnippets.fetchEntitiesWithKeys("fetch_key_1", "fetch_key_2");
  assertNull(deletedEntities.get(0));
  assertNull(deletedEntities.get(1));
}
 

@Test
public void testPutGetMultipleDeferredId() {
  Transaction transaction = datastore.newTransaction();
  TransactionSnippets transactionSnippets = new TransactionSnippets(transaction);
  List<Key> keys = transactionSnippets.multiplePutEntitiesDeferredId();
  assertEquals(2, keys.size());
  Key key1 = keys.get(0);
  registerKey(key1);
  Entity entity1 = datastore.get(key1);
  assertEquals("value1", entity1.getString("propertyName"));
  Key key2 = keys.get(1);
  registerKey(key2);
  Entity entity2 = datastore.get(key2);
  assertEquals("value2", entity2.getString("propertyName"));
}