Java Code Examples for com.google.cloud.datastore.Transaction#isActive()

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

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

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

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

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

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

/**
 * 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 testTransactionalSingleEntityGroupReadOnly() {
  setUpQueryTests();
  Key taskListKey = datastore.newKeyFactory().setKind("TaskList").newKey("default");
  Entity taskListEntity = Entity.newBuilder(taskListKey).build();
  datastore.put(taskListEntity);
  // [START datastore_transactional_single_entity_group_read_only]
  Entity taskList;
  QueryResults<Entity> tasks;
  Transaction txn = datastore.newTransaction(
      TransactionOptions.newBuilder()
          .setReadOnly(ReadOnly.newBuilder().build())
          .build()
  );
  try {
    taskList = txn.get(taskListKey);
    Query<Entity> query = Query.newEntityQueryBuilder()
        .setKind("Task")
        .setFilter(PropertyFilter.hasAncestor(taskListKey))
        .build();
    tasks = txn.run(query);
    txn.commit();
  } finally {
    if (txn.isActive()) {
      txn.rollback();
    }
  }
  // [END datastore_transactional_single_entity_group_read_only]
  assertEquals(taskListEntity, taskList);
  assertNotNull(tasks.next());
  assertFalse(tasks.hasNext());
}
 

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

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

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

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

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

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

/**
 * Rolls back the given transaction, if it is still active.
 * 
 * @param transaction
 *          the transaction to roll back.
 */
static void rollbackIfActive(Transaction transaction) {
  try {
    if (transaction != null && transaction.isActive()) {
      transaction.rollback();
    }
  } catch (DatastoreException exp) {
    throw new EntityManagerException(exp);
  }
}
 

void flush() {
  if (!pending.isEmpty()) {
    log.debug("flush started, attempting to acquire permit");
    double wait = rateLimiter.acquire();
    log.debug("permit acquired after {} seconds", wait);
    Multimap<String, Mutation> toWrite = ArrayListMultimap.create();

    // drain the queue into a Map<shard, list<mutations_for_the_shard>>
    Mutation queued = pending.poll();
    while(queued != null) {
      toWrite.put(queued.getShard(), queued);
      queued = pending.poll();
    }

    // merge multimap of mutations into a list of single entities per shard to write as a batch
    List<FullEntity> list = new ArrayList<>();
    for (String shard: toWrite.keySet()) {
      Collection<Mutation> deltas = toWrite.get(shard);
      deltas.stream().reduce((deltaA, deltaB) ->
          new Mutation(shard,
              deltaA.getSizeDelta() + deltaB.getSizeDelta(),
              deltaA.getCountDelta() + deltaB.getCountDelta())
      ).ifPresent(merged -> {
        Entity shardCounter = getShardCounter(merged.getShard());
        FullEntity<Key> entity = FullEntity.newBuilder(shardCounter.getKey())
            .set(SIZE, shardCounter.getLong(SIZE) + merged.getSizeDelta())
            .set(COUNT, shardCounter.getLong(COUNT) + merged.getCountDelta())
            .build();
        list.add(entity);
      });
    }
    log.debug("sending {} mutations to datastore", list.size());
    // write the batch off to datastore
    if (!list.isEmpty()) {
      Transaction txn = datastore.newTransaction();
      try {
        txn.put(list.toArray(new FullEntity[list.size()]));
        txn.commit();
      } finally {
        if (txn.isActive()) {
          txn.rollback();
        }
      }
      log.debug("drained {} mutations to datastore", list.size());
    }
  }
}