com.google.cloud.datastore.PathElement Java Examples

private Entity getShardCounter(final String location) {
  KeyFactory keyFactory = datastore.newKeyFactory().addAncestors(
      NXRM_ROOT,
      PathElement.of(METRICS_STORE, 1L)
  );
  Key key = keyFactory.setNamespace(namespace)
      .setKind(SHARD).newKey(location);
  Entity exists = datastore.get(key);
  if (exists != null) {
    log.trace("counter for {} already present", location);
    return exists;
  }

  log.debug("creating metrics store counter shard for {}", location);
  // otherwise make it
  Entity entity = Entity.newBuilder(key)
      .set(SIZE, LongValue.newBuilder(0L).build())
      .set(COUNT, LongValue.newBuilder(0L).build())
      .build();

  return datastore.put(entity);
}
 

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

/**
 * Save the specified Session into this Store. Any previously saved information for
 * the associated session identifier is replaced.
 *
 * <p>Attempt to serialize the session and send it to the datastore.</p>
 *
 * @throws IOException If an error occurs during the serialization of the session.
 *
 * @param session Session to be saved
 */
@Override
public void save(Session session) throws IOException {
  log.debug("Persisting session: " + session.getId());

  if (!(session instanceof DatastoreSession)) {
    throw new IOException(
        "The session must be an instance of DatastoreSession to be serialized");
  }
  DatastoreSession datastoreSession = (DatastoreSession) session;
  Key sessionKey = newKey(session.getId());
  KeyFactory attributeKeyFactory = datastore.newKeyFactory()
      .setKind(sessionKind)
      .addAncestor(PathElement.of(sessionKind, sessionKey.getName()));

  List<Entity> entities = serializeSession(datastoreSession, sessionKey, attributeKeyFactory);

  TraceContext datastoreSaveContext = startSpan("Storing the session in the Datastore");
  datastore.put(entities.toArray(new FullEntity[0]));
  datastore.delete(datastoreSession.getSuppressedAttributes().stream()
      .map(attributeKeyFactory::newKey)
      .toArray(Key[]::new));
  endSpan(datastoreSaveContext);
}
 

@Test
public void saveTestKeyWithAncestor() {
	Key key0 = createFakeKey("key0");
	Key keyA = Key.newBuilder(key0)
			.addAncestor(PathElement.of(key0.getKind(), key0.getName())).setName("keyA").build();
	ChildEntity childEntity = new ChildEntity();
	childEntity.id = keyA;
	when(this.objectToKeyFactory.getKeyFromObject(eq(childEntity), any())).thenReturn(keyA);
	List<Object[]> callsArgs = gatherVarArgCallsArgs(this.datastore.put(ArgumentMatchers.<FullEntity[]>any()),
			Collections.singletonList(this.e1));

	this.datastoreTemplate.save(childEntity, key0);

	Entity writtenChildEntity = Entity.newBuilder(keyA).build();

	assertArgs(callsArgs, new MapBuilder<List, Integer>()
			.put(Collections.singletonList(writtenChildEntity), 1)
			.buildModifiable());
}
 

private void setUpQueryTests() {
  Key taskKey = datastore.newKeyFactory()
      .setKind("Task")
      .addAncestors(PathElement.of("TaskList", "default"))
      .newKey("someTask");
  datastore.put(Entity.newBuilder(taskKey)
      .set("category", "Personal")
      .set("done", false)
      .set("completed", false)
      .set("priority", 4)
      .set("created", includedDate)
      .set("percent_complete", 10.0)
      .set("description",
          StringValue.newBuilder("Learn Cloud Datastore").setExcludeFromIndexes(true).build())
      .set("tag", "fun", "l", "programming")
      .build());
}
 

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

private Key getErrorKey(UUID jobId, String errorId) {
  // Use the main job as the ancestor to all the errors, see:
  // https://cloud.google.com/datastore/docs/concepts/entities#ancestor_paths
  return datastore
      .newKeyFactory()
      .setKind(ERROR_KIND)
      .addAncestor(PathElement.of(JOB_KIND, jobId.toString()))
      .newKey(errorId);
}
 

@Test
public void testSerializationCycleWithAttributeRemoval() throws Exception {
  DatastoreSession initialSession = new DatastoreSession(manager);
  initialSession.setValid(true);
  initialSession.setId(keyId);
  initialSession.setAttribute("count", 5);
  initialSession.setAttribute("map", Collections.singletonMap("key", "value"));
  KeyFactory attributeKeyFactory = datastore.newKeyFactory()
      .setKind("kind")
      .addAncestor(PathElement.of("kind", key.getName()));

  List<Entity> initialSessionEntities = store.serializeSession(initialSession, key,
      attributeKeyFactory);

  // Load the session and remove the map attribute
  when(datastore.<Entity>run(any())).thenReturn(
      new IteratorQueryResults<>(initialSessionEntities.iterator()));
  DatastoreSession session = (DatastoreSession)store.load(keyId);
  session.getSession().setAttribute("map", null);

  // Save and reload the session to ensure that the attribute map is not serialized
  store.save(session);

  ArgumentCaptor<Key> keyCaptors = ArgumentCaptor.forClass(Key.class);
  verify(datastore).delete(keyCaptors.capture());

  assertNotNull(keyCaptors.getValue());
  assertEquals("map", keyCaptors.getValue().getName());
}
 

@Test
public void testDecomposedSessionLoad() throws Exception {
  DatastoreSession session = new DatastoreSession(manager);
  session.setValid(true);
  session.setId(keyId);
  session.setAttribute("count", 2);
  session.setAttribute("map", Collections.singletonMap("key", "value"));

  KeyFactory attributeKeyFactory = datastore.newKeyFactory()
      .setKind("kind")
      .addAncestor(PathElement.of("kind", key.getName()));
  List<Entity> entities = session.saveToEntities(key, attributeKeyFactory);

  QueryResults<Entity> queryResults = new IteratorQueryResults<>(entities.iterator());
  when(datastore.<Entity>run(any())).thenReturn(queryResults);

  Session restored = store.load(keyId);

  assertEquals(keyId, restored.getId());
  assertEquals(2, restored.getSession().getAttribute("count"));
  assertEquals("value",
      ((Map<String, String>)session.getSession().getAttribute("map")).get("key"));
}
 

@Test
public void testPropertyFilteringRunQuery() {
  setUpQueryTests();
  // [START datastore_property_filtering_run_query]
  Key startKey = datastore.newKeyFactory()
      .setKind("__property__")
      .addAncestors(PathElement.of("__kind__", "Task"))
      .newKey("priority");
  Query<Key> query = Query.newKeyQueryBuilder()
      .setKind("__property__")
      .setFilter(PropertyFilter.ge("__key__", startKey))
      .build();
  Map<String, Collection<String>> propertiesByKind = new HashMap<>();
  QueryResults<Key> keys = datastore.run(query);
  while (keys.hasNext()) {
    Key key = keys.next();
    String kind = key.getParent().getName();
    String propertyName = key.getName();
    Collection<String> properties = propertiesByKind.get(kind);
    if (properties == null) {
      properties = new HashSet<String>();
      propertiesByKind.put(kind, properties);
    }
    properties.add(propertyName);
  }
  // [END datastore_property_filtering_run_query]
  Map<String, ImmutableSet<String>> expected =
      ImmutableMap.of("Task", ImmutableSet.of("priority", "tag"));
  assertEquals(expected, propertiesByKind);
}
 

@Test
public void testKeyWithMultilevelParent() {
  // [START datastore_key_with_multilevel_parent]
  KeyFactory keyFactory = datastore.newKeyFactory()
      .addAncestors(PathElement.of("User", "Alice"), PathElement.of("TaskList", "default"))
      .setKind("Task");
  Key taskKey = keyFactory.newKey("sampleTask");
  // [END datastore_key_with_multilevel_parent]
  assertValidKey(taskKey);
}
 

@Test
public void testKeyWithParent() {
  // [START datastore_key_with_parent]
  Key taskKey = datastore.newKeyFactory()
      .addAncestors(PathElement.of("TaskList", "default"))
      .setKind("Task")
      .newKey("sampleTask");
  // [END datastore_key_with_parent]
  assertValidKey(taskKey);
}
 

@Test
public void testRemoveAncestors_IdKeys() {
	Key idKey = Key.newBuilder("project", "person", 46L)
			.addAncestor(PathElement.of("person", 22L))
			.addAncestor(PathElement.of("person", 18L))
			.build();

	Key processedKey = KeyUtil.getKeyWithoutAncestors(idKey);
	assertThat(processedKey.getAncestors()).isEmpty();
}
 

@Test
public void testRemoveAncestors_NamedKeys() {
	Key namedKey = Key.newBuilder("project", "person", "Smith")
			.addAncestor(PathElement.of("person", "GrandParent"))
			.addAncestor(PathElement.of("person", "Parent"))
			.build();

	Key processedKey = KeyUtil.getKeyWithoutAncestors(namedKey);
	assertThat(processedKey.getAncestors()).isEmpty();
}
 

@Test
public void allocateIdForObjectTest() {
	TestEntityWithKeyId testEntityWithKeyId = new TestEntityWithKeyId();

	doAnswer((invocation) -> {
		IncompleteKey incompleteKey = (IncompleteKey) invocation.getArguments()[0];
		long id = 123L;
		if (incompleteKey.getAncestors().size() > 0) {
			id = 456L;
		}
		return Key.newBuilder(incompleteKey, id).build();
	}).when(this.datastore).allocateId((IncompleteKey) any());

	when(this.datastore.newKeyFactory()).thenReturn(new KeyFactory("project"));
	Key allocatedKey = this.datastoreServiceObjectToKeyFactory
			.allocateKeyForObject(testEntityWithKeyId, this.datastoreMappingContext
					.getPersistentEntity(testEntityWithKeyId.getClass()));
	Key key = new KeyFactory("project").setKind("custom_test_kind").newKey(123L);
	assertThat(allocatedKey).isEqualTo(key);
	assertThat(testEntityWithKeyId.id).isEqualTo(key);

	Key allocatedKeyWithAncestor = this.datastoreServiceObjectToKeyFactory
			.allocateKeyForObject(testEntityWithKeyId, this.datastoreMappingContext
					.getPersistentEntity(testEntityWithKeyId.getClass()), allocatedKey);
	Key keyWithAncestor = new KeyFactory("project").setKind("custom_test_kind")
			.addAncestor(PathElement.of(key.getKind(), key.getId()))
			.newKey(456L);
	assertThat(allocatedKeyWithAncestor).isEqualTo(keyWithAncestor);
	assertThat(testEntityWithKeyId.id).isEqualTo(keyWithAncestor);
}
 

private void validateKey(Object entity, PathElement ancestorPE) {
	DatastorePersistentEntity datastorePersistentEntity =
			this.datastoreMappingContext.getPersistentEntity(entity.getClass());
	DatastorePersistentProperty idProp = datastorePersistentEntity.getIdPropertyOrFail();

	if (!TypeUtils.isAssignable(BaseKey.class, idProp.getType())) {
		throw new DatastoreDataException("Only Key types are allowed for descendants id");
	}

	Key key = getKey(entity, false);
	if (key == null || key.getAncestors().stream().anyMatch((pe) -> pe.equals(ancestorPE))) {
		return;
	}
	throw new DatastoreDataException("Descendant object has a key without current ancestor");
}
 

private Key getCountsKey(UUID jobId, String dataType) {
  // Use the main job as the ancestor to all the counts, see:
  // https://cloud.google.com/datastore/docs/concepts/entities#ancestor_paths
  return datastore
      .newKeyFactory()
      .setKind(COUNTS_KIND)
      .addAncestor(PathElement.of(JOB_KIND, jobId.toString()))
      .newKey(dataType);
}
 

private static Key activeWorkflowInstanceIndexShardEntryKey(KeyFactory keyFactory, String workflowInstanceKey) {
  return keyFactory.setKind(KIND_ACTIVE_WORKFLOW_INSTANCE_INDEX_SHARD_ENTRY)
      .addAncestor(PathElement.of(KIND_ACTIVE_WORKFLOW_INSTANCE_INDEX_SHARD,
          activeWorkflowInstanceIndexShardName(workflowInstanceKey)))
      .newKey(workflowInstanceKey);
}
 

public static PathElement keyToPathElement(Key key) {
	Assert.notNull(key, "A non-null key is required");
	return (key.getName() != null)
			? PathElement.of(key.getKind(), key.getName())
			: PathElement.of(key.getKind(), key.getId());
}