From 3d9a9541f1ff37ea80dd9d7d44fe501909269495 Mon Sep 17 00:00:00 2001
From: david-leifker <114954101+david-leifker@users.noreply.github.com>
Date: Thu, 8 Aug 2024 13:54:14 -0500
Subject: [PATCH] fix(entiy-client): handle null entityUrn case for restli
 (#11122)

---
 .../entity/client/RestliEntityClient.java       | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/RestliEntityClient.java b/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/RestliEntityClient.java
index 2a3ae5d006..780c6c6a00 100644
--- a/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/RestliEntityClient.java
+++ b/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/RestliEntityClient.java
@@ -50,6 +50,7 @@ import com.linkedin.metadata.aspect.VersionedAspect;
 import com.linkedin.metadata.browse.BrowseResult;
 import com.linkedin.metadata.browse.BrowseResultV2;
 import com.linkedin.metadata.graph.LineageDirection;
+import com.linkedin.metadata.models.EntitySpec;
 import com.linkedin.metadata.query.AutoCompleteResult;
 import com.linkedin.metadata.query.LineageFlags;
 import com.linkedin.metadata.query.ListResult;
@@ -66,6 +67,7 @@ import com.linkedin.metadata.search.LineageScrollResult;
 import com.linkedin.metadata.search.LineageSearchResult;
 import com.linkedin.metadata.search.ScrollResult;
 import com.linkedin.metadata.search.SearchResult;
+import com.linkedin.metadata.utils.EntityKeyUtils;
 import com.linkedin.mxe.MetadataChangeProposal;
 import com.linkedin.mxe.MetadataChangeProposalArray;
 import com.linkedin.mxe.PlatformEvent;
@@ -1063,7 +1065,20 @@ public class RestliEntityClient extends BaseClient implements EntityClient {
     String result =
         sendClientRequest(requestBuilder, opContext.getSessionAuthentication()).getEntity();
     return metadataChangeProposals.stream()
-        .map(proposal -> "success".equals(result) ? proposal.getEntityUrn().toString() : null)
+        .map(
+            proposal -> {
+              if ("success".equals(result)) {
+                if (proposal.getEntityUrn() != null) {
+                  return proposal.getEntityUrn().toString();
+                } else {
+                  EntitySpec entitySpec =
+                      opContext.getEntityRegistry().getEntitySpec(proposal.getEntityType());
+                  return EntityKeyUtils.getUrnFromProposal(proposal, entitySpec.getKeyAspectSpec())
+                      .toString();
+                }
+              }
+              return null;
+            })
         .collect(Collectors.toList());
   }