more fixes

datahub-project · Aug 9, 2024 · 4da3865 · 4da3865
1 parent 9471759
commit 4da3865
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 9 deletions.
diff --git a/...-registry/src/main/java/com/linkedin/metadata/models/registry/SnapshotEntityRegistry.java b/...-registry/src/main/java/com/linkedin/metadata/models/registry/SnapshotEntityRegistry.java
@@ -40,7 +40,6 @@
 import javax.annotation.Nullable;
 import lombok.Getter;
 
-
 /**
  * Implementation of {@link EntityRegistry} that builds {@link DefaultEntitySpec} objects from the a
  * {@link Snapshot} Record Template present on the classpath
@@ -52,9 +51,7 @@ public class SnapshotEntityRegistry implements EntityRegistry {
   private final AspectTemplateEngine _aspectTemplateEngine;
   private final Map<String, AspectSpec> _aspectNameToSpec;
 
-
-  @Getter
-  @Nullable
+  @Getter @Nullable
   private BiFunction<PluginConfiguration, List<ClassLoader>, PluginFactory> pluginFactoryProvider;
 
   private static final SnapshotEntityRegistry INSTANCE = new SnapshotEntityRegistry();
@@ -72,9 +69,10 @@ public SnapshotEntityRegistry() {
 
   public SnapshotEntityRegistry(
       BiFunction<PluginConfiguration, List<ClassLoader>, PluginFactory> pluginFactoryProvider) {
-    entityNameToSpec = new EntitySpecBuilder().buildEntitySpecs(new Snapshot().schema())
-        .stream()
-        .collect(Collectors.toMap(spec -> spec.getName().toLowerCase(), spec -> spec));
+    entityNameToSpec =
+        new EntitySpecBuilder()
+            .buildEntitySpecs(new Snapshot().schema()).stream()
+                .collect(Collectors.toMap(spec -> spec.getName().toLowerCase(), spec -> spec));
     entitySpecs = new ArrayList<>(entityNameToSpec.values());
     _aspectNameToSpec = populateAspectMap(entitySpecs);
     _aspectTemplateEngine = populateTemplateEngine(_aspectNameToSpec);

diff --git a/...adata-io-api/src/main/java/com/linkedin/metadata/entity/ebean/batch/AspectsBatchImpl.java b/...adata-io-api/src/main/java/com/linkedin/metadata/entity/ebean/batch/AspectsBatchImpl.java
@@ -48,7 +48,7 @@ public Pair<Map<String, Set<String>>, List<ChangeMCP>> toUpsertBatchItems(
       final Map<String, Map<String, SystemAspect>> latestAspects) {
 
     // Process proposals to change items
-    Stream<ChangeMCP> mutatedProposalsStream =
+    Stream<? extends BatchItem> mutatedProposalsStream =
         proposedItemsToChangeItemStream(
             items.stream()
                 .filter(item -> item instanceof ProposedItem)
@@ -93,8 +93,9 @@ public Pair<Map<String, Set<String>>, List<ChangeMCP>> toUpsertBatchItems(
 
     LinkedList<ChangeMCP> newItems =
         applyMCPSideEffects(upsertBatchItems).collect(Collectors.toCollection(LinkedList::new));
-    Map<String, Set<String>> newUrnAspectNames = getNewUrnAspectsMap(getUrnAspectsMap(), newItems);
     upsertBatchItems.addAll(newItems);
+    Map<String, Set<String>> newUrnAspectNames =
+        getNewUrnAspectsMap(getUrnAspectsMap(), upsertBatchItems);
 
     return Pair.of(newUrnAspectNames, upsertBatchItems);
   }