fix(gms): Return empty snapshots when one does not exist (#2646)

datahub-project · Jun 4, 2021 · d26903f · d26903f
1 parent cc3caa9
commit d26903f
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 17 deletions.
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java
@@ -170,7 +170,17 @@ public Entity getEntity(@Nonnull final Urn urn, @Nonnull final Set<String> aspec
         .orElse(null);
   }
 
-  public Map<Urn, Entity> getEntities(@Nonnull final Set<Urn> urns, @Nonnull final Set<String> aspectNames) {
+  /**
+   * Retrieves multiple entities.
+   *
+   * @param urns set of urns to fetch
+   * @param aspectNames set of aspects to fetch
+   * @return a map of {@link Urn} to {@link Entity} object
+   */
+  public Map<Urn, Entity> getEntities(@Nonnull final Set<Urn> urns, @Nonnull Set<String> aspectNames) {
+    if (urns.isEmpty()) {
+      return Collections.emptyMap();
+    }
     return getSnapshotUnions(urns, aspectNames).entrySet().stream()
         .collect(Collectors.toMap(Map.Entry::getKey, entry -> toEntity(entry.getValue())));
   }

diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/EbeanEntityService.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/EbeanEntityService.java
@@ -55,17 +55,29 @@ public Map<Urn, List<RecordTemplate>> getLatestAspects(@Nonnull final Set<Urn> u
     final Set<EbeanAspectV2.PrimaryKey> dbKeys = urns.stream()
         .map(urn -> {
           final Set<String> aspectsToFetch = aspectNames.isEmpty()
-              ? getEntityAspectNames(urn)
-              : aspectNames;
+            ? getEntityAspectNames(urn)
+            : aspectNames;
           return aspectsToFetch.stream()
-              .map(aspectName -> new EbeanAspectV2.PrimaryKey(urn.toString(), aspectName, LATEST_ASPECT_VERSION))
-              .collect(Collectors.toList());
+            .map(aspectName -> new EbeanAspectV2.PrimaryKey(urn.toString(), aspectName, LATEST_ASPECT_VERSION))
+            .collect(Collectors.toList());
         })
         .flatMap(List::stream)
         .collect(Collectors.toSet());
 
     // Fetch from db and populate urn -> aspect map.
     final Map<Urn, List<RecordTemplate>> urnToAspects = new HashMap<>();
+
+    // Each urn should have some result, regardless of whether aspects are found in the DB.
+    for (Urn urn: urns) {
+      urnToAspects.putIfAbsent(urn, new ArrayList<>());
+    }
+
+    // Add "key" aspects for each urn. TODO: Replace this with a materialized key aspect.
+    urnToAspects.keySet().forEach(key -> {
+      final RecordTemplate keyAspect = buildKeyAspect(key);
+      urnToAspects.get(key).add(keyAspect);
+    });
+
     _entityDao.batchGet(dbKeys).forEach((key, aspectEntry) -> {
       final Urn urn = toUrn(key.getUrn());
       final String aspectName = key.getAspect();
@@ -74,12 +86,6 @@ public Map<Urn, List<RecordTemplate>> getLatestAspects(@Nonnull final Set<Urn> u
       urnToAspects.get(urn).add(aspectRecord);
     });
 
-    // Add "key" aspects to any non null keys.
-    urnToAspects.keySet().forEach(key -> {
-      final RecordTemplate keyAspect = buildKeyAspect(key);
-      urnToAspects.get(key).add(keyAspect);
-    });
-
     return urnToAspects;
   }
 

diff --git a/metadata-io/src/test/java/com/linkedin/metadata/entity/EbeanEntityServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/entity/EbeanEntityServiceTest.java
@@ -70,12 +70,12 @@ public void testIngestGetEntity() throws Exception {
     assertEquals(2, readEntity.getValue().getCorpUserSnapshot().getAspects().size()); // Key + Info aspect.
     assertTrue(DataTemplateUtil.areEqual(
         writeEntity.getValue().getCorpUserSnapshot().getAspects().get(0),
-        readEntity.getValue().getCorpUserSnapshot().getAspects().get(0)));
+        readEntity.getValue().getCorpUserSnapshot().getAspects().get(1)));
     CorpUserKey expectedKey = new CorpUserKey();
     expectedKey.setUsername("test");
     assertTrue(DataTemplateUtil.areEqual(
         expectedKey,
-        readEntity.getValue().getCorpUserSnapshot().getAspects().get(1).getCorpUserKey()
+        readEntity.getValue().getCorpUserSnapshot().getAspects().get(0).getCorpUserKey()
     )); // Key + Info aspect.
 
     verify(_mockProducer, times(1)).produceMetadataAuditEvent(
@@ -108,25 +108,25 @@ public void testIngestGetEntities() throws Exception {
     assertEquals(2, readEntity1.getValue().getCorpUserSnapshot().getAspects().size()); // Key + Info aspect.
     assertTrue(DataTemplateUtil.areEqual(
         writeEntity1.getValue().getCorpUserSnapshot().getAspects().get(0),
-        readEntity1.getValue().getCorpUserSnapshot().getAspects().get(0)));
+        readEntity1.getValue().getCorpUserSnapshot().getAspects().get(1)));
     CorpUserKey expectedKey1 = new CorpUserKey();
     expectedKey1.setUsername("tester1");
     assertTrue(DataTemplateUtil.areEqual(
         expectedKey1,
-        readEntity1.getValue().getCorpUserSnapshot().getAspects().get(1).getCorpUserKey()
+        readEntity1.getValue().getCorpUserSnapshot().getAspects().get(0).getCorpUserKey()
     )); // Key + Info aspect.
 
     // Entity 2
     com.linkedin.entity.Entity readEntity2 = readEntities.get(entityUrn2);
     assertEquals(2, readEntity2.getValue().getCorpUserSnapshot().getAspects().size()); // Key + Info aspect.
     assertTrue(DataTemplateUtil.areEqual(
         writeEntity2.getValue().getCorpUserSnapshot().getAspects().get(0),
-        readEntity2.getValue().getCorpUserSnapshot().getAspects().get(0)));
+        readEntity2.getValue().getCorpUserSnapshot().getAspects().get(1)));
     CorpUserKey expectedKey2 = new CorpUserKey();
     expectedKey2.setUsername("tester2");
     assertTrue(DataTemplateUtil.areEqual(
         expectedKey2,
-        readEntity2.getValue().getCorpUserSnapshot().getAspects().get(1).getCorpUserKey()
+        readEntity2.getValue().getCorpUserSnapshot().getAspects().get(0).getCorpUserKey()
     )); // Key + Info aspect.
 
     verify(_mockProducer, times(1)).produceMetadataAuditEvent(