fix: Datastore lazy reference save/update behavior

spring-cloud-gcp-data-datastore/src/main/java/com/google/cloud/spring/data/datastore/core/DatastoreTemplate.java

@@ -165,6 +165,10 @@ private <T> List<Entity> getEntitiesForSave(
       Iterable<T> entities, Set<Key> persisted, Key... ancestors) {
     List<Entity> entitiesForSave = new LinkedList<>();
     for (T entity : entities) {
+      // If entity is lazy referenced, now load all properties before attempt to save
+      if (LazyUtil.isLazy(entity)) {
+        entity = (T) LazyUtil.getLazyValue(entity);
+      }
       Key key = getKey(entity, true, ancestors);
       if (!persisted.contains(key)) {
         persisted.add(key);
@@ -523,7 +527,7 @@ private static StructuredQuery.OrderBy createOrderBy(
 
   private List<Entity> convertToEntityForSave(
       Object entity, Set<Key> persistedEntities, Key... ancestors) {
-    if (ancestors != null) {
+    if (ancestors.length != 0) {
       for (Key ancestor : ancestors) {
         validateKey(entity, keyToPathElement(ancestor));
       }
@@ -565,6 +569,10 @@ private List<Entity> getReferenceEntitiesForSave(
                 value = ListValue.of(keyValues);
 
               } else {
+                // If property is lazy referenced, now load it before attempt to save.
+                if (LazyUtil.isLazy(val)) {
+                  val = LazyUtil.getLazyValue(val);
+                }
                 entitiesToSave.addAll(
                     getEntitiesForSave(Collections.singletonList(val), persistedEntities));
                 Key key = getKey(val, false);

spring-cloud-gcp-data-datastore/src/main/java/com/google/cloud/spring/data/datastore/core/LazyUtil.java

@@ -90,6 +90,36 @@ public static boolean isLazyAndNotLoaded(Object object) {
     return false;
   }
 
+  /**
+   * Check if the object is a lazy loaded proxy.
+   *
+   * @param object an object
+   * @return true if the object is a proxy
+   */
+  static boolean isLazy(Object object) {
+    SimpleLazyDynamicInvocationHandler handler = getProxy(object);
+    if (handler != null) {
+      return handler.getKeys() != null;
+    }
+    return false;
+  }
+
+  /**
+   * Loads the value provided by the proxy supplier.
+   *
+   * @param object an object that is a proxy
+   * @return value provided by the proxy supplier.
+   */
+  static Object getLazyValue(Object object) {
+    Assert.isTrue(isLazy(object), "A proxy object is required.");
+
+    SimpleLazyDynamicInvocationHandler handler = getProxy(object);
+    if (handler != null) {
+      return handler.getValue();
+    }
+    return null;
+  }
+
   /**
    * Extract keys from a proxy object.
    *
@@ -145,6 +175,10 @@ public Value getKeys() {
       return this.keys;
     }
 
+    T getValue() {
+      return value;
+    }
+
     @Override
     public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
       if (!this.isEvaluated) {

spring-cloud-gcp-data-datastore/src/test/java/com/google/cloud/spring/data/datastore/it/ChildEntity.java

@@ -0,0 +1,48 @@
+/*
+ * Copyright 2022-2022 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 com.google.cloud.spring.data.datastore.it;
+
+import com.google.cloud.spring.data.datastore.core.mapping.Entity;
+import org.springframework.data.annotation.Id;
+
+@Entity
+public class ChildEntity {
+  @Id String id;
+
+  String value;
+
+  public ChildEntity(String id, String value) {
+    this.id = id;
+    this.value = value;
+  }
+
+  public String getId() {
+    return id;
+  }
+
+  // public void setId(String id) {
+  //   this.id = id;
+  // }
+
+  public String getValue() {
+    return value;
+  }
+
+  public void setValue(String value) {
+    this.value = value;
+  }
+}

spring-cloud-gcp-data-datastore/src/test/java/com/google/cloud/spring/data/datastore/it/DatastoreIntegrationTests.java

@@ -41,6 +41,7 @@
 import com.google.cloud.spring.data.datastore.core.mapping.DiscriminatorField;
 import com.google.cloud.spring.data.datastore.core.mapping.DiscriminatorValue;
 import com.google.cloud.spring.data.datastore.core.mapping.Entity;
+import com.google.cloud.spring.data.datastore.core.mapping.LazyReference;
 import com.google.cloud.spring.data.datastore.core.mapping.Unindexed;
 import com.google.cloud.spring.data.datastore.entities.CustomMap;
 import com.google.cloud.spring.data.datastore.entities.Product;
@@ -145,6 +146,8 @@ void deleteAll() {
     this.datastoreTemplate.deleteAll(LazyEntity.class);
     this.datastoreTemplate.deleteAll(Product.class);
     this.datastoreTemplate.deleteAll(Store.class);
+    this.datastoreTemplate.deleteAll(ChildEntity.class);
+    this.datastoreTemplate.deleteAll(ParentEntityWithLazyChild.class);
     this.testEntityRepository.deleteAll();
     if (this.keyForMap != null) {
       this.datastore.delete(this.keyForMap);
@@ -739,6 +742,45 @@ void lazyReferenceTest() throws InterruptedException {
     assertThat(loadedParent).isEqualTo(lazyParentEntity);
   }
 
+  @Test
+  void lazyReferenceUpdateAndSaveParentTest() {
+    LazyEntity lazyParentEntity = new LazyEntity(new LazyEntity(new LazyEntity()));
+    this.datastoreTemplate.save(lazyParentEntity);
+
+    LazyEntity loadedParent =
+        this.datastoreTemplate.findById(lazyParentEntity.id, LazyEntity.class);
+
+    // update lazyReferenced property
+    LazyEntity newChild = new LazyEntity();
+    loadedParent.setLazyChild(newChild);
+    this.datastoreTemplate.save(loadedParent);
+
+    loadedParent = this.datastoreTemplate.findById(loadedParent.id, LazyEntity.class);
+    assertThat(loadedParent.getLazyChild()).isEqualTo(newChild);
+  }
+
+  @Test
+  void lazyReferenceWithStringIdUpdateChildTest() {
+
+    ChildEntity child = new ChildEntity("child-entity-id-string", "child value");
+    // child.id = "abcdse";
+    ParentEntityWithLazyChild parentEntityWithLazyChild = new ParentEntityWithLazyChild(child);
+    this.datastoreTemplate.save(parentEntityWithLazyChild);
+
+    ParentEntityWithLazyChild loadedParent = this.datastoreTemplate.findById(
+        parentEntityWithLazyChild.id, ParentEntityWithLazyChild.class);
+    ChildEntity lazyChild = loadedParent.getLazyChild();
+
+    lazyChild.setValue("updated child");
+    this.datastoreTemplate.save(lazyChild);
+
+    loadedParent = this.datastoreTemplate.findById(
+        parentEntityWithLazyChild.id, ParentEntityWithLazyChild.class);
+    assertThat(loadedParent.getLazyChild().getId()).isEqualTo("child-entity-id-string");
+    assertThat(loadedParent.getLazyChild().getValue()).isEqualTo("updated child");
+
+  }
+
   @Test
   void singularLazyPropertyTest() {
     LazyEntity lazyParentEntity = new LazyEntity(new LazyEntity(new LazyEntity()));
@@ -1396,6 +1438,24 @@ public String toString() {
   }
 }
 
+@Entity
+class ParentEntityWithLazyChild {
+  @Id Long id;
+
+  @LazyReference
+  ChildEntity lazyChild;
+
+  public ParentEntityWithLazyChild() {}
+
+  ParentEntityWithLazyChild(ChildEntity child) {
+    this.lazyChild = child;
+  }
+
+  ChildEntity getLazyChild() {
+    return this.lazyChild;
+  }
+}
+
 enum CommunicationChannels {
   EMAIL,
   SMS;

spring-cloud-gcp-data-datastore/src/test/java/com/google/cloud/spring/data/datastore/it/LazyEntity.java

@@ -41,6 +41,10 @@ LazyEntity getLazyChild() {
     return this.lazyChild;
   }
 
+  public void setLazyChild(LazyEntity lazyChild) {
+    this.lazyChild = lazyChild;
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) {