From 4732694780e12350edc03d68385fd982288897f9 Mon Sep 17 00:00:00 2001
From: Pedro Silva <pedro@acryl.io>
Date: Thu, 2 Feb 2023 19:45:44 +0000
Subject: [PATCH] fix(gms): Corrects MCP generation in async mode (#7214)

Co-authored-by: John Joyce <john@acryl.io>
---
 .../linkedin/metadata/dao/producer/KafkaEventProducer.java | 7 ++-----
 .../java/com/linkedin/metadata/entity/EntityService.java   | 2 +-
 .../java/com/linkedin/metadata/event/EventProducer.java    | 3 ++-
 .../com/linkedin/metadata/entity/EntityServiceTest.java    | 6 ++++--
 4 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/metadata-dao-impl/kafka-producer/src/main/java/com/linkedin/metadata/dao/producer/KafkaEventProducer.java b/metadata-dao-impl/kafka-producer/src/main/java/com/linkedin/metadata/dao/producer/KafkaEventProducer.java
index 40db7100cc..95307c2a56 100644
--- a/metadata-dao-impl/kafka-producer/src/main/java/com/linkedin/metadata/dao/producer/KafkaEventProducer.java
+++ b/metadata-dao-impl/kafka-producer/src/main/java/com/linkedin/metadata/dao/producer/KafkaEventProducer.java
@@ -118,13 +118,10 @@ public class KafkaEventProducer implements EventProducer {
 
   @Override
   @WithSpan
-  public void produceMetadataChangeProposal(@Nonnull final MetadataChangeProposal metadataChangeProposal) {
+  public void produceMetadataChangeProposal(@Nonnull final Urn urn, @Nonnull final MetadataChangeProposal
+      metadataChangeProposal) {
     GenericRecord record;
 
-    Urn urn = metadataChangeProposal.getEntityUrn();
-    if (urn == null) {
-      throw new IllegalArgumentException("Urn for proposal cannot be null.");
-    }
     try {
       log.debug(String.format("Converting Pegasus snapshot to Avro snapshot urn %s\nMetadataChangeProposal: %s",
           urn,
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
index c2c0bd9014..dd319a33f1 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java
@@ -882,7 +882,7 @@ private Map<Urn, List<EnvelopedAspect>> getCorrespondingAspects(Set<EntityAspect
         newSystemMetadata = result != null ? result.getNewSystemMetadata() : null;
       } else {
         // When async is turned on, we write to proposal log and return without waiting
-        _producer.produceMetadataChangeProposal(mcp);
+        _producer.produceMetadataChangeProposal(entityUrn, mcp);
         return new IngestProposalResult(mcp.getEntityUrn(), false, true);
       }
     } else {
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/event/EventProducer.java b/metadata-io/src/main/java/com/linkedin/metadata/event/EventProducer.java
index fafaaefa5e..eafa8191f4 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/event/EventProducer.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/event/EventProducer.java
@@ -63,7 +63,8 @@ public interface EventProducer {
    * @param metadataChangeProposal metadata change proposal to push into MCP kafka topic
    */
   @WithSpan
-  void produceMetadataChangeProposal(@Nonnull MetadataChangeProposal metadataChangeProposal);
+  void produceMetadataChangeProposal(@Nonnull final Urn urn, @Nonnull final MetadataChangeProposal
+      metadataChangeProposal);
 
   /**
    * Produces a generic platform "event".
diff --git a/metadata-io/src/test/java/com/linkedin/metadata/entity/EntityServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/entity/EntityServiceTest.java
index 0ab2fa34b8..4659fa65a9 100644
--- a/metadata-io/src/test/java/com/linkedin/metadata/entity/EntityServiceTest.java
+++ b/metadata-io/src/test/java/com/linkedin/metadata/entity/EntityServiceTest.java
@@ -461,7 +461,8 @@ abstract public class EntityServiceTest<T_AD extends AspectDao, T_RS extends Ret
         _entityService.ingestProposal(gmce, TEST_AUDIT_STAMP, true);
         verify(_mockProducer, times(0)).produceMetadataChangeLog(Mockito.eq(entityUrn),
             Mockito.any(), Mockito.any());
-        verify(_mockProducer, times(1)).produceMetadataChangeProposal(Mockito.eq(gmce));
+        verify(_mockProducer, times(1)).produceMetadataChangeProposal(Mockito.eq(entityUrn),
+            Mockito.eq(gmce));
     }
 
 
@@ -486,7 +487,8 @@ abstract public class EntityServiceTest<T_AD extends AspectDao, T_RS extends Ret
         _entityService.ingestProposal(gmce, TEST_AUDIT_STAMP, true);
         verify(_mockProducer, times(1)).produceMetadataChangeLog(Mockito.eq(entityUrn),
             Mockito.any(), Mockito.any());
-        verify(_mockProducer, times(0)).produceMetadataChangeProposal(Mockito.eq(gmce));
+        verify(_mockProducer, times(0)).produceMetadataChangeProposal(Mockito.eq(entityUrn),
+            Mockito.eq(gmce));
     }
 
     @Test