support lazy loading

docs/src/main/asciidoc/datastore.adoc

@@ -425,6 +425,7 @@ There are three ways to represent relationships between entities that are descri
 * Embedded entities stored directly in the field of the containing entity
 * `@Descendant` annotated properties for one-to-many relationships
 * `@Reference` annotated properties for general relationships without hierarchy
+* `@LazyReferenceCollection` similar to `@Reference`, but the entities are lazy loaded when the property is accessed (note that the children keys retrieved when the parent entity is loaded)
 
 ==== Embedded Entities
 

spring-cloud-gcp-data-datastore/src/main/java/org/springframework/cloud/gcp/data/datastore/core/DatastoreTemplate.java

@@ -462,13 +462,18 @@ private List<Entity> getReferenceEntitiesForSave(Object entity, Builder builder,
 				return;
 			}
 			Value<?> value;
-			if (persistentProperty.isCollectionLike()) {
-				Iterable<?> iterableVal = (Iterable<?>) ValueUtil.toListIfArray(val);
-				entitiesToSave.addAll(getEntitiesForSave(iterableVal, persistedEntities));
-				List<KeyValue> keyValues = StreamSupport.stream((iterableVal).spliterator(), false)
-						.map((o) -> KeyValue.of(this.getKey(o, false)))
-						.collect(Collectors.toList());
-				value = ListValue.of(keyValues);
+			if (persistentProperty.isCollectionLike() || persistentProperty.getType() == Optional.class) {
+				if (LazyUtil.isLazyAndNotLoaded(val)) {
+					value = ListValue.of(LazyUtil.getKeys(val));
+				}
+				else {
+					Iterable<?> iterableVal = (Iterable<?>) ValueUtil.toListIfArray(val);
+					entitiesToSave.addAll(getEntitiesForSave(iterableVal, persistedEntities));
+					List<KeyValue> keyValues = StreamSupport.stream((iterableVal).spliterator(), false)
+							.map((o) -> KeyValue.of(this.getKey(o, false)))
+							.collect(Collectors.toList());
+					value = ListValue.of(keyValues);
+				}
 			}
 			else {
 				entitiesToSave.addAll(getEntitiesForSave(Collections.singletonList(val), persistedEntities));
@@ -593,35 +598,48 @@ private <T> void resolveReferenceProperties(DatastorePersistentEntity datastoreP
 			BaseEntity entity, T convertedObject, ReadContext context) {
 		datastorePersistentEntity.doWithAssociations(
 				(AssociationHandler) (association) -> {
-					DatastorePersistentProperty referencePersistentProperty = (DatastorePersistentProperty) association
+					DatastorePersistentProperty referenceProperty = (DatastorePersistentProperty) association
 							.getInverse();
-					Object referenced = findReferenced(entity, referencePersistentProperty, context);
-					if (referenced != null) {
-						datastorePersistentEntity.getPropertyAccessor(convertedObject)
-								.setProperty(referencePersistentProperty, referenced);
+					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() && referenceProperty.isCollectionLike()) {
+			List<Value<Key>> keyList = entity.getList(fieldName);
+			DatastoreReaderWriter originalTx = getDatastoreReadWriter();
+			referenced = LazyUtil.wrapSimpleLazyProxy((List<Value<Key>> storedKeys) -> {
+				if (getDatastoreReadWriter() != originalTx) {
+					throw new DatastoreDataException("Lazy load should be invoked within the same transaction");
+				}
+				return (T) fetchReferenced(referenceProperty, context, valuesToKeys(storedKeys));
+			}, type, keyList);
+		}
+		else {
+			referenced = (T) findReferenced(entity, referenceProperty, context);
+		}
+		return referenced;
+	}
+
+	// Extracts key(s) from a property, fetches and if necessary, converts values to the required type
 	private Object findReferenced(BaseEntity entity, DatastorePersistentProperty referencePersistentProperty,
 			ReadContext context) {
 		String fieldName = referencePersistentProperty.getFieldName();
 		try {
 			Object referenced;
-			if (!entity.contains(fieldName)) {
-				referenced = null;
-			}
-			else if (referencePersistentProperty.isCollectionLike()) {
-				Class referencedType = referencePersistentProperty.getComponentType();
-				List<Value<Key>> keyValues = entity.getList(fieldName);
-				referenced = this.datastoreEntityConverter.getConversions()
-						.convertOnRead(
-								findAllById(
-										keyValues.stream().map(Value::get).collect(Collectors.toSet()),
-										referencedType, context),
-								referencePersistentProperty.getType(),
-								referencedType);
+			if (referencePersistentProperty.isCollectionLike()) {
+				referenced = fetchReferenced(referencePersistentProperty, context,
+						valuesToKeys(entity.getList(fieldName)));
 			}
 			else {
 				List referencedList = findAllById(Collections.singleton(entity.getKey(fieldName)),
@@ -638,6 +656,23 @@ else if (referencePersistentProperty.isCollectionLike()) {
 			}
 	}
 
+	// Given keys, fetches and converts values to the required collection type
+	private Object fetchReferenced(DatastorePersistentProperty referencePersistentProperty, ReadContext context,
+			Set<Key> keys) {
+		Class referencedType = referencePersistentProperty.getComponentType();
+		return this.datastoreEntityConverter.getConversions()
+				.convertOnRead(
+						findAllById(
+								keys,
+								referencedType, context),
+						referencePersistentProperty.getType(),
+						referencedType);
+	}
+
+	private Set<Key> valuesToKeys(List<Value<Key>> keyValues) {
+		return keyValues.stream().map(Value::get).collect(Collectors.toSet());
+	}
+
 	private <T> void resolveDescendantProperties(DatastorePersistentEntity datastorePersistentEntity,
 			BaseEntity entity, T convertedObject, ReadContext context) {
 		datastorePersistentEntity

spring-cloud-gcp-data-datastore/src/main/java/org/springframework/cloud/gcp/data/datastore/core/LazyUtil.java

@@ -0,0 +1,126 @@
+/*
+ * Copyright 2017-2019 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.cloud.gcp.data.datastore.core;
+
+import java.lang.reflect.InvocationHandler;
+import java.lang.reflect.Method;
+import java.lang.reflect.Proxy;
+import java.util.List;
+import java.util.function.Function;
+
+import com.google.cloud.datastore.Key;
+import com.google.cloud.datastore.Value;
+
+import org.springframework.cloud.gcp.data.datastore.core.mapping.DatastoreDataException;
+import org.springframework.util.Assert;
+
+/**
+ * Utilities used to support lazy loaded properties.
+ *
+ * @author Dmitry Solomakha
+ *
+ * @since 1.3
+ */
+final class LazyUtil {
+
+	private LazyUtil() {
+	}
+
+	public static <T> T wrapSimpleLazyProxy(Function<List<Value<Key>>, T> supplierFunc, Class<T> type, List keys) {
+		return (T) Proxy.newProxyInstance(type.getClassLoader(), new Class[] {type},
+				new SimpleLazyDynamicInvocationHandler<T>(supplierFunc, keys));
+	}
+
+	/**
+	 * Check if the object is a lazy loaded proxy that hasn't been evaluated.
+	 * @param object an object
+	 * @return true if the object is a proxy that was not evaluated
+	 */
+	public static boolean isLazyAndNotLoaded(Object object) {
+		SimpleLazyDynamicInvocationHandler handler = getProxy(object);
+		if (handler != null) {
+			return !handler.isEvaluated() && handler.getKeys() != null;
+		}
+		return false;
+	}
+
+	/**
+	 * Extract keys from a proxy object.
+	 * @param object a proxy object
+	 * @return list of keys if the object is a proxy, null otherwise
+	 */
+	public static List getKeys(Object object) {
+		SimpleLazyDynamicInvocationHandler handler = getProxy(object);
+		if (handler != null) {
+			if (!handler.isEvaluated()) {
+				return handler.getKeys();
+			}
+		}
+		return null;
+	}
+
+	private static SimpleLazyDynamicInvocationHandler getProxy(Object object) {
+		if (Proxy.isProxyClass(object.getClass())
+				&& (Proxy.getInvocationHandler(object) instanceof SimpleLazyDynamicInvocationHandler)) {
+			return (SimpleLazyDynamicInvocationHandler) Proxy
+					.getInvocationHandler(object);
+		}
+		return null;
+	}
+
+	/**
+	 * Proxy class used for lazy loading.
+	 */
+	public static final class SimpleLazyDynamicInvocationHandler<T> implements InvocationHandler {
+
+		private final Function<List<Value<Key>>, T> supplierFunc;
+
+		private final List<Value<Key>> keys;
+
+		private boolean isEvaluated = false;
+
+		private T value;
+
+		private SimpleLazyDynamicInvocationHandler(Function<List<Value<Key>>, T> supplierFunc, List<Value<Key>> keys) {
+			Assert.notNull(supplierFunc, "A non-null supplier function is required for a lazy proxy.");
+			this.supplierFunc = supplierFunc;
+			this.keys = keys;
+		}
+
+		private boolean isEvaluated() {
+			return this.isEvaluated;
+		}
+
+		public List getKeys() {
+			return this.keys;
+		}
+
+		@Override
+		public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
+			if (!this.isEvaluated) {
+				T value = this.supplierFunc.apply(this.keys);
+				if (value == null) {
+					throw new DatastoreDataException("Can't load referenced entity");
+				}
+				this.value = value;
+
+				this.isEvaluated = true;
+			}
+			return method.invoke(this.value, args);
+		}
+	}
+}

spring-cloud-gcp-data-datastore/src/main/java/org/springframework/cloud/gcp/data/datastore/core/mapping/DatastorePersistentProperty.java

@@ -60,4 +60,10 @@ public interface DatastorePersistentProperty
 	 * @return true if the property is stored within Datastore entity
 	 */
 	boolean isColumnBacked();
+
+	/**
+	 * Return whether this property is a lazily-fetched one.
+	 * @return {@code true} if the property is lazily-fetched. {@code false} otherwise.
+	 */
+	boolean isLazyLoaded();
 }

spring-cloud-gcp-data-datastore/src/main/java/org/springframework/cloud/gcp/data/datastore/core/mapping/DatastorePersistentPropertyImpl.java

@@ -80,6 +80,10 @@ private void verify() {
 					"Only collection-like properties can contain the "
 							+ "descendant entity objects can be annotated @Descendants.");
 		}
+		if (isLazyLoaded() && !isCollectionLike()) {
+			throw new DatastoreDataException(
+					"Only collection-like properties can be lazy-loaded");
+		}
 	}
 
 	@Override
@@ -140,4 +144,9 @@ public Iterable<? extends TypeInformation<?>> getPersistentEntityTypes() {
 				.filter((typeInfo) -> typeInfo.getType().isAnnotationPresent(Entity.class))
 				.collect(Collectors.toList());
 	}
+
+	@Override
+	public boolean isLazyLoaded() {
+		return findAnnotation(LazyReferenceCollection.class) != null;
+	}
 }

spring-cloud-gcp-data-datastore/src/main/java/org/springframework/cloud/gcp/data/datastore/core/mapping/LazyReferenceCollection.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2017-2018 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.cloud.gcp.data.datastore.core.mapping;
+
+import java.lang.annotation.Documented;
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+import org.springframework.data.annotation.Reference;
+
+/**
+ * Annotation for a class that indicates that a property is a collection of lazy loaded Datastore entities.
+ *
+ * @author Dmitry Solomakha
+ *
+ * @since 1.3
+ */
+@Documented
+@Target(ElementType.FIELD)
+@Retention(RetentionPolicy.RUNTIME)
+@Reference
+public @interface LazyReferenceCollection {
+}