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

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

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

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

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

/**
 * Verifies a {@code PortabilityJob} already exists for {@code jobId}, and updates the entry to
 * {@code job}, within a {@code Transaction}. If {@code validator} is non-null,
 * validator.validate() is called first in the transaction.
 *
 * @throws IOException if a job didn't already exist for {@code jobId} or there was a problem
 *     updating it @throws IllegalStateException if validator.validate() failed
 */
@Override
protected void updateJob(UUID jobId, PortabilityJob job, JobUpdateValidator validator)
    throws IOException {
  Preconditions.checkNotNull(jobId);
  Transaction transaction = datastore.newTransaction();
  Key key = getJobKey(jobId);

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

    if (validator != null) {
      PortabilityJob previousJob = PortabilityJob.fromMap(getProperties(previousEntity));
      validator.validate(previousJob, job);
    }

    Entity newEntity = createUpdatedEntity(key, job.toMap());
    transaction.put(newEntity);
    transaction.commit();
  } catch (Throwable t) {
    transaction.rollback();
    throw new IOException("Failed atomic update of jobId: " + jobId, t);
  }
}
 

@Override
public <T extends DataModel> void create(UUID jobId, String key, T model) throws IOException {
  Preconditions.checkNotNull(jobId);
  Transaction transaction = datastore.newTransaction();
  Key fullKey = getDataKey(jobId, key);
  Entity shouldNotExist = transaction.get(fullKey);
  if (shouldNotExist != null) {
    transaction.rollback();
    throw new IOException(
        "Record already exists for key: " + fullKey.getName() + ". Record: " + shouldNotExist);
  }

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

  try {
    transaction.put(entity);
  } catch (DatastoreException e) {
    throw new IOException(
        "Could not create initial record for jobID: " + jobId + ". Record: " + entity, e);
  }
  transaction.commit();
}
 

private void initializeDS() throws IOException {
  Transaction t = datastore.newTransaction();
  t.put(googleExecutor.createResultEntity("id1", JOB_ID, "idempotentId1"));
  t.put(googleExecutor.createResultEntity("id2", JOB_ID, "idempotentId2"));
  t.put(googleExecutor.createResultEntity("id3", JOB_ID, "idempotentId3"));
  t.put(googleExecutor.createErrorEntity("id4", JOB_ID,
      ErrorDetail.builder().setId("id4").setTitle("title").setException("error").build()));

  t.put(googleExecutor.createResultEntity("id1_job2", JOB_ID_2, "idempotentId1_job2"));
  t.commit();
}
 

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

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

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