feat(graphql): More forgiving for unknown data platforms during reads (…

…#2185) Co-authored-by: John Joyce <[email protected]>
datahub-project · Mar 7, 2021 · 6756c2d · 6756c2d
1 parent 3537b2e
commit 6756c2d
Showing 1 changed file with 14 additions and 8 deletions.
diff --git a/...-core/src/main/java/com/linkedin/datahub/graphql/types/dataplatform/DataPlatformType.java b/...-core/src/main/java/com/linkedin/datahub/graphql/types/dataplatform/DataPlatformType.java
@@ -3,10 +3,11 @@
 import com.linkedin.common.urn.DataPlatformUrn;
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.types.EntityType;
-import com.linkedin.datahub.graphql.types.mappers.DataPlatformInfoMapper;
 import com.linkedin.datahub.graphql.generated.DataPlatform;
+import com.linkedin.datahub.graphql.types.mappers.DataPlatformMapper;
 import com.linkedin.dataplatform.client.DataPlatforms;
 
+import java.net.URISyntaxException;
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
@@ -30,16 +31,11 @@ public List<DataPlatform> batchLoad(final List<String> urns, final QueryContext
         try {
             if (_urnToPlatform == null) {
                 _urnToPlatform = _dataPlatformsClient.getAllPlatforms().stream()
-                        .filter(com.linkedin.dataPlatforms.DataPlatform::hasDataPlatformInfo)
-                        .map(platform -> new DataPlatform(
-                                new DataPlatformUrn(platform.getName()).toString(),
-                                com.linkedin.datahub.graphql.generated.EntityType.DATA_PLATFORM,
-                                platform.getName(),
-                                DataPlatformInfoMapper.map(platform.getDataPlatformInfo())))
+                        .map(DataPlatformMapper::map)
                         .collect(Collectors.toMap(DataPlatform::getUrn, platform -> platform));
             }
             return urns.stream()
-                    .map(key -> _urnToPlatform.get(key))
+                    .map(key -> _urnToPlatform.containsKey(key) ? _urnToPlatform.get(key) : getUnknownDataPlatform(key))
                     .collect(Collectors.toList());
         } catch (Exception e) {
             throw new RuntimeException("Failed to batch load DataPlatforms", e);
@@ -50,4 +46,14 @@ public List<DataPlatform> batchLoad(final List<String> urns, final QueryContext
     public com.linkedin.datahub.graphql.generated.EntityType type() {
         return com.linkedin.datahub.graphql.generated.EntityType.DATA_PLATFORM;
     }
+
+    private DataPlatform getUnknownDataPlatform(final String urnStr) {
+        try {
+            final com.linkedin.dataPlatforms.DataPlatform platform = new com.linkedin.dataPlatforms.DataPlatform()
+                    .setName(DataPlatformUrn.createFromString(urnStr).getPlatformNameEntity());
+            return DataPlatformMapper.map(platform);
+        } catch (URISyntaxException e) {
+            throw new RuntimeException(String.format("Invalid DataPlatformUrn %s provided", urnStr), e);
+        }
+    }
 }