com.google.cloud.datastore.Key Java Examples

/**
 * Restore the attributes and metadata of the session from Datastore Entities.
 *
 * @param entities An iterator of entity, containing the metadata and attributes of the session.
 * @throws ClassNotFoundException The class in attempt to be deserialized is not present in the
 *                                application.
 * @throws IOException Error during the deserialization of the object.
 */
public void restoreFromEntities(Key sessionKey, Iterable<Entity> entities) throws
    ClassNotFoundException, IOException {
  Entity metadataEntity = null;
  List<Entity> attributeEntities = new LinkedList<>();
  for (Entity entity : entities) {
    if (entity.getKey().equals(sessionKey)) {
      metadataEntity = entity;
    } else {
      attributeEntities.add(entity);
    }
  }

  if (metadataEntity == null) {
    throw new IOException("The serialized session is missing the metadata entity");
  }

  restoreMetadataFromEntity(metadataEntity);
  restoreAttributesFromEntity(attributeEntities);
  setId(sessionKey.getName());
  initialAttributes.addAll(Collections.list(getAttributeNames()));
}
 

/** Example of getting multiple entity objects. */
// [TARGET get(Iterable, ReadOption...)]
// [VARIABLE "my_first_key_name"]
// [VARIABLE "my_second_key_name"]
public List<Entity> getEntitiesWithKeys(String firstKeyName, String secondKeyName) {
  // TODO change so that it's not necessary to hold the entities in a list for integration testing
  // [START getEntitiesWithKeys]
  KeyFactory keyFactory = datastore.newKeyFactory().setKind("MyKind");
  Key firstKey = keyFactory.newKey(firstKeyName);
  Key secondKey = keyFactory.newKey(secondKeyName);
  Iterator<Entity> entitiesIterator = datastore.get(Lists.newArrayList(firstKey, secondKey));
  List<Entity> entities = Lists.newArrayList();
  while (entitiesIterator.hasNext()) {
    Entity entity = entitiesIterator.next();
    // do something with the entity
    entities.add(entity);
  }
  // [END getEntitiesWithKeys]
  return entities;
}
 

/**
 * Store the metadata of the session in an entity.
 * @param sessionKey Identifier of the session on the Datastore
 * @return An entity containing the metadata.
 */
@VisibleForTesting
Entity saveMetadataToEntity(Key sessionKey) {
  Entity.Builder sessionEntity = Entity.newBuilder(sessionKey)
      .set(SessionMetadata.CREATION_TIME, getCreationTime())
      .set(SessionMetadata.LAST_ACCESSED_TIME, getLastAccessedTime())
      .set(SessionMetadata.MAX_INACTIVE_INTERVAL, getMaxInactiveInterval())
      .set(SessionMetadata.IS_NEW, isNew())
      .set(SessionMetadata.IS_VALID, isValid())
      .set(SessionMetadata.THIS_ACCESSED_TIME, getThisAccessedTime());

  // A negative time indicates that the session should never time out
  if (getMaxInactiveInterval() >= 0) {
    sessionEntity.set(SessionMetadata.EXPIRATION_TIME,
        getLastAccessedTime() + getMaxInactiveInterval() * 1000);
  }

  return sessionEntity.build();
}
 

@Test
public void testAttributesDeserialization() throws Exception {
  Entity metadata = mock(Entity.class);
  when(metadata.getBoolean(any())).thenReturn(true);
  when(sessionKey.getName()).thenReturn("count");
  when(metadata.getKey()).thenReturn(sessionKey);
  int count = 5;
  ByteArrayOutputStream bos = new ByteArrayOutputStream();
  ObjectOutputStream oos = new ObjectOutputStream(bos);
  oos.writeObject(count);

  Entity valueEntity = Entity.newBuilder(
      (Key)when(mock(Key.class).getName()).thenReturn("count").getMock())
      .set("value", Blob.copyFrom(bos.toByteArray()))
      .build();

  DatastoreSession session = new DatastoreSession(sessionManager);
  session.restoreFromEntities(sessionKey, Arrays.asList(metadata, valueEntity));

  assertEquals(count, session.getAttribute("count"));
}
 

/**
 * Updates a user in Cloud Datastore.
 */
public User updateUser(String id, String name, String email) {
  failIfInvalid(name, email);
  Key key = keyFactory.newKey(id);
  Entity entity = datastore.get(key);
  if (entity == null) {
    throw new IllegalArgumentException("No user with id '" + id + "' found");
  } else {
    entity = Entity.newBuilder(entity)
        .set("id", id)
        .set("name", name)
        .set("email", email)
        .build();
    datastore.update(entity);
  }
  return new User(id, name, email);
}
 

/** Example of adding multiple entities. */
// [TARGET add(FullEntity...)]
// [VARIABLE "my_key_name1"]
// [VARIABLE "my_key_name2"]
public void batchAddEntities(String keyName1, String keyName2) {
  // [START batchAddEntities]
  Key key1 = datastore.newKeyFactory().setKind("MyKind").newKey(keyName1);
  Entity.Builder entityBuilder1 = Entity.newBuilder(key1);
  entityBuilder1.set("propertyName", "value1");
  Entity entity1 = entityBuilder1.build();

  Key key2 = datastore.newKeyFactory().setKind("MyKind").newKey(keyName2);
  Entity.Builder entityBuilder2 = Entity.newBuilder(key2);
  entityBuilder2.set("propertyName", "value2");
  Entity entity2 = entityBuilder2.build();

  try {
    datastore.add(entity1, entity2);
  } catch (DatastoreException ex) {
    if ("ALREADY_EXISTS".equals(ex.getReason())) {
      // at least one of entity1.getKey() and entity2.getKey() already exists
    }
  }
  // [END batchAddEntities]
}
 

@Override
public void save(Message message) {
  // Save message to "messages"
  Datastore datastore = getDatastoreInstance();
  Key key = datastore.allocateId(keyFactory.newKey());

  Entity.Builder messageEntityBuilder = Entity.newBuilder(key)
      .set("messageId", message.getMessageId());

  if (message.getData() != null) {
    messageEntityBuilder = messageEntityBuilder.set("data", message.getData());
  }

  if (message.getPublishTime() != null) {
    messageEntityBuilder = messageEntityBuilder.set("publishTime", message.getPublishTime());
  }
  datastore.put(messageEntityBuilder.build());
}
 

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

@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 testEntityWithParent() {
  // [START datastore_entity_with_parent]
  Key taskKey = datastore.newKeyFactory()
      .addAncestors(PathElement.of("TaskList", "default"))
      .setKind("Task")
      .newKey("sampleTask");
  Entity task = Entity.newBuilder(taskKey)
      .set("category", "Personal")
      .set("done", false)
      .set("priority", 4)
      .set("description", "Learn Cloud Datastore")
      .build();
  // [END datastore_entity_with_parent]
  assertValidEntity(task);
}
 

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

/** Example of verifying if a transaction is active. */
// [TARGET isActive()]
public Key isActive() {
  Datastore datastore = transaction.getDatastore();
  // [START isActive]
  // create an entity
  KeyFactory keyFactory = datastore.newKeyFactory().setKind("MyKind");
  Key key = datastore.allocateId(keyFactory.newKey());
  Entity entity = Entity.newBuilder(key).set("description", "active()").build();
  // calling transaction.active() now would return true
  try {
    // add the entity and commit
    transaction.put(entity);
    transaction.commit();
  } finally {
    // if committing succeeded
    // then transaction.active() will be false
    if (transaction.isActive()) {
      // otherwise it's true and we need to rollback
      transaction.rollback();
    }
  }
  // [END isActive]
  return key;
}
 

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

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

/** Example of updating multiple entities. */
// [TARGET update(Entity...)]
// [VARIABLE "my_key_name_1"]
// [VARIABLE "my_key_name_2"]
public void batchUpdateEntities(String keyName1, String keyName2) {
  // [START batchUpdateEntities]
  Key key1 = datastore.newKeyFactory().setKind("MyKind").newKey(keyName1);
  Entity.Builder entityBuilder1 = Entity.newBuilder(key1);
  entityBuilder1.set("propertyName", "updatedValue1");
  Entity entity1 = entityBuilder1.build();

  Key key2 = datastore.newKeyFactory().setKind("MyKind").newKey(keyName2);
  Entity.Builder entityBuilder2 = Entity.newBuilder(key2);
  entityBuilder2.set("propertyName", "updatedValue2");
  Entity entity2 = entityBuilder2.build();

  datastore.update(entity1, entity2);
  // [END batchUpdateEntities]
}
 

@Override
public String convert(BaseKey key) {
	String name = null;
	// embedded entities have IncompleteKey, and have no inner value
	if (key instanceof Key) {
		name = ((Key) key).getName();
		if (name == null) {
			throw new DatastoreDataException("The given key doesn't have a String name value but " +
					"a conversion to String was attempted: " + key);
		}
	}
	return name;
}
 

@Override
public void save(Message message) {
  // Save message to "messages"
  Datastore datastore = getDatastoreInstance();
  Key key = datastore.allocateId(keyFactory.newKey());

  Entity.Builder messageEntityBuilder =
      Entity.newBuilder(key).set("messageId", message.getMessageId());

  String translated = message.getTranslated();
  if (translated != null) {
    messageEntityBuilder = messageEntityBuilder.set("data", translated);
  }

  if (message.getPublishTime() != null) {
    messageEntityBuilder = messageEntityBuilder.set("publishTime", message.getPublishTime());
  }

  if (message.getSourceLang() != null) {
    messageEntityBuilder = messageEntityBuilder.set("sourceLang", message.getSourceLang());
  }

  if (message.getTargetLang() != null) {
    messageEntityBuilder = messageEntityBuilder.set("targetLang", message.getTargetLang());
  }
  datastore.put(messageEntityBuilder.build());
}
 

Map<Integer, Long> shardsForCounter(String counterId) throws IOException {
  final List<Key> shardKeys = IntStream.range(0, NUM_SHARDS).mapToObj(
      index -> datastore.newKeyFactory().setKind(KIND_COUNTER_SHARD).newKey(
          String.format("%s-%d", counterId, index)))
      .collect(toList());

  final Map<Integer, Long> fetchedShards = new HashMap<>();
  datastore.get(shardKeys, shard -> fetchedShards.put(
      (int) shard.getLong(PROPERTY_SHARD_INDEX),
      shard.getLong(PROPERTY_SHARD_VALUE)));
  return fetchedShards;
}
 

@Before
public void setUp() {
  StructuredQuery<Key> query = Query.newKeyQueryBuilder().build();
  QueryResults<Key> result = DATASTORE.run(query);
  DATASTORE.delete(Iterators.toArray(result, Key.class));
  DATASTORE.add(USER_RECORD);
}
 

long getLimitForCounter(String counterId) throws IOException {
  if (GLOBAL_RESOURCE_ID.equals(counterId)) {
    // missing global resource means free to go
    return config().globalConcurrency().orElse(Long.MAX_VALUE);
  }

  final Key limitKey = datastore.newKeyFactory().setKind(KIND_COUNTER_LIMIT).newKey(counterId);
  final Entity limitEntity = datastore.get(limitKey);
  if (limitEntity == null) {
    throw new IllegalArgumentException("No limit found in Datastore for " + counterId);
  } else {
    return limitEntity.getLong(PROPERTY_LIMIT);
  }
}
 

@Override
public Long convert(BaseKey key) {
	Long id = null;
	// embedded entities have IncompleteKey, and have no inner value
	if (key instanceof Key) {
		id = ((Key) key).getId();
		if (id == null) {
			throw new DatastoreDataException("The given key doesn't have a numeric ID but a conversion" +
					" to Long was attempted: " + key);
		}
	}
	return id;
}
 

private void deleteShardsForCounter(String counterId) throws IOException {
  final List<Key> shards = new ArrayList<>();
  datastore.query(EntityQuery.newEntityQueryBuilder()
      .setKind(KIND_COUNTER_SHARD)
      .setFilter(PropertyFilter.eq(PROPERTY_COUNTER_ID, counterId))
      .build(), entity -> shards.add(entity.getKey()));

  // this is a safe guard to not to exceed max number of entities in one batch write
  // because in practice number of shards is much smaller
  for (List<Key> batch : Lists.partition(shards, MAX_NUMBER_OF_ENTITIES_IN_ONE_BATCH_WRITE)) {
      datastore.delete(batch.toArray(new Key[0]));
  }
}
 

@Override
public Backfill store(Backfill backfill) throws IOException {
  final Key key = DatastoreStorage.backfillKey(tx.getDatastore().newKeyFactory(), backfill.id());
  Entity.Builder builder = Entity.newBuilder(key)
      .set(PROPERTY_CONCURRENCY, backfill.concurrency())
      .set(PROPERTY_START, instantToTimestamp(backfill.start()))
      .set(PROPERTY_END, instantToTimestamp(backfill.end()))
      .set(PROPERTY_COMPONENT, backfill.workflowId().componentId())
      .set(PROPERTY_WORKFLOW, backfill.workflowId().id())
      .set(PROPERTY_SCHEDULE, backfill.schedule().toString())
      .set(PROPERTY_NEXT_TRIGGER, instantToTimestamp(backfill.nextTrigger()))
      .set(PROPERTY_ALL_TRIGGERED, backfill.allTriggered())
      .set(PROPERTY_HALTED, backfill.halted())
      .set(PROPERTY_REVERSE, backfill.reverse());

  backfill.created().ifPresent(x -> builder.set(PROPERTY_CREATED,
      instantToTimestamp(x)));
  backfill.lastModified().ifPresent(x -> builder.set(PROPERTY_LAST_MODIFIED,
      instantToTimestamp(x)));
  backfill.description().ifPresent(x -> builder.set(PROPERTY_DESCRIPTION, StringValue
      .newBuilder(x).setExcludeFromIndexes(true).build()));

  if (backfill.triggerParameters().isPresent()) {
    final String json = OBJECT_MAPPER.writeValueAsString(backfill.triggerParameters().get());
    builder.set(PROPERTY_TRIGGER_PARAMETERS,
        StringValue.newBuilder(json).setExcludeFromIndexes(true).build());
  }

  tx.put(builder.build());

  return backfill;
}
 

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

/**
 * 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 testBatchLookup() {
  Key taskKey1 = keyFactory.newKey(1);
  Key taskKey2 = keyFactory.newKey(2);
  List<Entity> expectedTasks = setUpBatchTests(taskKey1, taskKey2);
  // [START datastore_batch_lookup]
  Iterator<Entity> tasks = datastore.get(taskKey1, taskKey2);
  // [END datastore_batch_lookup]
  assertEquals(expectedTasks.get(0), tasks.next());
  assertEquals(expectedTasks.get(1), tasks.next());
}
 

/** Example of deleting multiple entities. */
// [TARGET delete(Key...)]
// [VARIABLE "my_key_name1"]
// [VARIABLE "my_key_name2"]
public void multipleDeleteEntities(String keyName1, String keyName2) {
  Datastore datastore = transaction.getDatastore();
  // [START multipleDeleteEntities]
  Key key1 = datastore.newKeyFactory().setKind("MyKind").newKey(keyName1);
  Key key2 = datastore.newKeyFactory().setKind("MyKind").newKey(keyName2);
  transaction.delete(key1, key2);
  transaction.commit();
  // [END multipleDeleteEntities]
}
 

/**
 * Executes the given {@link KeyQueryRequest} and returns the response.
 * 
 * @param request
 *          the key query request
 * @return the query response
 */
public QueryResponse<DatastoreKey> executeKeyQueryRequest(KeyQueryRequest request) {
  try {
    GqlQuery.Builder<Key> queryBuilder = Query.newGqlQueryBuilder(ResultType.KEY,
        request.getQuery());
    queryBuilder.setNamespace(entityManager.getEffectiveNamespace());
    queryBuilder.setAllowLiteral(request.isAllowLiterals());
    QueryUtils.applyNamedBindings(queryBuilder, request.getNamedBindings());
    QueryUtils.applyPositionalBindings(queryBuilder, request.getPositionalBindings());
    GqlQuery<Key> gqlQuery = queryBuilder.build();
    QueryResults<Key> results = nativeReader.run(gqlQuery);
    List<DatastoreKey> entities = new ArrayList<>();
    DefaultQueryResponse<DatastoreKey> response = new DefaultQueryResponse<>();
    response.setStartCursor(new DefaultDatastoreCursor(results.getCursorAfter().toUrlSafe()));
    while (results.hasNext()) {
      Key result = results.next();
      DatastoreKey datastoreKey = new DefaultDatastoreKey(result);
      entities.add(datastoreKey);
    }
    response.setResults(entities);
    response.setEndCursor(new DefaultDatastoreCursor(results.getCursorAfter().toUrlSafe()));
    response.setQueryResponseMetadata(
        new DefaultQueryResponseMetadata(
            QueryResponseMetadata.QueryState.forMoreResultsType(results.getMoreResults())));
    return response;
  } catch (DatastoreException exp) {
    throw new EntityManagerException(exp);
  }
}
 

@Override
public <T> Iterable<T> saveAll(Iterable<T> entities, Key... ancestors) {
	List<T> instances;
	if (entities instanceof List) {
		instances = (List<T>) entities;
	}
	else {
		instances = new ArrayList<>();
		entities.forEach(instances::add);
	}
	saveEntities(instances, ancestors);
	return entities;
}
 

@Test
public void testKeysOnlyQuery() {
  setUpQueryTests();
  // [START datastore_keys_only_query]
  Query<Key> query = Query.newKeyQueryBuilder().setKind("Task").build();
  // [END datastore_keys_only_query]
  assertValidQuery(query);
}