From b2b61d7b62b83d4a36934d131e19293e97b900bf Mon Sep 17 00:00:00 2001
From: sonika-shah <58761340+sonika-shah@users.noreply.github.com>
Date: Wed, 3 Sep 2025 08:26:10 +0530
Subject: [PATCH] Fix Data Product to asset mapping lost after upgrade (#23195)

---
 .../service/jdbi3/EntityRepository.java       | 36 ++++++++++-
 .../service/resources/EntityResourceTest.java | 63 +++++++++++++++++--
 2 files changed, 93 insertions(+), 6 deletions(-)

diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityRepository.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityRepository.java
index b5958897f32..3fafc88f4ab 100644
--- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityRepository.java
+++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityRepository.java
@@ -5376,7 +5376,7 @@ public abstract class EntityRepository<T extends EntityInterface> {
   }
 
   private Map<UUID, List<EntityReference>> batchFetchDataProducts(List<T> entities) {
-    return batchFetchFromIdsManyToOne(entities, Relationship.HAS, Entity.DATA_PRODUCT);
+    return batchFetchToIdsOneToMany(entities, Relationship.HAS, Entity.DATA_PRODUCT);
   }
 
   private Map<UUID, AssetCertification> batchFetchCertification(List<T> entities) {
@@ -5682,6 +5682,40 @@ public abstract class EntityRepository<T extends EntityInterface> {
     return resultMap;
   }
 
+  protected Map<UUID, List<EntityReference>> batchFetchToIdsOneToMany(
+      List<T> entities, Relationship relationship, String fromEntityType) {
+    var resultMap = new HashMap<UUID, List<EntityReference>>();
+    if (entities == null || entities.isEmpty()) {
+      return resultMap;
+    }
+
+    // Use Include.ALL to get all relationships including those for soft-deleted entities
+    var records =
+        daoCollection
+            .relationshipDAO()
+            .findFromBatch(
+                entityListToStrings(entities), relationship.ordinal(), fromEntityType, ALL);
+
+    var idReferenceMap =
+        Entity.getEntityReferencesByIds(
+                fromEntityType,
+                records.stream().map(e -> UUID.fromString(e.getFromId())).distinct().toList(),
+                ALL)
+            .stream()
+            .collect(Collectors.toMap(e -> e.getId().toString(), Function.identity()));
+
+    records.forEach(
+        record -> {
+          var entityId = UUID.fromString(record.getToId());
+          var relatedRef = idReferenceMap.get(record.getFromId());
+          if (relatedRef != null) {
+            resultMap.computeIfAbsent(entityId, k -> new ArrayList<>()).add(relatedRef);
+          }
+        });
+
+    return resultMap;
+  }
+
   protected Map<UUID, EntityReference> batchFetchFromIdsAndRelationSingleRelation(
       List<T> entities, Relationship relationship) {
     var resultMap = new HashMap<UUID, EntityReference>();
diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/resources/EntityResourceTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/resources/EntityResourceTest.java
index ecbaa34d125..7264213faed 100644
--- a/openmetadata-service/src/test/java/org/openmetadata/service/resources/EntityResourceTest.java
+++ b/openmetadata-service/src/test/java/org/openmetadata/service/resources/EntityResourceTest.java
@@ -866,6 +866,18 @@ public abstract class EntityResourceTest<T extends EntityInterface, K extends Cr
     Domain testDomain2 =
         domainResourceTest.createEntity(
             domainResourceTest.createRequest(test, 4), ADMIN_AUTH_HEADERS);
+
+    // Create DataProduct for testing if supported
+    DataProductResourceTest dataProductResourceTest = new DataProductResourceTest();
+    DataProduct testDataProduct = null;
+    if (supportsDataProducts && supportsDomains) {
+      CreateDataProduct createDataProduct =
+          dataProductResourceTest
+              .createRequest(test, 5)
+              .withDomains(List.of(testDomain1.getFullyQualifiedName()));
+      testDataProduct = dataProductResourceTest.createEntity(createDataProduct, ADMIN_AUTH_HEADERS);
+    }
+
     final String testName = "000_" + getEntityName(test);
     final K createRequest =
         createRequest(testName, "Test entity for field fetching", testName, null);
@@ -909,6 +921,13 @@ public abstract class EntityResourceTest<T extends EntityInterface, K extends Cr
                     .withName(testDomain1.getName()));
         entity.setDomains(domains);
         entity = patchEntity(entityId, originalJson, entity, ADMIN_AUTH_HEADERS);
+        originalJson = JsonUtils.pojoToJson(entity);
+
+        // Add DataProducts if supported
+        if (supportsDataProducts && testDataProduct != null) {
+          entity.setDataProducts(List.of(testDataProduct.getEntityReference()));
+          entity = patchEntity(entityId, originalJson, entity, ADMIN_AUTH_HEADERS);
+        }
       }
 
       Map<String, String> params = new HashMap<>();
@@ -926,12 +945,20 @@ public abstract class EntityResourceTest<T extends EntityInterface, K extends Cr
       if (supportsOwners) fieldCombinationsList.add("owners");
       if (supportsTags) fieldCombinationsList.add("tags");
       if (supportsFollowers) fieldCombinationsList.add("followers");
-      if (supportsDomains)
-        fieldCombinationsList.add("domains"); // Always test fetching domains if supported
+      if (supportsDomains) fieldCombinationsList.add("domains");
+      if (supportsDataProducts) fieldCombinationsList.add("dataProducts");
+      if (supportsExperts) fieldCombinationsList.add("experts");
+      if (supportsReviewers) fieldCombinationsList.add("reviewers");
+      if (supportsVotes) fieldCombinationsList.add("votes");
+
+      // Test combinations
       if (supportsOwners && supportsTags) fieldCombinationsList.add("owners,tags");
       if (supportsFollowers && supportsOwners) fieldCombinationsList.add("followers,owners");
-      if (supportsDomains && supportsTags)
-        fieldCombinationsList.add("domains,tags"); // Always test fetching domains if supported
+      if (supportsDomains && supportsTags) fieldCombinationsList.add("domains,tags");
+      if (supportsDataProducts && supportsDomains)
+        fieldCombinationsList.add("dataProducts,domains");
+
+      // Always test with all allowed fields
       fieldCombinationsList.add(getAllowedFields());
 
       for (String fields : fieldCombinationsList) {
@@ -1011,7 +1038,6 @@ public abstract class EntityResourceTest<T extends EntityInterface, K extends Cr
           assertFalse(indivDomains.isEmpty(), "Individual domains should not be empty");
 
           final UUID domain1Id = testDomain1.getId();
-          final UUID domain2Id = testDomain2.getId();
 
           assertTrue(
               batchDomains.stream().anyMatch(d -> d.getId().equals(domain1Id)),
@@ -1021,12 +1047,39 @@ public abstract class EntityResourceTest<T extends EntityInterface, K extends Cr
               "Should find test domain 1 in individual domains");
         }
 
+        if (fields.contains("dataProducts") && supportsDataProducts && testDataProduct != null) {
+          List<EntityReference> batchDataProducts = listOrEmpty(batchEntityFound.getDataProducts());
+          List<EntityReference> indivDataProducts = listOrEmpty(individualEntity.getDataProducts());
+
+          assertEquals(
+              indivDataProducts.size(),
+              batchDataProducts.size(),
+              "DataProducts count mismatch between batch and individual fetch - This is the bug!");
+
+          if (!indivDataProducts.isEmpty()) {
+            assertFalse(
+                batchDataProducts.isEmpty(),
+                "Batch DataProducts empty when individual has them - This is the exact bug!");
+
+            final UUID dataProductId = testDataProduct.getId();
+            assertTrue(
+                batchDataProducts.stream().anyMatch(dp -> dp.getId().equals(dataProductId)),
+                "Should find test data product in batch DataProducts");
+            assertTrue(
+                indivDataProducts.stream().anyMatch(dp -> dp.getId().equals(dataProductId)),
+                "Should find test data product in individual DataProducts");
+          }
+        }
+
         LOG.info("Successfully verified field combination: {}", fields);
       }
 
     } finally {
       if (supportsOwners) userResourceTest.deleteEntity(testUser.getId(), ADMIN_AUTH_HEADERS);
       if (supportsTags) tagResourceTest.deleteEntity(testTag.getId(), ADMIN_AUTH_HEADERS);
+      if (supportsDataProducts && testDataProduct != null) {
+        dataProductResourceTest.deleteEntity(testDataProduct.getId(), ADMIN_AUTH_HEADERS);
+      }
       if (supportsDomains) {
         domainResourceTest.deleteEntity(testDomain1.getId(), ADMIN_AUTH_HEADERS);
         domainResourceTest.deleteEntity(testDomain2.getId(), ADMIN_AUTH_HEADERS);