com.google.cloud.datastore.FullEntity Java Examples

@Test
public void saveTestNullDescendantsAndReferences() {
	// making sure save works when descendants are null
	assertThat(this.ob2.childEntities).isNull();
	assertThat(this.ob2.singularReference).isNull();
	assertThat(this.ob2.multipleReference).isNull();

	List<Object[]> callsArgs = gatherVarArgCallsArgs(this.datastore.put(ArgumentMatchers.<FullEntity[]>any()),
			Collections.singletonList(this.e1));

	this.datastoreTemplate.save(this.ob2);

	assertArgs(callsArgs, new MapBuilder<List, Integer>()
			.put(Collections.singletonList(Entity.newBuilder(this.key2).build()), 1).buildModifiable());

}
 

@Test
public void testBatchUpsert() {
  // [START datastore_batch_upsert]
  FullEntity<IncompleteKey> task1 = FullEntity.newBuilder(keyFactory.newKey())
      .set("category", "Personal")
      .set("done", false)
      .set("priority", 4)
      .set("description", "Learn Cloud Datastore")
      .build();
  FullEntity<IncompleteKey> task2 = Entity.newBuilder(keyFactory.newKey())
      .set("category", "Personal")
      .set("done", false)
      .set("priority", 5)
      .set("description", "Integrate Cloud Datastore")
      .build();
  List<Entity> tasks = datastore.add(task1, task2);
  Key taskKey1 = tasks.get(0).getKey();
  Key taskKey2 = tasks.get(1).getKey();
  // [END datastore_batch_upsert]
  assertEquals(Entity.newBuilder(taskKey1, task1).build(), datastore.get(taskKey1));
  assertEquals(Entity.newBuilder(taskKey2, task2).build(), datastore.get(taskKey2));
}
 

@Test
public void performTransactionTest() {

	DatastoreReaderWriter transactionContext = mock(DatastoreReaderWriter.class);

	when(this.datastore.runInTransaction(any())).thenAnswer((invocation) -> {
		TransactionCallable<String> callable = invocation.getArgument(0);
		return callable.run(transactionContext);
	});

	List<Entity> e1 = Collections
			.singletonList(this.e1);
	when(transactionContext.fetch(ArgumentMatchers.<Key[]>any())).thenReturn(e1);

	String finalResult = this.datastoreTemplate
			.performTransaction((datastoreOperations) -> {
				datastoreOperations.save(this.ob2);
				datastoreOperations.findById("ignored", TestEntity.class);
				return "all done";
			});

	assertThat(finalResult).isEqualTo("all done");
	verify(transactionContext, times(1)).put(ArgumentMatchers.<FullEntity[]>any());
	verify(transactionContext, times(2)).fetch((Key[]) any());
}
 

private Value setExcludeFromIndexes(Value convertedVal) {
	// ListValues must have its contents individually excluded instead.
	// the entire list must NOT be excluded or there will be an exception.
	// Same for maps and embedded entities which are stored as EntityValue.
	if (convertedVal.getClass().equals(EntityValue.class)) {
		FullEntity.Builder<IncompleteKey> builder = FullEntity.newBuilder();
		((EntityValue) convertedVal).get().getProperties()
						.forEach((key, value) -> builder.set(key, setExcludeFromIndexes(value)));
		return EntityValue.of(builder.build());
	}
	else if (convertedVal.getClass().equals(ListValue.class)) {
		return ListValue.of((List) ((ListValue) convertedVal).get().stream()
						.map(this::setExcludeFromIndexes).collect(Collectors.toList()));
	}
	else {
		return convertedVal.toBuilder().setExcludeFromIndexes(true).build();
	}
}
 

@Test
public void saveReferenceLoopTest() {
	ReferenceTestEntity referenceTestEntity = new ReferenceTestEntity();
	referenceTestEntity.id = 1L;
	referenceTestEntity.sibling = referenceTestEntity;
	when(this.objectToKeyFactory.getKeyFromObject(eq(referenceTestEntity), any())).thenReturn(this.key1);

	List<Object[]> callsArgs = gatherVarArgCallsArgs(this.datastore.put(ArgumentMatchers.<FullEntity[]>any()),
			Collections.singletonList(this.e1));

	assertThat(this.datastoreTemplate.save(referenceTestEntity)).isInstanceOf(ReferenceTestEntity.class);

	Entity writtenEntity = Entity.newBuilder(this.key1)
			.set("sibling", this.key1)
			.build();

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

public void save() {
  if (key == null) {
    key = getDatastore().allocateId(makeIncompleteKey()); // Give this greeting a unique ID
  }

  Builder<Key> builder = FullEntity.newBuilder(key);

  if (authorEmail != null) {
    builder.set("authorEmail", authorEmail);
  }

  if (authorId != null) {
    builder.set("authorId", authorId);
  }

  builder.set("content", content);
  builder.set("date", Timestamp.of(date));

  getDatastore().put(builder.build());
}
 

private EntityValue applyEntityValueBuilder(Object val, String kindName,
		Consumer<Builder> consumer, boolean createKey) {

	FullEntity.Builder<IncompleteKey> builder;
	if (!createKey) {
		builder = FullEntity.newBuilder();
	}
	else {
		/* The following does 3 sequential null checks. We only want an ID value if the object isn't null,
			has an ID property, and the ID property isn't null.
		* */
		Optional idProp = Optional.ofNullable(val)
				.map(v -> this.datastoreMappingContext.getPersistentEntity(v.getClass()))
				.map(datastorePersistentEntity -> datastorePersistentEntity.getIdProperty())
				.map(id -> this.datastoreMappingContext.getPersistentEntity(val.getClass())
						.getPropertyAccessor(val).getProperty(id));

		IncompleteKey key = idProp.isPresent() ? this.objectToKeyFactory.getKeyFromId(idProp.get(), kindName)
				: this.objectToKeyFactory.getIncompleteKey(kindName);
		builder = FullEntity.newBuilder(key);
	}
	consumer.accept(builder);
	return EntityValue.of(builder.build());
}
 

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

@SuppressWarnings("unchecked")
@Override
public Object toModel(Value<?> input) {
  if (input instanceof NullValue) {
    return null;
  }
  EntityValue entityValue = (EntityValue) input;
  FullEntity<?> entity = entityValue.get();
  Map<String, Object> map;
  if (Modifier.isAbstract(mapClass.getModifiers())) {
    if (SortedMap.class.equals(mapClass)) {
      map = new TreeMap<>();
    } else {
      map = new HashMap<>();
    }
  } else {
    map = (Map<String, Object>) IntrospectionUtils.instantiateObject(mapClass);
  }
  for (String property : entity.getNames()) {
    map.put(property, valueMapper.toModel(entity.getValue(property)));
  }
  return map;
}
 

@Override
public Object toModel(Value<?> input) {
  if (input.getType() == ValueType.NULL) {
    return null;
  }
  try {
    FullEntity<?> entity = ((EntityValue) input).get();
    ConstructorMetadata constructorMetadata = metadata.getConstructorMetadata();
    Object embeddedObject = constructorMetadata.getConstructorMethodHandle().invoke();
    for (PropertyMetadata propertyMetadata : metadata.getPropertyMetadataCollection()) {
      String mappedName = propertyMetadata.getMappedName();
      if (entity.contains(mappedName)) {
        Value<?> propertyValue = entity.getValue(mappedName);
        Object fieldValue = propertyMetadata.getMapper().toModel(propertyValue);
        propertyMetadata.getWriteMethod().invoke(embeddedObject, fieldValue);
      }
    }
    if (constructorMetadata.isBuilderConstructionStrategy()) {
      embeddedObject = metadata.getConstructorMetadata().getBuildMethodHandle()
          .invoke(embeddedObject);
    }
    return embeddedObject;
  } catch (Throwable exp) {
    throw new MappingException(exp);
  }
}
 

/** Example of putting multiple entities with deferred id allocation. */
// [TARGET putWithDeferredIdAllocation(FullEntity...)]
public List<Key> multiplePutEntitiesDeferredId() {
  Datastore datastore = transaction.getDatastore();
  // [START multiplePutEntitiesDeferredId]
  IncompleteKey key1 = datastore.newKeyFactory().setKind("MyKind").newKey();
  FullEntity.Builder entityBuilder1 = FullEntity.newBuilder(key1);
  entityBuilder1.set("propertyName", "value1");
  FullEntity entity1 = entityBuilder1.build();

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

  transaction.putWithDeferredIdAllocation(entity1, entity2);
  Response response = transaction.commit();
  // [END multiplePutEntitiesDeferredId]
  return response.getGeneratedKeys();
}
 

@Test
public void testSessionSave() throws Exception {
  DatastoreSession session = spy(new DatastoreSession(manager));
  session.setValid(true);
  session.setId(keyId);
  session.setAttribute("count", 5);

  store.save(session);
  ArgumentCaptor captor = ArgumentCaptor.forClass(Entity.class);
  verify(datastore).put((FullEntity<?>[]) captor.capture());
  verify(session).saveAttributesToEntity(any());

  List<Entity> entities = captor.getAllValues();
  assertEquals(2, entities.size());

  assertTrue(entities.stream()
      .map(e -> e.getKey().getName())
      .collect(Collectors.toList())
      .containsAll(Arrays.asList("count", keyId)));

}
 

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

/**
 * Marshals the given entity and and returns the equivalent Entity needed for the underlying Cloud
 * Datastore API.
 * 
 * @return A native entity that is equivalent to the POJO being marshalled. The returned value
 *         could either be a FullEntity or Entity.
 */
private BaseEntity<?> marshal() {
  marshalKey();
  if (key instanceof Key) {
    entityBuilder = Entity.newBuilder((Key) key);
  } else {
    entityBuilder = FullEntity.newBuilder(key);
  }
  marshalFields();
  marshalAutoTimestampFields();
  if (intent == Intent.UPDATE) {
    marshalVersionField();
  }
  marshalEmbeddedFields();
  return entityBuilder.build();
}
 

/**
 * Inserts the given list of entities into the Cloud Datastore.
 * 
 * @param entities
 *          the entities to insert.
 * @return the inserted entities. The inserted entities will not be same as the passed in
 *         entities. For example, the inserted entities may contain generated ID, key, parent key,
 *         etc.
 * @throws EntityManagerException
 *           if any error occurs while inserting.
 */
@SuppressWarnings("unchecked")
public <E> List<E> insert(List<E> entities) {
  if (entities == null || entities.isEmpty()) {
    return new ArrayList<>();
  }
  try {
    entityManager.executeEntityListeners(CallbackType.PRE_INSERT, entities);
    FullEntity<?>[] nativeEntities = toNativeFullEntities(entities, entityManager, Intent.INSERT);
    Class<?> entityClass = entities.get(0).getClass();
    List<Entity> insertedNativeEntities = nativeWriter.add(nativeEntities);
    List<E> insertedEntities = (List<E>) toEntities(entityClass, insertedNativeEntities);
    entityManager.executeEntityListeners(CallbackType.POST_INSERT, insertedEntities);
    return insertedEntities;
  } catch (DatastoreException exp) {
    throw DatastoreUtils.wrap(exp);
  }
}
 

/**
 * Updates or inserts the given list of entities in the Cloud Datastore. If the entities do not
 * have a valid ID, IDs may be generated.
 * 
 * @param entities
 *          the entities to update/or insert.
 * @return the updated or inserted entities
 * @throws EntityManagerException
 *           if any error occurs while saving.
 */
@SuppressWarnings("unchecked")
public <E> List<E> upsert(List<E> entities) {
  if (entities == null || entities.isEmpty()) {
    return new ArrayList<>();
  }
  try {
    entityManager.executeEntityListeners(CallbackType.PRE_UPSERT, entities);
    FullEntity<?>[] nativeEntities = toNativeFullEntities(entities, entityManager, Intent.UPSERT);
    Class<?> entityClass = entities.get(0).getClass();
    List<Entity> upsertedNativeEntities = nativeWriter.put(nativeEntities);
    List<E> upsertedEntities = (List<E>) toEntities(entityClass, upsertedNativeEntities);
    entityManager.executeEntityListeners(CallbackType.POST_UPSERT, upsertedEntities);
    return upsertedEntities;
  } catch (DatastoreException exp) {
    throw DatastoreUtils.wrap(exp);
  }
}
 

@Override
public Long createBook(Book book) {
  IncompleteKey key = keyFactory.newKey();          // Key will be assigned once written
  FullEntity<IncompleteKey> incBookEntity = Entity.newBuilder(key)  // Create the Entity
      .set(Book.AUTHOR, book.getAuthor())           // Add Property ("author", book.getAuthor())
      .set(Book.DESCRIPTION, book.getDescription())
      .set(Book.PUBLISHED_DATE, book.getPublishedDate())
      .set(Book.TITLE, book.getTitle())
      .set(Book.IMAGE_URL, book.getImageUrl())
      .set(Book.CREATED_BY, book.getCreatedBy())
      .set(Book.CREATED_BY_ID, book.getCreatedById())
      .build();
  Entity bookEntity = datastore.add(incBookEntity); // Save the Entity
  return bookEntity.getKey().getId();                     // The ID of the Key
}
 

@Override
public Long createBook(Book book) {
  IncompleteKey key = keyFactory.newKey();          // Key will be assigned once written
  FullEntity<IncompleteKey> incBookEntity = Entity.newBuilder(key)  // Create the Entity
      .set(Book.AUTHOR, book.getAuthor())           // Add Property ("author", book.getAuthor())
      .set(Book.DESCRIPTION, book.getDescription())
      .set(Book.PUBLISHED_DATE, book.getPublishedDate())
      .set(Book.TITLE, book.getTitle())
      .set(Book.IMAGE_URL, book.getImageUrl())
      .build();
  Entity bookEntity = datastore.add(incBookEntity); // Save the Entity
  return bookEntity.getKey().getId();                     // The ID of the Key
}
 

@Override
default List<Entity> add(FullEntity<?>... entities) {
  for (FullEntity<?> entity : entities) {
    stats().recordDatastoreEntityWrites(entity.getKey().getKind(), 1);
  }
  return writer().add(entities);
}
 

@Override
public Long createBook(Book book) {
  IncompleteKey key = keyFactory.newKey();          // Key will be assigned once written
  FullEntity<IncompleteKey> incBookEntity = Entity.newBuilder(key)  // Create the Entity
      .set(Book.AUTHOR, book.getAuthor())           // Add Property ("author", book.getAuthor())
      .set(Book.DESCRIPTION, book.getDescription())
      .set(Book.PUBLISHED_DATE, book.getPublishedDate())
      .set(Book.TITLE, book.getTitle())
      .set(Book.IMAGE_URL, book.getImageUrl())
      .set(Book.CREATED_BY, book.getCreatedBy())
      .set(Book.CREATED_BY_ID, book.getCreatedById())
      .build();
  Entity bookEntity = datastore.add(incBookEntity); // Save the Entity
  return bookEntity.getKey().getId();                     // The ID of the Key
}
 

@Override
default List<Entity> put(FullEntity<?>... entities) {
  for (FullEntity<?> entity : entities) {
    stats().recordDatastoreEntityWrites(entity.getKey().getKind(), 1);
  }
  return writer().put(entities);
}
 

@Override
public Long createBook(Book book) {
  IncompleteKey key = keyFactory.newKey();          // Key will be assigned once written
  FullEntity<IncompleteKey> incBookEntity = Entity.newBuilder(key)  // Create the Entity
      .set(Book.AUTHOR, book.getAuthor())           // Add Property ("author", book.getAuthor())
      .set(Book.DESCRIPTION, book.getDescription())
      .set(Book.PUBLISHED_DATE, book.getPublishedDate())
      .set(Book.TITLE, book.getTitle())
      .build();
  Entity bookEntity = datastore.add(incBookEntity); // Save the Entity
  return bookEntity.getKey().getId();                     // The ID of the Key
}
 

@Test
public void testMarshal_WrappedLongId2() {
  WrappedLongIdEntity entity = WrappedLongIdEntity.getSample2();
  FullEntity<?> nativeEntity = (FullEntity<?>) Marshaller.marshal(em, entity, Intent.INSERT);
  IncompleteKey incompleteKey = nativeEntity.getKey();
  assertNotNull(incompleteKey);
}
 

@Test
public void testMarshal_WrappedLongId1() {
  WrappedLongIdEntity entity = WrappedLongIdEntity.getSample1();
  FullEntity<?> nativeEntity = (FullEntity<?>) Marshaller.marshal(em, entity, Intent.INSERT);
  IncompleteKey incompleteKey = nativeEntity.getKey();
  assertNotNull(incompleteKey);
}
 

@Test
public void testMarshal_Embedded_Imploded() {
  Contact contact = Contact.createContact1();
  FullEntity<?> entity = (FullEntity<?>) Marshaller.marshal(em, contact, Intent.INSERT);
  assertNull(entity.getValue("cellNumber").get());
  assertNull(entity.getValue("homeAddress").get());
}
 

@Test
public void testMarshal_Embedded_Imploded2() {
  Contact contact = Contact.createContact2();
  FullEntity<?> entity = (FullEntity<?>) Marshaller.marshal(em, contact, Intent.INSERT);
  assertEquals("55555",
      entity.getEntity("homeAddress").getEntity("postal_code").getString("zip"));
}
 

@Test
public void testMarshal_Embedded() {
  Customer customer = Customer.createSampleCustomer2();
  FullEntity<?> entity = (FullEntity<?>) Marshaller.marshal(em, customer, Intent.INSERT);
  assertNull(entity.getString("ba_line1"));
  assertNull(entity.getString("ba_line2"));
  assertNull(entity.getString("ba_zip"));
  assertNull(entity.getString("ba_zipx"));
  assertNull(entity.getString("ba_zip"));
}
 

/**
 * Unmarshals the embedded field represented by the given metadata.
 * 
 * @param embeddedMetadata
 *          the metadata of the field to unmarshal
 * @param target
 *          the object in which the embedded field is declared/accessible from
 * @param nativeEntity
 *          the native entity from which the embedded entity is to be extracted
 * @throws Throwable
 *           propagated
 */
private static void unmarshalWithImplodedStrategy(EmbeddedMetadata embeddedMetadata,
    Object target, BaseEntity<?> nativeEntity) throws Throwable {
  Object embeddedObject = null;
  ConstructorMetadata constructorMetadata = embeddedMetadata.getConstructorMetadata();
  FullEntity<?> nativeEmbeddedEntity = null;
  String propertyName = embeddedMetadata.getMappedName();
  if (nativeEntity.contains(propertyName)) {
    Value<?> nativeValue = nativeEntity.getValue(propertyName);
    if (nativeValue instanceof NullValue) {
      embeddedMetadata.getWriteMethod().invoke(target, embeddedObject);
    } else {
      nativeEmbeddedEntity = ((EntityValue) nativeValue).get();
      embeddedObject = constructorMetadata.getConstructorMethodHandle().invoke();
    }
  }
  if (embeddedObject == null) {
    return;
  }
  for (PropertyMetadata propertyMetadata : embeddedMetadata.getPropertyMetadataCollection()) {
    unmarshalProperty(propertyMetadata, embeddedObject, nativeEmbeddedEntity);
  }
  for (EmbeddedMetadata embeddedMetadata2 : embeddedMetadata.getEmbeddedMetadataCollection()) {
    unmarshalWithImplodedStrategy(embeddedMetadata2, embeddedObject, nativeEmbeddedEntity);
  }
  if (constructorMetadata.isBuilderConstructionStrategy()) {
    embeddedObject = constructorMetadata.getBuildMethodHandle().invoke(embeddedObject);
  }
  embeddedMetadata.getWriteMethod().invoke(target, embeddedObject);
}
 

@Override
public <E> void upsertWithDeferredIdAllocation(List<E> entities) {
  if (entities == null || entities.isEmpty()) {
    return;
  }
  try {
    DatastoreUtils.validateDeferredIdAllocation(entities.get(0));
    FullEntity<?>[] nativeEntities = toNativeFullEntities(entities, entityManager, Intent.UPSERT);
    nativeTransaction.putWithDeferredIdAllocation(nativeEntities);
  } catch (DatastoreException exp) {
    throw DatastoreUtils.wrap(exp);
  }
}