From 95b73c4adcd721189a39d63119c2fbea45846c86 Mon Sep 17 00:00:00 2001
From: Gabe Lyons <itsgabelyons@gmail.com>
Date: Fri, 9 Jul 2021 17:04:14 -0700
Subject: [PATCH] fix(frontend): handling null aspects in AspectType (#2860)

---
 .../linkedin/datahub/graphql/types/aspect/AspectType.java | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/aspect/AspectType.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/aspect/AspectType.java
index 21d9bed26a7bb..1a226e189cf38 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/aspect/AspectType.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/aspect/AspectType.java
@@ -6,6 +6,7 @@
 import com.linkedin.entity.client.AspectClient;
 import com.linkedin.metadata.aspect.VersionedAspect;
 import com.linkedin.r2.RemoteInvocationException;
+import com.linkedin.restli.client.RestLiResponseException;
 import graphql.execution.DataFetcherResult;
 import java.util.List;
 import java.util.stream.Collectors;
@@ -31,6 +32,13 @@ public List<DataFetcherResult<Aspect>> batchLoad(@Nonnull List<VersionedAspectKe
           VersionedAspect entity = _aspectClient.getAspect(key.getUrn(), key.getAspectName(), key.getVersion());
           return DataFetcherResult.<Aspect>newResult().data(AspectMapper.map(entity)).build();
         } catch (RemoteInvocationException e) {
+          if (e instanceof RestLiResponseException) {
+            // if no aspect is found, restli will return a 404 rather than null
+            // https://linkedin.github.io/rest.li/user_guide/restli_server#returning-nulls
+            if(((RestLiResponseException) e).getStatus() == 404) {
+              return DataFetcherResult.<Aspect>newResult().data(null).build();
+            }
+          }
           throw new RuntimeException(String.format("Failed to load Aspect for entity %s", key.getUrn()), e);
         }
       }).collect(Collectors.toList());