support lazy loading

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

@@ -66,6 +66,7 @@
 import org.springframework.cloud.gcp.data.datastore.core.mapping.event.BeforeDeleteEvent;
 import org.springframework.cloud.gcp.data.datastore.core.mapping.event.BeforeSaveEvent;
 import org.springframework.cloud.gcp.data.datastore.core.util.KeyUtil;
+import org.springframework.cloud.gcp.data.datastore.core.util.LazyUtil;
 import org.springframework.cloud.gcp.data.datastore.core.util.SliceUtil;
 import org.springframework.cloud.gcp.data.datastore.core.util.ValueUtil;
 import org.springframework.context.ApplicationEvent;
@@ -463,12 +464,17 @@ private List<Entity> getReferenceEntitiesForSave(Object entity, Builder builder,
 			}
 			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 (LazyUtil.hasUsableKeys(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,14 +599,31 @@ 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;
+						if (referenceProperty.isLazyLoaded() && referenceProperty.isCollectionLike()) {
+							List keyList = entity.getList(fieldName);
+							DatastoreReaderWriter originalTx = getDatastoreReadWriter();
+							referenced = LazyUtil.wrapSimpleLazyProxy(() -> {
+								if (getDatastoreReadWriter() != originalTx) {
+									throw new DatastoreDataException("Lazy load should be invoked within the same transaction");
+								}
+								return fetchReferenced(referenceProperty, context,
+										valuesToKeys(keyList));
+							}, type, keyList);
+						}
+						else {
+							referenced = findReferenced(entity, referenceProperty, context);
+						}
+						if (referenced != null) {
+							datastorePersistentEntity.getPropertyAccessor(convertedObject)
+									.setProperty(referenceProperty, referenced);
+						}
 					}
-
 				});
 	}
 
@@ -609,19 +632,8 @@ private Object findReferenced(BaseEntity entity, DatastorePersistentProperty ref
 		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 +650,21 @@ else if (referencePersistentProperty.isCollectionLike()) {
 			}
 	}
 
+	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/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

@@ -140,4 +140,10 @@ public Iterable<? extends TypeInformation<?>> getPersistentEntityTypes() {
 				.filter((typeInfo) -> typeInfo.getType().isAnnotationPresent(Entity.class))
 				.collect(Collectors.toList());
 	}
+
+	@Override
+	public boolean isLazyLoaded() {
+		ReferenceCollection annotation = findAnnotation(ReferenceCollection.class);
+		return annotation != null && annotation.lazy();
+	}
 }

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

@@ -0,0 +1,47 @@
+/*
+ * 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 it is an entity stored in a Datastore Entity.
+ *
+ * @author Dmitry Solomakha
+ *
+ * @since 1.3
+ */
+@Documented
+@Target(ElementType.FIELD)
+@Retention(RetentionPolicy.RUNTIME)
+@Reference
+public @interface ReferenceCollection {
+
+	/**
+	 * Controls whether the referenced entity should be loaded lazily. This defaults to
+	 * {@literal false}.
+	 *
+	 * @return whether the interleaved property is retrieved lazily.
+	 */
+	boolean lazy() default false;
+}

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

@@ -0,0 +1,102 @@
+/*
+ * 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.util;
+
+import java.lang.reflect.InvocationHandler;
+import java.lang.reflect.Method;
+import java.lang.reflect.Proxy;
+import java.util.List;
+import java.util.function.Supplier;
+
+import org.springframework.cloud.gcp.data.datastore.core.mapping.DatastoreDataException;
+import org.springframework.util.Assert;
+
+/**
+ * @author Dmitry Solomakha
+ */
+final public class LazyUtil {
+
+	private LazyUtil() {
+	}
+
+	public static <T> T wrapSimpleLazyProxy(Supplier<T> supplierFunc, Class<T> type, List keys) {
+		return (T) Proxy.newProxyInstance(type.getClassLoader(), new Class[] {type},
+				new SimpleLazyDynamicInvocationHandler<T>(supplierFunc, keys));
+	}
+
+	public static boolean hasUsableKeys(Object object) {
+		if (Proxy.isProxyClass(object.getClass())
+				&& (Proxy.getInvocationHandler(object) instanceof SimpleLazyDynamicInvocationHandler)) {
+			SimpleLazyDynamicInvocationHandler handler = (SimpleLazyDynamicInvocationHandler) Proxy
+					.getInvocationHandler(object);
+			return !handler.isEvaluated() && handler.getKeys() != null;
+		}
+		return false;
+	}
+
+	public static List getKeys(Object object) {
+		if (Proxy.isProxyClass(object.getClass())
+				&& (Proxy.getInvocationHandler(object) instanceof SimpleLazyDynamicInvocationHandler)) {
+			SimpleLazyDynamicInvocationHandler handler = (SimpleLazyDynamicInvocationHandler) Proxy
+					.getInvocationHandler(object);
+			if (!handler.isEvaluated()) {
+				return handler.getKeys();
+			}
+		}
+		return null;
+	}
+
+	public static final class SimpleLazyDynamicInvocationHandler<T> implements InvocationHandler {
+
+		private final Supplier<T> supplierFunc;
+
+		private final List keys;
+
+		private boolean isEvaluated = false;
+
+		private T value;
+
+		private SimpleLazyDynamicInvocationHandler(Supplier<T> supplierFunc, List 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.get();
+				if (value == null) {
+					throw new DatastoreDataException("Can't load referenced entity");
+				}
+				this.value = value;
+
+				this.isEvaluated = true;
+			}
+			return method.invoke(this.value, args);
+		}
+	}
+
+}