Return shell entities when they don't exist

metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java

@@ -18,6 +18,7 @@
 import com.linkedin.metadata.snapshot.Snapshot;
 import java.net.URISyntaxException;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -170,8 +171,22 @@ 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) {
-    return getSnapshotUnions(urns, aspectNames).entrySet().stream()
+  /**
+   * Retrieves multiple entities of <b>the same type</b>.
+   *
+   * @param urns set of urns of the same entity type 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();
+    }
+    final String entityType = getEntityTypeFromUrns(urns);
+    final Set<String> aspectsToFetch = aspectNames.isEmpty()
+        ? getEntityAspectNames(entityType)
+        : aspectNames;
+    return getSnapshotUnions(urns, aspectsToFetch).entrySet().stream()
         .collect(Collectors.toMap(Map.Entry::getKey, entry -> toEntity(entry.getValue())));
   }
 
@@ -332,4 +347,18 @@ protected Set<String> getEntityAspectNames(final String entityName) {
   }
 
   public abstract void setWritable();
+
+  private String getEntityTypeFromUrns(final Set<Urn> urns) {
+    String entityType = null;
+    for (final Urn urn : urns) {
+      if (entityType == null) {
+        entityType = urnToEntityName(urn);
+      } else {
+        if (!entityType.equals(urnToEntityName(urn))) {
+          throw new IllegalArgumentException("All entities being retrieved must be of the same entity type.");
+        }
+      }
+    }
+    return entityType;
+  }
 }

metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/EbeanEntityService.java

@@ -53,19 +53,27 @@ public EbeanEntityService(
   public Map<Urn, List<RecordTemplate>> getLatestAspects(@Nonnull final Set<Urn> urns, @Nonnull final Set<String> aspectNames) {
     // Create DB keys
     final Set<EbeanAspectV2.PrimaryKey> dbKeys = urns.stream()
-        .map(urn -> {
-          final Set<String> aspectsToFetch = aspectNames.isEmpty()
-              ? getEntityAspectNames(urn)
-              : aspectNames;
-          return aspectsToFetch.stream()
-              .map(aspectName -> new EbeanAspectV2.PrimaryKey(urn.toString(), aspectName, LATEST_ASPECT_VERSION))
-              .collect(Collectors.toList());
-        })
+        .map(urn -> aspectNames.stream()
+            .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 +82,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;
   }