From 02a1e65a4bd5c33039d4ee438591fc3632bb0c1f Mon Sep 17 00:00:00 2001
From: ldetmer <1771267+ldetmer@users.noreply.github.com>
Date: Wed, 4 Sep 2024 09:48:12 -0400
Subject: [PATCH] fix: Descendants entities losing database ID from parent
 (#3142) (#3156)

* fix: need to also send database ID when migrating keys from ancestor to child
---
 .../data/datastore/core/util/KeyUtil.java     |  1 +
 .../data/datastore/core/util/KeyUtilTest.java | 25 +++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/spring-cloud-gcp-data-datastore/src/main/java/com/google/cloud/spring/data/datastore/core/util/KeyUtil.java b/spring-cloud-gcp-data-datastore/src/main/java/com/google/cloud/spring/data/datastore/core/util/KeyUtil.java
index 9b8da5034b..8236b041ec 100644
--- a/spring-cloud-gcp-data-datastore/src/main/java/com/google/cloud/spring/data/datastore/core/util/KeyUtil.java
+++ b/spring-cloud-gcp-data-datastore/src/main/java/com/google/cloud/spring/data/datastore/core/util/KeyUtil.java
@@ -46,6 +46,7 @@ public static Key getKeyWithoutAncestors(Key entityKey) {
           Key.newBuilder(entityKey.getProjectId(), entityKey.getKind(), entityKey.getId());
     }
     ancestorLookupKey.setNamespace(entityKey.getNamespace());
+    ancestorLookupKey.setDatabaseId(entityKey.getDatabaseId());
 
     return ancestorLookupKey.build();
   }
diff --git a/spring-cloud-gcp-data-datastore/src/test/java/com/google/cloud/spring/data/datastore/core/util/KeyUtilTest.java b/spring-cloud-gcp-data-datastore/src/test/java/com/google/cloud/spring/data/datastore/core/util/KeyUtilTest.java
index 964e1340c0..d851de8f99 100644
--- a/spring-cloud-gcp-data-datastore/src/test/java/com/google/cloud/spring/data/datastore/core/util/KeyUtilTest.java
+++ b/spring-cloud-gcp-data-datastore/src/test/java/com/google/cloud/spring/data/datastore/core/util/KeyUtilTest.java
@@ -47,4 +47,29 @@ void testRemoveAncestors_IdKeys() {
     Key processedKey = KeyUtil.getKeyWithoutAncestors(idKey);
     assertThat(processedKey.getAncestors()).isEmpty();
   }
+
+  @Test
+  void testAncestorKeys_containsAllDataStoreMetaData() {
+    String projectId = "project-id";
+    String kind = "kind";
+    Long id = 13L;
+    String databaseId = "database-id";
+    String namespace = "namespace";
+
+    Key idKey =
+        Key.newBuilder(projectId, kind, id, databaseId)
+            .setNamespace(namespace)
+            .addAncestor(PathElement.of("person", 22L))
+            .addAncestor(PathElement.of("person", 18L))
+            .build();
+
+    Key processedKey = KeyUtil.getKeyWithoutAncestors(idKey);
+
+    assertThat(processedKey.getAncestors()).isEmpty();
+    assertThat(processedKey.getId()).isEqualTo(id);
+    assertThat(processedKey.getKind()).isEqualTo(kind);
+    assertThat(processedKey.getDatabaseId()).isEqualTo(databaseId);
+    assertThat(processedKey.getNamespace()).isEqualTo(namespace);
+    assertThat(processedKey.getProjectId()).isEqualTo(projectId);
+  }
 }