Skip to content

Commit ddd0d21

Browse files
yoonhyejinshirshankaRyanHolstien
authored
feat(models): update mlflow-related mappers (#12263)
Co-authored-by: Shirshanka Das <[email protected]> Co-authored-by: RyanHolstien <[email protected]>
1 parent e34b2e4 commit ddd0d21

File tree

35 files changed

+2046
-243
lines changed

35 files changed

+2046
-243
lines changed

‎datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java

+36-1
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@
5656
import com.linkedin.datahub.graphql.generated.DataJobInputOutput;
5757
import com.linkedin.datahub.graphql.generated.DataPlatform;
5858
import com.linkedin.datahub.graphql.generated.DataPlatformInstance;
59+
import com.linkedin.datahub.graphql.generated.DataProcessInstance;
5960
import com.linkedin.datahub.graphql.generated.DataQualityContract;
6061
import com.linkedin.datahub.graphql.generated.Dataset;
6162
import com.linkedin.datahub.graphql.generated.DatasetStatsSummary;
@@ -346,6 +347,7 @@
346347
import com.linkedin.datahub.graphql.types.datajob.DataJobType;
347348
import com.linkedin.datahub.graphql.types.dataplatform.DataPlatformType;
348349
import com.linkedin.datahub.graphql.types.dataplatforminstance.DataPlatformInstanceType;
350+
import com.linkedin.datahub.graphql.types.dataprocessinst.DataProcessInstanceType;
349351
import com.linkedin.datahub.graphql.types.dataprocessinst.mappers.DataProcessInstanceRunEventMapper;
350352
import com.linkedin.datahub.graphql.types.dataproduct.DataProductType;
351353
import com.linkedin.datahub.graphql.types.dataset.DatasetType;
@@ -530,6 +532,7 @@ public class GmsGraphQLEngine {
530532
private final FormType formType;
531533
private final IncidentType incidentType;
532534
private final RestrictedType restrictedType;
535+
private final DataProcessInstanceType dataProcessInstanceType;
533536

534537
private final int graphQLQueryComplexityLimit;
535538
private final int graphQLQueryDepthLimit;
@@ -649,6 +652,7 @@ public GmsGraphQLEngine(final GmsGraphQLEngineArgs args) {
649652
this.formType = new FormType(entityClient);
650653
this.incidentType = new IncidentType(entityClient);
651654
this.restrictedType = new RestrictedType(entityClient, restrictedService);
655+
this.dataProcessInstanceType = new DataProcessInstanceType(entityClient, featureFlags);
652656

653657
this.graphQLQueryComplexityLimit = args.graphQLQueryComplexityLimit;
654658
this.graphQLQueryDepthLimit = args.graphQLQueryDepthLimit;
@@ -699,7 +703,8 @@ public GmsGraphQLEngine(final GmsGraphQLEngineArgs args) {
699703
formType,
700704
incidentType,
701705
restrictedType,
702-
businessAttributeType));
706+
businessAttributeType,
707+
dataProcessInstanceType));
703708
this.loadableTypes = new ArrayList<>(entityTypes);
704709
// Extend loadable types with types from the plugins
705710
// This allows us to offer search and browse capabilities out of the box for
@@ -1024,6 +1029,7 @@ private void configureQueryResolvers(final RuntimeWiring.Builder builder) {
10241029
.dataFetcher("tag", getResolver(tagType))
10251030
.dataFetcher("dataFlow", getResolver(dataFlowType))
10261031
.dataFetcher("dataJob", getResolver(dataJobType))
1032+
.dataFetcher("dataProcessInstance", getResolver(dataProcessInstanceType))
10271033
.dataFetcher("glossaryTerm", getResolver(glossaryTermType))
10281034
.dataFetcher("glossaryNode", getResolver(glossaryNodeType))
10291035
.dataFetcher("domain", getResolver((domainType)))
@@ -3058,6 +3064,35 @@ private void configureDataProcessInstanceResolvers(final RuntimeWiring.Builder b
30583064
"DataProcessInstance",
30593065
typeWiring ->
30603066
typeWiring
3067+
.dataFetcher(
3068+
"dataPlatformInstance",
3069+
new LoadableTypeResolver<>(
3070+
dataPlatformInstanceType,
3071+
(env) -> {
3072+
final DataProcessInstance dataProcessInstance = env.getSource();
3073+
return dataProcessInstance.getDataPlatformInstance() != null
3074+
? dataProcessInstance.getDataPlatformInstance().getUrn()
3075+
: null;
3076+
}))
3077+
.dataFetcher(
3078+
"platform",
3079+
new LoadableTypeResolver<>(
3080+
dataPlatformType,
3081+
(env) -> {
3082+
final DataProcessInstance dataProcessInstance = env.getSource();
3083+
return dataProcessInstance.getPlatform() != null
3084+
? dataProcessInstance.getPlatform().getUrn()
3085+
: null;
3086+
}))
3087+
.dataFetcher("parentContainers", new ParentContainersResolver(entityClient))
3088+
.dataFetcher(
3089+
"container",
3090+
new LoadableTypeResolver<>(
3091+
containerType,
3092+
(env) -> {
3093+
final DataProcessInstance dpi = env.getSource();
3094+
return dpi.getContainer() != null ? dpi.getContainer().getUrn() : null;
3095+
}))
30613096
.dataFetcher("relationships", new EntityRelationshipsResultResolver(graphClient))
30623097
.dataFetcher(
30633098
"lineage",

‎datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/DataPlatformInstanceAspectMapper.java

+6
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package com.linkedin.datahub.graphql.types.common.mappers;
22

33
import com.linkedin.datahub.graphql.QueryContext;
4+
import com.linkedin.datahub.graphql.generated.DataPlatform;
45
import com.linkedin.datahub.graphql.generated.DataPlatformInstance;
56
import com.linkedin.datahub.graphql.generated.EntityType;
67
import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
@@ -28,6 +29,11 @@ public DataPlatformInstance apply(
2829
result.setType(EntityType.DATA_PLATFORM_INSTANCE);
2930
result.setUrn(input.getInstance().toString());
3031
}
32+
result.setPlatform(
33+
DataPlatform.builder()
34+
.setUrn(input.getPlatform().toString())
35+
.setType(EntityType.DATA_PLATFORM)
36+
.build());
3137
return result;
3238
}
3339
}

‎datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/TimeStampToAuditStampMapper.java

+24
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
package com.linkedin.datahub.graphql.types.common.mappers;
2+
3+
import com.linkedin.common.TimeStamp;
4+
import com.linkedin.datahub.graphql.QueryContext;
5+
import com.linkedin.datahub.graphql.generated.AuditStamp;
6+
import javax.annotation.Nullable;
7+
8+
public class TimeStampToAuditStampMapper {
9+
10+
public static final TimeStampToAuditStampMapper INSTANCE = new TimeStampToAuditStampMapper();
11+
12+
public static AuditStamp map(
13+
@Nullable final QueryContext context, @Nullable final TimeStamp input) {
14+
if (input == null) {
15+
return null;
16+
}
17+
final AuditStamp result = new AuditStamp();
18+
result.setTime(input.getTime());
19+
if (input.hasActor()) {
20+
result.setActor(input.getActor().toString());
21+
}
22+
return result;
23+
}
24+
}

‎datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/UrnToEntityMapper.java

+6
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import com.linkedin.datahub.graphql.generated.DataJob;
1919
import com.linkedin.datahub.graphql.generated.DataPlatform;
2020
import com.linkedin.datahub.graphql.generated.DataPlatformInstance;
21+
import com.linkedin.datahub.graphql.generated.DataProcessInstance;
2122
import com.linkedin.datahub.graphql.generated.DataProduct;
2223
import com.linkedin.datahub.graphql.generated.Dataset;
2324
import com.linkedin.datahub.graphql.generated.Domain;
@@ -225,6 +226,11 @@ public Entity apply(@Nullable QueryContext context, Urn input) {
225226
((BusinessAttribute) partialEntity).setUrn(input.toString());
226227
((BusinessAttribute) partialEntity).setType(EntityType.BUSINESS_ATTRIBUTE);
227228
}
229+
if (input.getEntityType().equals(DATA_PROCESS_INSTANCE_ENTITY_NAME)) {
230+
partialEntity = new DataProcessInstance();
231+
((DataProcessInstance) partialEntity).setUrn(input.toString());
232+
((DataProcessInstance) partialEntity).setType(EntityType.DATA_PROCESS_INSTANCE);
233+
}
228234
return partialEntity;
229235
}
230236
}

‎datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataprocessinst/DataProcessInstanceType.java

+102
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
package com.linkedin.datahub.graphql.types.dataprocessinst;
2+
3+
import static com.linkedin.metadata.Constants.*;
4+
5+
import com.google.common.collect.ImmutableSet;
6+
import com.linkedin.common.urn.Urn;
7+
import com.linkedin.common.urn.UrnUtils;
8+
import com.linkedin.datahub.graphql.QueryContext;
9+
import com.linkedin.datahub.graphql.featureflags.FeatureFlags;
10+
import com.linkedin.datahub.graphql.generated.DataProcessInstance;
11+
import com.linkedin.datahub.graphql.generated.Entity;
12+
import com.linkedin.datahub.graphql.generated.EntityType;
13+
import com.linkedin.datahub.graphql.types.dataprocessinst.mappers.DataProcessInstanceMapper;
14+
import com.linkedin.entity.EntityResponse;
15+
import com.linkedin.entity.client.EntityClient;
16+
import graphql.execution.DataFetcherResult;
17+
import java.util.ArrayList;
18+
import java.util.HashMap;
19+
import java.util.HashSet;
20+
import java.util.List;
21+
import java.util.Map;
22+
import java.util.Set;
23+
import java.util.function.Function;
24+
import java.util.stream.Collectors;
25+
import javax.annotation.Nonnull;
26+
import lombok.RequiredArgsConstructor;
27+
28+
@RequiredArgsConstructor
29+
public class DataProcessInstanceType
30+
implements com.linkedin.datahub.graphql.types.EntityType<DataProcessInstance, String> {
31+
32+
public static final Set<String> ASPECTS_TO_FETCH =
33+
ImmutableSet.of(
34+
DATA_PROCESS_INSTANCE_KEY_ASPECT_NAME,
35+
DATA_PLATFORM_INSTANCE_ASPECT_NAME,
36+
DATA_PROCESS_INSTANCE_PROPERTIES_ASPECT_NAME,
37+
DATA_PROCESS_INSTANCE_INPUT_ASPECT_NAME,
38+
DATA_PROCESS_INSTANCE_OUTPUT_ASPECT_NAME,
39+
DATA_PROCESS_INSTANCE_RUN_EVENT_ASPECT_NAME,
40+
TEST_RESULTS_ASPECT_NAME,
41+
DATA_PROCESS_INSTANCE_RELATIONSHIPS_ASPECT_NAME,
42+
ML_TRAINING_RUN_PROPERTIES_ASPECT_NAME,
43+
SUB_TYPES_ASPECT_NAME,
44+
CONTAINER_ASPECT_NAME);
45+
46+
private final EntityClient _entityClient;
47+
private final FeatureFlags _featureFlags;
48+
49+
@Override
50+
public EntityType type() {
51+
return EntityType.DATA_PROCESS_INSTANCE;
52+
}
53+
54+
@Override
55+
public Function<Entity, String> getKeyProvider() {
56+
return Entity::getUrn;
57+
}
58+
59+
@Override
60+
public Class<DataProcessInstance> objectClass() {
61+
return DataProcessInstance.class;
62+
}
63+
64+
@Override
65+
public List<DataFetcherResult<DataProcessInstance>> batchLoad(
66+
@Nonnull List<String> urns, @Nonnull QueryContext context) throws Exception {
67+
final List<Urn> dataProcessInstanceUrns =
68+
urns.stream().map(UrnUtils::getUrn).collect(Collectors.toList());
69+
70+
try {
71+
Map<Urn, EntityResponse> entities = new HashMap<>();
72+
if (_featureFlags.isDataProcessInstanceEntityEnabled()) {
73+
entities =
74+
_entityClient.batchGetV2(
75+
context.getOperationContext(),
76+
DATA_PROCESS_INSTANCE_ENTITY_NAME,
77+
new HashSet<>(dataProcessInstanceUrns),
78+
ASPECTS_TO_FETCH);
79+
}
80+
81+
final List<EntityResponse> gmsResults = new ArrayList<>();
82+
for (Urn urn : dataProcessInstanceUrns) {
83+
if (_featureFlags.isDataProcessInstanceEntityEnabled()) {
84+
gmsResults.add(entities.getOrDefault(urn, null));
85+
}
86+
}
87+
88+
return gmsResults.stream()
89+
.map(
90+
gmsResult ->
91+
gmsResult == null
92+
? null
93+
: DataFetcherResult.<DataProcessInstance>newResult()
94+
.data(DataProcessInstanceMapper.map(context, gmsResult))
95+
.build())
96+
.collect(Collectors.toList());
97+
98+
} catch (Exception e) {
99+
throw new RuntimeException("Failed to load Data Process Instance entity", e);
100+
}
101+
}
102+
}

0 commit comments

Comments
 (0)