yujunjun/datahub
2
0
Fork 0
mirror of https://github.com/datahub-project/datahub.git synced 2025-12-25 00:48:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(graphql): fetching data platforms using standard procedure (#2998)

Browse Source
This commit is contained in:
Gabe Lyons 2021-08-02 11:35:09 -07:00 committed by GitHub
parent 8e9b5f4b7a
commit 9b0f8bc0d3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 69 additions and 38 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsClientFactory.java
View File

@ -41,17 +41,6 @@ public class GmsClientFactory {
private GmsClientFactory() { }
public static DataPlatforms getDataPlatformsClient() {
if (_dataPlatforms == null) {
synchronized (GmsClientFactory.class) {
if (_dataPlatforms == null) {
_dataPlatforms = new DataPlatforms(REST_CLIENT);
}
}
}
return _dataPlatforms;
}
public static Lineages getLineagesClient() {
if (_lineages == null) {
synchronized (GmsClientFactory.class) {

2
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java
View File

@ -112,7 +112,7 @@ public class GmsGraphQLEngine {
public static final CorpGroupType CORP_GROUP_TYPE = new CorpGroupType(GmsClientFactory.getEntitiesClient());
public static final ChartType CHART_TYPE = new ChartType(GmsClientFactory.getEntitiesClient());
public static final DashboardType DASHBOARD_TYPE = new DashboardType(GmsClientFactory.getEntitiesClient());
public static final DataPlatformType DATA_PLATFORM_TYPE = new DataPlatformType(GmsClientFactory.getDataPlatformsClient());
public static final DataPlatformType DATA_PLATFORM_TYPE = new DataPlatformType(GmsClientFactory.getEntitiesClient());
public static final DownstreamLineageType DOWNSTREAM_LINEAGE_TYPE = new DownstreamLineageType(
GmsClientFactory.getLineagesClient()
);

63
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataplatform/DataPlatformType.java
View File

@ -1,25 +1,27 @@
package com.linkedin.datahub.graphql.types.dataplatform;
import com.linkedin.common.urn.DataPlatformUrn;
import com.linkedin.common.urn.Urn;
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 com.linkedin.datahub.graphql.types.dataplatform.mappers.DataPlatformSnapshotMapper;
import com.linkedin.entity.client.EntityClient;
import com.linkedin.metadata.extractor.SnapshotToAspectMap;
import graphql.execution.DataFetcherResult;
import java.net.URISyntaxException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;
public class DataPlatformType implements EntityType<DataPlatform> {
private final DataPlatforms _dataPlatformsClient;
private Map<String, DataPlatform> _urnToPlatform;
private final EntityClient _entityClient;
public DataPlatformType(final DataPlatforms dataPlatformsClient) {
_dataPlatformsClient = dataPlatformsClient;
public DataPlatformType(final EntityClient entityClient) {
_entityClient = entityClient;
}
@Override
@ -29,18 +31,37 @@ public class DataPlatformType implements EntityType<DataPlatform> {
@Override
public List<DataFetcherResult<DataPlatform>> batchLoad(final List<String> urns, final QueryContext context) {
final List<Urn> dataPlatformUrns = urns.stream()
.map(urnStr -> {
try {
return Urn.createFromString(urnStr);
} catch (URISyntaxException e) {
throw new RuntimeException(String.format("Failed to retrieve entity with urn %s", urnStr));
}
})
.collect(Collectors.toList());
try {
if (_urnToPlatform == null) {
_urnToPlatform = _dataPlatformsClient.getAllPlatforms().stream()
.map(DataPlatformMapper::map)
.collect(Collectors.toMap(DataPlatform::getUrn, platform -> platform));
final Map<Urn, com.linkedin.entity.Entity> dataPlatformMap = _entityClient.batchGet(dataPlatformUrns
.stream()
.filter(Objects::nonNull)
.collect(Collectors.toSet()));
final List<com.linkedin.entity.Entity> gmsResults = new ArrayList<>();
for (Urn urn : dataPlatformUrns) {
gmsResults.add(dataPlatformMap.getOrDefault(urn, null));
}
return urns.stream()
.map(key -> _urnToPlatform.containsKey(key) ? _urnToPlatform.get(key) : getUnknownDataPlatform(key))
.map(dataPlatform -> DataFetcherResult.<DataPlatform>newResult().data(dataPlatform).build())
.collect(Collectors.toList());
return gmsResults.stream()
.map(gmsPlatform -> gmsPlatform == null ? null
: DataFetcherResult.<DataPlatform>newResult()
.data(DataPlatformSnapshotMapper.map(gmsPlatform.getValue().getDataPlatformSnapshot()))
.localContext(SnapshotToAspectMap.extractAspectMap(gmsPlatform.getValue().getDataPlatformSnapshot()))
.build())
.collect(Collectors.toList());
} catch (Exception e) {
throw new RuntimeException("Failed to batch load DataPlatforms", e);
throw new RuntimeException("Failed to batch load Data Platforms", e);
}
}
@ -48,14 +69,4 @@ public class DataPlatformType implements EntityType<DataPlatform> {
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);
}
}
}

31
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataplatform/mappers/DataPlatformSnapshotMapper.java Normal file
View File

@ -0,0 +1,31 @@
package com.linkedin.datahub.graphql.types.dataplatform.mappers;
import com.linkedin.datahub.graphql.generated.DataPlatform;
import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
import com.linkedin.metadata.aspect.DataPlatformAspect;
import com.linkedin.metadata.snapshot.DataPlatformSnapshot;
import javax.annotation.Nonnull;
public class DataPlatformSnapshotMapper implements ModelMapper<DataPlatformSnapshot, DataPlatform> {
public static final DataPlatformSnapshotMapper INSTANCE = new DataPlatformSnapshotMapper();
public static DataPlatform map(@Nonnull final DataPlatformSnapshot platform) {
return INSTANCE.apply(platform);
}
@Override
public DataPlatform apply(@Nonnull final DataPlatformSnapshot input) {
final DataPlatform result = new DataPlatform();
result.setUrn(input.getUrn().toString());
result.setName(input.getUrn().getPlatformNameEntity());
for (DataPlatformAspect aspect : input.getAspects()) {
if (aspect.isDataPlatformInfo()) {
result.setInfo(DataPlatformInfoMapper.map(aspect.getDataPlatformInfo()));
}
}
return result;
}
}