com.google.cloud.datastore.BaseEntity Java Examples

@Test
public void testAttributesSerializationKey() throws Exception {
  DatastoreSession session = new DatastoreSession(sessionManager);
  session.setValid(true);
  session.setAttribute("count", 2);
  session.setAttribute("map", new HashMap<>());

  KeyFactory factory = new KeyFactory("project").setKind("kind");
  List<Entity> entities = session.saveAttributesToEntity(factory);

  assertTrue(entities.stream()
      .map(BaseEntity::getKey)
      .map(Key::getName)
      .collect(Collectors.toSet())
      .containsAll(Arrays.asList("count", "map")));
}
 

private <T extends BaseEntity<?>> void testInstrumentedQueryResults(QueryResults<T> instrumented,
    QueryResults<T> delegate) {
  when(delegate.hasNext()).thenReturn(true);
  when(delegate.next()).thenReturn((T) entity);
  when(delegate.getSkippedResults()).thenReturn(17);
  when(delegate.getMoreResults()).thenReturn(MORE_RESULTS_AFTER_CURSOR);
  assertThat(instrumented.hasNext(), is(true));
  verify(delegate).hasNext();
  assertThat(instrumented.next(), is(entity));
  verify(delegate).next();
  assertThat(instrumented.getSkippedResults(), is(17));
  assertThat(instrumented.getMoreResults(), is(MORE_RESULTS_AFTER_CURSOR));
  verify(delegate).getSkippedResults();
  verify(delegate).getMoreResults();
  verify(stats).recordDatastoreEntityReads(TEST_KIND, 1);
}
 

private GqlQuery<? extends BaseEntity> bindArgsToGqlQuery() {
	Builder builder = GqlQuery.newGqlQueryBuilder(this.finalGql);
	builder.setAllowLiteral(true);
	if (this.tagsOrdered.size() != this.params.size()) {
		throw new DatastoreDataException("Annotated GQL Query Method "
				+ GqlDatastoreQuery.this.queryMethod.getName() + " has " + this.tagsOrdered.size()
				+ " tags but a different number of parameter values: " + this.params.size());
	}
	for (int i = 0; i < this.tagsOrdered.size(); i++) {
		Object val = this.params.get(i);
		Object boundVal;
		if (val instanceof Cursor) {
			boundVal = val;
		}
		else if (ValueUtil.isCollectionLike(val.getClass())) {
			boundVal = convertCollectionParamToCompatibleArray((List) ValueUtil.toListIfArray(val));
		}
		else {
			boundVal = GqlDatastoreQuery.this.datastoreOperations.getDatastoreEntityConverter().getConversions()
					.convertOnWriteSingle(convertEntitiesToKeys(val)).get();
		}
		DatastoreNativeTypes.bindValueToGqlBuilder(builder, this.tagsOrdered.get(i), boundVal);
	}
	return builder.build();
}
 

private <T> T convertEntityResolveDescendantsAndReferences(Class<T> entityClass,
		DatastorePersistentEntity datastorePersistentEntity,
		BaseKey key, ReadContext context) {
	T convertedObject;
	if (context.converted(key)) {
		convertedObject = (T) context.getConvertedEntity(key);
	}
	else {
		BaseEntity readEntity = context.getReadEntity(key);
		convertedObject = this.datastoreEntityConverter.read(entityClass, readEntity);

		// the parent entity should be put into context BEFORE referenced and descendant entities
		// are being resolved to avoid infinite loops
		context.putConvertedEntity(key, convertedObject);

		//raw Datastore entity is no longer needed
		context.removeReadEntity(key);
		if (convertedObject != null) {
			resolveDescendantProperties(datastorePersistentEntity, readEntity, convertedObject, context);
			resolveReferenceProperties(datastorePersistentEntity, readEntity, convertedObject, context);
		}
	}

	return convertedObject;
}
 

private <T> void resolveReferenceProperties(DatastorePersistentEntity datastorePersistentEntity,
		BaseEntity entity, T convertedObject, ReadContext context) {
	datastorePersistentEntity.doWithAssociations(
			(AssociationHandler) (association) -> {
				DatastorePersistentProperty referenceProperty = (DatastorePersistentProperty) association
						.getInverse();
				String fieldName = referenceProperty.getFieldName();
				if (entity.contains(fieldName) && !entity.isNull(fieldName)) {
					Class<?> type = referenceProperty.getType();
					Object referenced = computeReferencedField(entity, context, referenceProperty, fieldName, type);
					if (referenced != null) {
						datastorePersistentEntity.getPropertyAccessor(convertedObject)
								.setProperty(referenceProperty, referenced);
					}
				}
			});
}
 

private <T> T computeReferencedField(BaseEntity entity, ReadContext context,
		DatastorePersistentProperty referenceProperty, String fieldName, Class<T> type) {
	T referenced;
	if (referenceProperty.isLazyLoaded()) {
		DatastoreReaderWriter originalTx = getDatastoreReadWriter();
		referenced = LazyUtil.wrapSimpleLazyProxy(() -> {
			if (getDatastoreReadWriter() != originalTx) {
				throw new DatastoreDataException("Lazy load should be invoked within the same transaction");
			}
			return (T) findReferenced(entity, referenceProperty, context);
		}, type, entity.getValue(fieldName));
	}
	else {
		referenced = (T) findReferenced(entity, referenceProperty, context);
	}
	return referenced;
}
 

private Object findReferenced(BaseEntity entity, DatastorePersistentProperty referencePersistentProperty,
		ReadContext context) {
	String fieldName = referencePersistentProperty.getFieldName();
	try {
		Object referenced;
		if (referencePersistentProperty.isCollectionLike()) {
			referenced = fetchReferenced(referencePersistentProperty, context,
					valuesToKeys(entity.getList(fieldName)));
		}
		else {
			List referencedList = findAllById(Collections.singleton(entity.getKey(fieldName)),
					referencePersistentProperty.getType(), context);
			referenced = referencedList.isEmpty() ? null : referencedList.get(0);
		}
		return referenced;
	}
	catch (ClassCastException ex) {
			throw new DatastoreDataException(
				"Error loading reference property " + fieldName + "."
						+ "Reference properties must be stored as Keys or lists of Keys"
						+ " in Cloud Datastore for singular or multiple references, respectively.");
		}
}
 

/**
 * Marshals the field with the given property metadata.
 * 
 * @param propertyMetadata
 *          the metadata of the field to be marshaled.
 * @param target
 *          the object in which the field is defined/accessible from
 * @param entityBuilder
 *          the native entity on which the marshaled field should be set
 */
private static void marshalField(PropertyMetadata propertyMetadata, Object target,
    BaseEntity.Builder<?, ?> entityBuilder) {
  Object fieldValue = IntrospectionUtils.getFieldValue(propertyMetadata, target);
  if (fieldValue == null && propertyMetadata.isOptional()) {
    return;
  }
  ValueBuilder<?, ?, ?> valueBuilder = propertyMetadata.getMapper().toDatastore(fieldValue);
  // ListValues cannot have indexing turned off. Indexing is turned on by
  // default, so we don't touch excludeFromIndexes for ListValues.
  if (valueBuilder.getValueType() != ValueType.LIST) {
    valueBuilder.setExcludeFromIndexes(!propertyMetadata.isIndexed());
  }
  Value<?> datastoreValue = valueBuilder.build();
  entityBuilder.set(propertyMetadata.getMappedName(), datastoreValue);
  Indexer indexer = propertyMetadata.getSecondaryIndexer();
  if (indexer != null) {
    entityBuilder.set(propertyMetadata.getSecondaryIndexName(), indexer.index(datastoreValue));
  }
}
 

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

private static Object getNonEntityObjectFromRow(Object x) {
	Object mappedResult;
	if (x instanceof Key) {
		mappedResult = x;
	}
	else {
		BaseEntity entity = (BaseEntity) x;
		Set<String> colNames = entity.getNames();
		if (colNames.size() > 1) {
			throw new DatastoreDataException(
					"The query method returns non-entity types, but the query result has "
							+ "more than one column. Use a Projection entity type instead.");
		}
		mappedResult = entity.getValue((String) colNames.toArray()[0]).get();
	}
	return mappedResult;
}
 

@SuppressWarnings("unchecked")
private <T> T convertOnReadSingleEmbedded(Object value, TypeInformation<?> targetTypeInformation) {
	Assert.notNull(value, "Cannot convert a null value.");
	if (value instanceof BaseEntity) {
		return (T) this.datastoreEntityConverter.read(targetTypeInformation.getType(),
				(BaseEntity) value);
	}
	throw new DatastoreDataException("Embedded entity was expected, but "
			+ value.getClass() + " found");
}
 

/**
 * Unmarshals the property with the given metadata and sets the unmarshalled value on the given
 * <code>target</code> object.
 * 
 * @param propertyMetadata
 *          the metadata of the property
 * @param target
 *          the target object to set the unmarshalled value on
 * @param nativeEntity
 *          the native entity containing the source property
 * @throws Throwable
 *           propagated
 */
private static void unmarshalProperty(PropertyMetadata propertyMetadata, Object target,
    BaseEntity<?> nativeEntity) throws Throwable {
  // The datastore may not have every property that the entity class has
  // defined. For example, if we are running a projection query or if the
  // entity class added a new field without updating existing data...So
  // make sure there is a property or else, we get an exception from the
  // datastore.
  if (nativeEntity.contains(propertyMetadata.getMappedName())) {
    Value<?> datastoreValue = nativeEntity.getValue(propertyMetadata.getMappedName());
    Object entityValue = propertyMetadata.getMapper().toModel(datastoreValue);
    MethodHandle writeMethod = propertyMetadata.getWriteMethod();
    writeMethod.invoke(target, entityValue);
  }
}
 

/**
 * 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
@SuppressWarnings("unchecked")
public void write(Object source, BaseEntity.Builder sink) {
	DatastorePersistentEntity<?> persistentEntity = this.mappingContext.getPersistentEntity(source.getClass());

	String discriminationFieldName = persistentEntity.getDiscriminationFieldName();
	List<String> discriminationValues = persistentEntity.getCompatibleDiscriminationValues();
	if (!discriminationValues.isEmpty() || discriminationFieldName != null) {
		sink.set(discriminationFieldName,
				discriminationValues.stream().map(StringValue::of).collect(Collectors.toList()));
	}
	PersistentPropertyAccessor accessor = persistentEntity.getPropertyAccessor(source);
	persistentEntity.doWithColumnBackedProperties(
			(DatastorePersistentProperty persistentProperty) -> {
				// Datastore doesn't store its Key as a regular field.
				if (persistentProperty.isIdProperty()) {
					return;
				}
				try {
					Object val = accessor.getProperty(persistentProperty);
					Value convertedVal = this.conversions.convertOnWrite(val, persistentProperty);

					if (persistentProperty.isUnindexed()) {
						convertedVal = setExcludeFromIndexes(convertedVal);
					}
					sink.set(persistentProperty.getFieldName(), convertedVal);
				}
				catch (DatastoreDataException ex) {
					throw new DatastoreDataException(
							"Unable to write "
									+ persistentEntity.kindName() + "." + persistentProperty.getFieldName(),
							ex);
				}
			});
}
 

@Override
public <T, R> Map<T, R> readAsMap(Class<T> keyType, TypeInformation<R> componentType,
		BaseEntity entity) {
	if (entity == null) {
		return null;
	}
	Map<T, R> result = new HashMap<>();
	return readAsMap(entity, ClassTypeInformation.from(result.getClass()));
}
 

@Override
public <T, R> Map<T, R> readAsMap(BaseEntity entity, TypeInformation mapTypeInformation) {
	Assert.notNull(mapTypeInformation, "mapTypeInformation can't be null");
	if (entity == null) {
		return null;
	}
	Map<T, R> result;
	if (!mapTypeInformation.getType().isInterface()) {
		try {
			result = (Map<T, R>) ((Constructor<?>) mapTypeInformation.getType().getConstructor()).newInstance();
		}
		catch (Exception e) {
			throw new DatastoreDataException("Unable to create an instance of a custom map type: "
					+ mapTypeInformation.getType()
					+ " (make sure the class is public and has a public no-args constructor)", e);
		}
	}
	else {
		result = new HashMap<>();
	}

	EntityPropertyValueProvider propertyValueProvider = new EntityPropertyValueProvider(
			entity, this.conversions);
	Set<String> fieldNames = entity.getNames();
	for (String field : fieldNames) {
		result.put(this.conversions.convertOnRead(field, NOT_EMBEDDED, mapTypeInformation.getComponentType()),
				propertyValueProvider.getPropertyValue(field,
						EmbeddedType.of(mapTypeInformation.getMapValueType()),
						mapTypeInformation.getMapValueType()));
	}
	return result;
}
 

private <T, R> Map<T, R> convertOnReadSingleEmbeddedMap(Object value,
		Class<T> keyType, TypeInformation<R> targetComponentType, TypeInformation componentType) {
	Assert.notNull(value, "Cannot convert a null value.");
	if (value instanceof BaseEntity) {
		return this.datastoreEntityConverter.readAsMap((BaseEntity) value, componentType);
	}
	throw new DatastoreDataException(
			"Embedded entity was expected, but " + value.getClass() + " found");
}
 

public EntityPropertyValueProvider(BaseEntity entity, ReadWriteConversions readWriteConversions) {
	if (entity == null) {
		throw new DatastoreDataException("A non-null entity is required");
	}
	this.entity = entity;
	this.conversion = readWriteConversions;
}
 

private <T> void resolveDescendantProperties(DatastorePersistentEntity datastorePersistentEntity,
		BaseEntity entity, T convertedObject, ReadContext context) {
	datastorePersistentEntity
			.doWithDescendantProperties((descendantPersistentProperty) -> {

				Class descendantType = descendantPersistentProperty
						.getComponentType();

				Key entityKey = (Key) entity.getKey();
				Key ancestorKey = KeyUtil.getKeyWithoutAncestors(entityKey);

				EntityQuery descendantQuery = Query.newEntityQueryBuilder()
						.setKind(this.datastoreMappingContext
								.getPersistentEntity(descendantType).kindName())
						.setFilter(PropertyFilter.hasAncestor(ancestorKey))
						.build();

				List entities = convertEntitiesForRead(
						getDatastoreReadWriter().run(descendantQuery), descendantType, context);

				datastorePersistentEntity.getPropertyAccessor(convertedObject)
						.setProperty(descendantPersistentProperty,
								// Converting the collection type.
								this.datastoreEntityConverter.getConversions()
										.convertOnRead(
												entities,
												descendantPersistentProperty.getType(),
												descendantType));
			});
}
 

private <T> List<T> convertEntitiesForRead(Iterator<? extends BaseEntity> entities, Class<T> entityClass, ReadContext context) {
	if (entities == null) {
		return Collections.emptyList();
	}
	List<BaseKey> keys = new ArrayList<>();
	entities.forEachRemaining(e -> {
		IncompleteKey key = e.getKey();
		context.putReadEntity(key, e);
		keys.add(key);
	});
	return convertEntitiesForRead(keys, entityClass, context);
}
 

private GqlQuery<? extends BaseEntity> bindArgsToGqlQueryNoLimit() {
	this.finalGql = this.noLimitQuery;
	this.tagsOrdered = this.tagsOrdered.subList(0, this.limitPosition);
	this.params = this.params.subList(0, this.limitPosition);

	return bindArgsToGqlQuery();
}
 

@Override
@SuppressWarnings("unchecked")
public <R> R read(Class<R> aClass, BaseEntity entity) {
	if (entity == null) {
		return null;
	}
	DatastorePersistentEntity<R> ostensiblePersistentEntity = (DatastorePersistentEntity<R>) this.mappingContext
			.getPersistentEntity(aClass);

	if (ostensiblePersistentEntity == null) {
		throw new DatastoreDataException("Unable to convert Datastore Entity to " + aClass);
	}

	EntityPropertyValueProvider propertyValueProvider = new EntityPropertyValueProvider(entity, this.conversions);

	DatastorePersistentEntity<?> persistentEntity = getDiscriminationPersistentEntity(ostensiblePersistentEntity,
			propertyValueProvider);

	ParameterValueProvider<DatastorePersistentProperty> parameterValueProvider =
			new PersistentEntityParameterValueProvider<>(persistentEntity, propertyValueProvider, null);

	EntityInstantiator instantiator = this.instantiators.getInstantiatorFor(persistentEntity);
	Object instance;
	try {
		instance = instantiator.createInstance(persistentEntity, parameterValueProvider);
		PersistentPropertyAccessor accessor = persistentEntity.getPropertyAccessor(instance);

		persistentEntity.doWithColumnBackedProperties((datastorePersistentProperty) -> {
			// if a property is a constructor argument, it was already computed on instantiation
			if (!persistentEntity.isConstructorArgument(datastorePersistentProperty)) {
				Object value = propertyValueProvider
						.getPropertyValue(datastorePersistentProperty);
				accessor.setProperty(datastorePersistentProperty, value);
			}
		});
	}
	catch (DatastoreDataException ex) {
		throw new DatastoreDataException("Unable to read " + persistentEntity.getName() + " entity", ex);
	}

	return (R) instance;
}
 

void putReadEntity(BaseKey key, BaseEntity entity) {
	this.readEntities.put(key, entity);
}
 

BaseEntity getReadEntity(BaseKey key) {
	return this.readEntities.get(key);
}
 

private GqlQuery<? extends BaseEntity> bindArgsToGqlQuery(Cursor newCursor, int newLimit) {
	this.params.set(this.cursorPosition, newCursor);
	this.params.set(this.limitPosition, newLimit);

	return bindArgsToGqlQuery();
}
 

/**
 * Marshals the given entity (POJO) into the format needed for the low level Cloud Datastore API.
 * 
 * @param entityManager
 *          the entity manager
 * @param entity
 *          the entity to marshal
 * @param intent
 *          the intent or purpose of marshalling. Marshalling process varies slightly depending on
 *          the purpose. For example, if the purpose if INSERT or UPSERT, the marshaller would
 *          auto generate any keys. Where as if the purpose is UPDATE, then then marshaller will
 *          NOT generate any keys.
 * @return the marshaled object
 */
@SuppressWarnings("rawtypes")
public static BaseEntity marshal(DefaultEntityManager entityManager, Object entity,
    Intent intent) {
  Marshaller marshaller = new Marshaller(entityManager, entity, intent);
  return marshaller.marshal();
}
 

/**
 * Unmarshals the given BaseEntity and returns the equivalent model object.
 * 
 * @param nativeEntity
 *          the native entity to unmarshal
 * @param entityClass
 *          the target type of the model class
 * @return the model object
 */
private static <T> T unmarshalBaseEntity(BaseEntity<?> nativeEntity, Class<T> entityClass) {
  if (nativeEntity == null) {
    return null;
  }
  Unmarshaller unmarshaller = new Unmarshaller(nativeEntity, entityClass);
  return unmarshaller.unmarshal();
}
 

/**
 * Read the entity as a {@link Map}.
 * @param <T> the type of the key in the map
 * @param <R> the type of the value in the map
 * @param entity the entity from Cloud Datastore.
 * @param mapTypeInformation the actual type of the map
 * @return a Map where the key values are the field names and the values the field
 * values.
 */
<T, R> Map<T, R> readAsMap(BaseEntity entity, TypeInformation mapTypeInformation);
 

/**
 * Read the entity as a {@link Map} when the map type is unknown.
 * @param <T> the type of the key in the map
 * @param <R> the type of the value in the map
 * @param keyType the key type of the map to be read.
 * @param componentType the value type of the map, into which each field value will be
 * converted.
 * @param entity the entity from Cloud Datastore.
 * @return a Map where the key values are the field names and the values the field
 * values.
 */
<T, R> Map<T, R> readAsMap(Class<T> keyType, TypeInformation<R> componentType,
		BaseEntity entity);
 

/**
 * Creates a new instance of <code>Unmarshaller</code>.
 * 
 * @param nativeEntity
 *          the native entity to unmarshal
 * @param entityClass
 *          the expected model type
 */
private Unmarshaller(BaseEntity<?> nativeEntity, Class<?> entityClass) {
  this.nativeEntity = nativeEntity;
  entityMetadata = EntityIntrospector.introspect(entityClass);

}