package com.linkedin.datahub.graphql.types.dataplatform;

import com.linkedin.common.urn.DataPlatformUrn;
import com.linkedin.datahub.graphql.QueryContext;
import com.linkedin.datahub.graphql.types.EntityType;
import com.linkedin.datahub.graphql.generated.DataPlatform;
import com.linkedin.datahub.graphql.types.dataplatform.mappers.DataPlatformMapper;
import com.linkedin.dataplatform.client.DataPlatforms;

import graphql.execution.DataFetcherResult;
import java.net.URISyntaxException;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

public class DataPlatformType implements EntityType<DataPlatform> {

    private final DataPlatforms _dataPlatformsClient;

    public DataPlatformType(final DataPlatforms dataPlatformsClient) {
        _dataPlatformsClient = dataPlatformsClient;
    }

    @Override
    public Class<DataPlatform> objectClass() {
        return DataPlatform.class;
    }

    @Override
    public List<DataFetcherResult<DataPlatform>> batchLoad(final List<String> urns, final QueryContext context) {
        try {
            Map<String, DataPlatform> urnToPlatform = _dataPlatformsClient.getAllPlatforms()
                .stream()
                .map(DataPlatformMapper::map)
                .collect(Collectors.toMap(DataPlatform::getUrn, platform -> platform));
            return urns.stream()
                    .map(key -> urnToPlatform.containsKey(key) ? urnToPlatform.get(key) : getUnknownDataPlatform(key))
                .map(dataPlatform -> DataFetcherResult.<DataPlatform>newResult().data(dataPlatform).build())
                    .collect(Collectors.toList());
        } catch (Exception e) {
            throw new RuntimeException("Failed to batch load DataPlatforms", e);
        }
    }

    @Override
    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);
        }
    }
}

@@ -686,16 +686,70 @@
 sClient;
+%0A    private Map%3CString, DataPlatform%3E _urnToPlatform;
 %0A%0A    pu
@@ -1133,34 +1133,55 @@
     
-Map%3CString, DataPlatform%3E 
+if (_urnToPlatform == null) %7B%0A                _
 urnT
@@ -1230,33 +1230,16 @@
 tforms()
-%0A                
 .stream(
@@ -1252,24 +1252,32 @@
             
+       
+ 
 .map(DataPla
@@ -1294,16 +1294,24 @@
 r::map)%0A
+        
         
@@ -1386,24 +1386,38 @@
 platform));%0A
+            %7D%0A
             
@@ -1469,16 +1469,17 @@
 (key -%3E 
+_
 urnToPla
@@ -1503,16 +1503,17 @@
 (key) ? 
+_
 urnToPla

Revert "fix(gql): removing data platform caching in gql (#2947)" (#2956)