Skip to content

Commit 92581f0

Browse files
jayacryljayacryl
authored
fix(gql) add incident assignee owner type resolver (datahub-project#12897)
1 parent 85a5b5c commit 92581f0

File tree

3 files changed

+266
-102
lines changed

3 files changed

+266
-102
lines changed

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

+9-102
Original file line numberDiff line numberDiff line change
@@ -24,108 +24,7 @@
2424
import com.linkedin.datahub.graphql.analytics.service.AnalyticsService;
2525
import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils;
2626
import com.linkedin.datahub.graphql.featureflags.FeatureFlags;
27-
import com.linkedin.datahub.graphql.generated.AccessToken;
28-
import com.linkedin.datahub.graphql.generated.AccessTokenMetadata;
29-
import com.linkedin.datahub.graphql.generated.ActorFilter;
30-
import com.linkedin.datahub.graphql.generated.AggregationMetadata;
31-
import com.linkedin.datahub.graphql.generated.Assertion;
32-
import com.linkedin.datahub.graphql.generated.AutoCompleteResultForEntity;
33-
import com.linkedin.datahub.graphql.generated.AutoCompleteResults;
34-
import com.linkedin.datahub.graphql.generated.BrowsePathEntry;
35-
import com.linkedin.datahub.graphql.generated.BrowseResultGroupV2;
36-
import com.linkedin.datahub.graphql.generated.BrowseResults;
37-
import com.linkedin.datahub.graphql.generated.BusinessAttribute;
38-
import com.linkedin.datahub.graphql.generated.BusinessAttributeAssociation;
39-
import com.linkedin.datahub.graphql.generated.Chart;
40-
import com.linkedin.datahub.graphql.generated.ChartInfo;
41-
import com.linkedin.datahub.graphql.generated.Container;
42-
import com.linkedin.datahub.graphql.generated.CorpGroup;
43-
import com.linkedin.datahub.graphql.generated.CorpGroupInfo;
44-
import com.linkedin.datahub.graphql.generated.CorpUser;
45-
import com.linkedin.datahub.graphql.generated.CorpUserEditableProperties;
46-
import com.linkedin.datahub.graphql.generated.CorpUserInfo;
47-
import com.linkedin.datahub.graphql.generated.CorpUserViewsSettings;
48-
import com.linkedin.datahub.graphql.generated.Dashboard;
49-
import com.linkedin.datahub.graphql.generated.DashboardInfo;
50-
import com.linkedin.datahub.graphql.generated.DashboardStatsSummary;
51-
import com.linkedin.datahub.graphql.generated.DashboardUserUsageCounts;
52-
import com.linkedin.datahub.graphql.generated.DataFlow;
53-
import com.linkedin.datahub.graphql.generated.DataHubConnection;
54-
import com.linkedin.datahub.graphql.generated.DataHubView;
55-
import com.linkedin.datahub.graphql.generated.DataJob;
56-
import com.linkedin.datahub.graphql.generated.DataJobInputOutput;
57-
import com.linkedin.datahub.graphql.generated.DataPlatform;
58-
import com.linkedin.datahub.graphql.generated.DataPlatformInstance;
59-
import com.linkedin.datahub.graphql.generated.DataProcessInstance;
60-
import com.linkedin.datahub.graphql.generated.DataQualityContract;
61-
import com.linkedin.datahub.graphql.generated.Dataset;
62-
import com.linkedin.datahub.graphql.generated.DatasetStatsSummary;
63-
import com.linkedin.datahub.graphql.generated.Deprecation;
64-
import com.linkedin.datahub.graphql.generated.Domain;
65-
import com.linkedin.datahub.graphql.generated.ERModelRelationship;
66-
import com.linkedin.datahub.graphql.generated.ERModelRelationshipProperties;
67-
import com.linkedin.datahub.graphql.generated.Entity;
68-
import com.linkedin.datahub.graphql.generated.EntityPath;
69-
import com.linkedin.datahub.graphql.generated.EntityRelationship;
70-
import com.linkedin.datahub.graphql.generated.EntityRelationshipLegacy;
71-
import com.linkedin.datahub.graphql.generated.FacetMetadata;
72-
import com.linkedin.datahub.graphql.generated.ForeignKeyConstraint;
73-
import com.linkedin.datahub.graphql.generated.FormActorAssignment;
74-
import com.linkedin.datahub.graphql.generated.FreshnessContract;
75-
import com.linkedin.datahub.graphql.generated.GetRootGlossaryNodesResult;
76-
import com.linkedin.datahub.graphql.generated.GetRootGlossaryTermsResult;
77-
import com.linkedin.datahub.graphql.generated.GlossaryNode;
78-
import com.linkedin.datahub.graphql.generated.GlossaryTerm;
79-
import com.linkedin.datahub.graphql.generated.GlossaryTermAssociation;
80-
import com.linkedin.datahub.graphql.generated.IncidentSource;
81-
import com.linkedin.datahub.graphql.generated.IngestionSource;
82-
import com.linkedin.datahub.graphql.generated.InstitutionalMemoryMetadata;
83-
import com.linkedin.datahub.graphql.generated.LineageRelationship;
84-
import com.linkedin.datahub.graphql.generated.ListAccessTokenResult;
85-
import com.linkedin.datahub.graphql.generated.ListBusinessAttributesResult;
86-
import com.linkedin.datahub.graphql.generated.ListDomainsResult;
87-
import com.linkedin.datahub.graphql.generated.ListGroupsResult;
88-
import com.linkedin.datahub.graphql.generated.ListOwnershipTypesResult;
89-
import com.linkedin.datahub.graphql.generated.ListQueriesResult;
90-
import com.linkedin.datahub.graphql.generated.ListTestsResult;
91-
import com.linkedin.datahub.graphql.generated.ListViewsResult;
92-
import com.linkedin.datahub.graphql.generated.MLFeature;
93-
import com.linkedin.datahub.graphql.generated.MLFeatureProperties;
94-
import com.linkedin.datahub.graphql.generated.MLFeatureTable;
95-
import com.linkedin.datahub.graphql.generated.MLFeatureTableProperties;
96-
import com.linkedin.datahub.graphql.generated.MLModel;
97-
import com.linkedin.datahub.graphql.generated.MLModelGroup;
98-
import com.linkedin.datahub.graphql.generated.MLModelProperties;
99-
import com.linkedin.datahub.graphql.generated.MLPrimaryKey;
100-
import com.linkedin.datahub.graphql.generated.MLPrimaryKeyProperties;
101-
import com.linkedin.datahub.graphql.generated.MatchedField;
102-
import com.linkedin.datahub.graphql.generated.MetadataAttribution;
103-
import com.linkedin.datahub.graphql.generated.Notebook;
104-
import com.linkedin.datahub.graphql.generated.Owner;
105-
import com.linkedin.datahub.graphql.generated.OwnershipTypeEntity;
106-
import com.linkedin.datahub.graphql.generated.ParentDomainsResult;
107-
import com.linkedin.datahub.graphql.generated.PolicyMatchCriterionValue;
108-
import com.linkedin.datahub.graphql.generated.QueryEntity;
109-
import com.linkedin.datahub.graphql.generated.QueryProperties;
110-
import com.linkedin.datahub.graphql.generated.QuerySubject;
111-
import com.linkedin.datahub.graphql.generated.QuickFilter;
112-
import com.linkedin.datahub.graphql.generated.RecommendationContent;
113-
import com.linkedin.datahub.graphql.generated.ResolvedAuditStamp;
114-
import com.linkedin.datahub.graphql.generated.SchemaContract;
115-
import com.linkedin.datahub.graphql.generated.SchemaField;
116-
import com.linkedin.datahub.graphql.generated.SchemaFieldEntity;
117-
import com.linkedin.datahub.graphql.generated.SearchAcrossLineageResult;
118-
import com.linkedin.datahub.graphql.generated.SearchResult;
119-
import com.linkedin.datahub.graphql.generated.SiblingProperties;
120-
import com.linkedin.datahub.graphql.generated.StructuredPropertiesEntry;
121-
import com.linkedin.datahub.graphql.generated.StructuredPropertyDefinition;
122-
import com.linkedin.datahub.graphql.generated.StructuredPropertyParams;
123-
import com.linkedin.datahub.graphql.generated.Test;
124-
import com.linkedin.datahub.graphql.generated.TestResult;
125-
import com.linkedin.datahub.graphql.generated.TypeQualifier;
126-
import com.linkedin.datahub.graphql.generated.UserUsageCounts;
127-
import com.linkedin.datahub.graphql.generated.VersionProperties;
128-
import com.linkedin.datahub.graphql.generated.VersionSet;
27+
import com.linkedin.datahub.graphql.generated.*;
12928
import com.linkedin.datahub.graphql.resolvers.MeResolver;
13029
import com.linkedin.datahub.graphql.resolvers.assertion.AssertionRunEventResolver;
13130
import com.linkedin.datahub.graphql.resolvers.assertion.DeleteAssertionResolver;
@@ -232,6 +131,7 @@
232131
import com.linkedin.datahub.graphql.resolvers.load.EntityTypeResolver;
233132
import com.linkedin.datahub.graphql.resolvers.load.LoadableTypeBatchResolver;
234133
import com.linkedin.datahub.graphql.resolvers.load.LoadableTypeResolver;
134+
import com.linkedin.datahub.graphql.resolvers.load.OwnerTypeBatchResolver;
235135
import com.linkedin.datahub.graphql.resolvers.load.OwnerTypeResolver;
236136
import com.linkedin.datahub.graphql.resolvers.load.TimeSeriesAspectResolver;
237137
import com.linkedin.datahub.graphql.resolvers.mutate.AddLinkResolver;
@@ -1767,6 +1667,13 @@ private void configureDatasetResolvers(final RuntimeWiring.Builder builder) {
17671667
"owner",
17681668
new OwnerTypeResolver<>(
17691669
ownerTypes, (env) -> ((Owner) env.getSource()).getOwner())))
1670+
.type(
1671+
"Incident",
1672+
typeWiring ->
1673+
typeWiring.dataFetcher(
1674+
"assignees",
1675+
new OwnerTypeBatchResolver(
1676+
ownerTypes, (env) -> ((Incident) env.getSource()).getAssignees())))
17701677
.type(
17711678
"SchemaField",
17721679
typeWiring ->

‎datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/OwnerTypeBatchResolver.java

+90
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
package com.linkedin.datahub.graphql.resolvers.load;
2+
3+
import com.google.common.collect.Iterables;
4+
import com.linkedin.datahub.graphql.generated.Entity;
5+
import com.linkedin.datahub.graphql.generated.OwnerType;
6+
import com.linkedin.datahub.graphql.types.LoadableType;
7+
import graphql.schema.DataFetcher;
8+
import graphql.schema.DataFetchingEnvironment;
9+
import java.util.ArrayList;
10+
import java.util.Collections;
11+
import java.util.HashMap;
12+
import java.util.List;
13+
import java.util.Map;
14+
import java.util.concurrent.CompletableFuture;
15+
import java.util.function.Function;
16+
import java.util.stream.Collectors;
17+
import org.dataloader.DataLoader;
18+
import org.dataloader.DataLoaderRegistry;
19+
20+
/**
21+
* GraphQL resolver responsible for
22+
*
23+
* <p>1. Retrieving a batch of OwnerTypes. 2. Resolving their Entities
24+
*/
25+
public class OwnerTypeBatchResolver implements DataFetcher<CompletableFuture<List<OwnerType>>> {
26+
27+
private final List<com.linkedin.datahub.graphql.types.LoadableType<?, ?>> _loadableTypes;
28+
private final Function<DataFetchingEnvironment, List<OwnerType>> _typesProvider;
29+
30+
public OwnerTypeBatchResolver(
31+
final List<com.linkedin.datahub.graphql.types.LoadableType<?, ?>> loadableTypes,
32+
final Function<DataFetchingEnvironment, List<OwnerType>> typesProvider) {
33+
_loadableTypes = loadableTypes;
34+
_typesProvider = typesProvider;
35+
}
36+
37+
@Override
38+
public CompletableFuture<List<OwnerType>> get(DataFetchingEnvironment environment) {
39+
// 1. get OwnerTypes
40+
final List<OwnerType> ownerTypes = _typesProvider.apply(environment);
41+
if (ownerTypes == null || ownerTypes.isEmpty()) {
42+
return CompletableFuture.completedFuture(Collections.emptyList());
43+
}
44+
45+
// 2. Filter and group OwnerTypes by entity types
46+
final Map<String, List<OwnerType>> filteredEntitiesMap = new HashMap<>();
47+
for (final OwnerType ownerType : ownerTypes) {
48+
final LoadableType<?, ?> filteredEntity =
49+
Iterables.getOnlyElement(
50+
_loadableTypes.stream()
51+
.filter(entity -> ownerType.getClass().isAssignableFrom(entity.objectClass()))
52+
.collect(Collectors.toList()),
53+
null);
54+
if (filteredEntity == null) {
55+
continue;
56+
}
57+
filteredEntitiesMap.putIfAbsent(filteredEntity.name(), new ArrayList<>());
58+
filteredEntitiesMap.get(filteredEntity.name()).add(ownerType);
59+
}
60+
61+
// 3. Generate batch load requests for each EntityType
62+
final List<CompletableFuture<List<OwnerType>>> batchLoadFutures =
63+
filteredEntitiesMap.entrySet().stream()
64+
.map(
65+
(set) ->
66+
loadUniformTypes(
67+
environment.getDataLoaderRegistry(), set.getKey(), set.getValue()))
68+
.toList();
69+
70+
// 4. Flatmap the results of the futures
71+
return CompletableFuture.allOf(batchLoadFutures.toArray(new CompletableFuture[0]))
72+
.thenApply(
73+
v ->
74+
batchLoadFutures.stream()
75+
.map(CompletableFuture::join)
76+
.flatMap(List::stream)
77+
.collect(Collectors.toList()));
78+
}
79+
80+
private CompletableFuture<List<OwnerType>> loadUniformTypes(
81+
DataLoaderRegistry dataLoaderRegistry,
82+
String filteredEntityName,
83+
List<OwnerType> ownerTypes) {
84+
final DataLoader<Object, OwnerType> loader =
85+
dataLoaderRegistry.getDataLoader(filteredEntityName);
86+
final List<Object> keyList =
87+
ownerTypes.stream().map(ownerType -> (Object) ((Entity) ownerType).getUrn()).toList();
88+
return loader.loadMany(keyList);
89+
}
90+
}

‎datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/load/OwnerTypeBatchResolverTest.java

+167
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,167 @@
1+
package com.linkedin.datahub.graphql.resolvers.load;
2+
3+
import static org.mockito.ArgumentMatchers.any;
4+
import static org.mockito.Mockito.*;
5+
import static org.testng.Assert.*;
6+
7+
import com.google.common.collect.ImmutableList;
8+
import com.linkedin.datahub.graphql.generated.CorpGroup;
9+
import com.linkedin.datahub.graphql.generated.CorpUser;
10+
import com.linkedin.datahub.graphql.generated.OwnerType;
11+
import com.linkedin.datahub.graphql.types.corpgroup.CorpGroupType;
12+
import com.linkedin.datahub.graphql.types.corpuser.CorpUserType;
13+
import com.linkedin.entity.client.EntityClient;
14+
import graphql.schema.DataFetchingEnvironment;
15+
import java.util.List;
16+
import java.util.concurrent.CompletableFuture;
17+
import java.util.function.Function;
18+
import java.util.stream.Collectors;
19+
import org.dataloader.DataLoader;
20+
import org.dataloader.DataLoaderRegistry;
21+
import org.testng.annotations.BeforeMethod;
22+
import org.testng.annotations.Test;
23+
24+
public class OwnerTypeBatchResolverTest {
25+
private EntityClient _entityClient;
26+
private DataFetchingEnvironment _dataFetchingEnvironment;
27+
private DataLoaderRegistry _mockDataLoaderRegistry;
28+
private Function<DataFetchingEnvironment, List<OwnerType>> _ownerTypesProvider;
29+
30+
@BeforeMethod
31+
public void setupTest() {
32+
_entityClient = mock(EntityClient.class);
33+
_dataFetchingEnvironment = mock(DataFetchingEnvironment.class);
34+
_mockDataLoaderRegistry = mock(DataLoaderRegistry.class);
35+
_ownerTypesProvider = mock(Function.class);
36+
when(_dataFetchingEnvironment.getDataLoaderRegistry()).thenReturn(_mockDataLoaderRegistry);
37+
}
38+
39+
private List<OwnerType> getRequestOwnerTypes(List<String> urnList) {
40+
return urnList.stream()
41+
.map(
42+
urn -> {
43+
if (urn.startsWith("urn:li:corpuser")) {
44+
CorpUser user = new CorpUser();
45+
user.setUrn(urn);
46+
return user;
47+
} else if (urn.startsWith("urn:li:corpGroup")) {
48+
CorpGroup group = new CorpGroup();
49+
group.setUrn(urn);
50+
return group;
51+
} else {
52+
throw new RuntimeException("Can't handle urn " + urn);
53+
}
54+
})
55+
.collect(Collectors.toList());
56+
}
57+
58+
@Test
59+
public void testEmptyOwnerTypesList() throws Exception {
60+
when(_ownerTypesProvider.apply(any())).thenReturn(ImmutableList.of());
61+
62+
OwnerTypeBatchResolver resolver =
63+
new OwnerTypeBatchResolver(
64+
ImmutableList.of(
65+
new CorpUserType(_entityClient, null), new CorpGroupType(_entityClient)),
66+
_ownerTypesProvider);
67+
68+
List<OwnerType> response = resolver.get(_dataFetchingEnvironment).join();
69+
assertEquals(response.size(), 0);
70+
}
71+
72+
@Test
73+
public void testNullOwnerTypesList() throws Exception {
74+
when(_ownerTypesProvider.apply(any())).thenReturn(null);
75+
76+
OwnerTypeBatchResolver resolver =
77+
new OwnerTypeBatchResolver(
78+
ImmutableList.of(
79+
new CorpUserType(_entityClient, null), new CorpGroupType(_entityClient)),
80+
_ownerTypesProvider);
81+
82+
List<OwnerType> response = resolver.get(_dataFetchingEnvironment).join();
83+
assertEquals(response.size(), 0);
84+
}
85+
86+
@Test
87+
public void testMixedOwnerTypes() throws Exception {
88+
List<OwnerType> inputOwners =
89+
getRequestOwnerTypes(ImmutableList.of("urn:li:corpuser:1", "urn:li:corpGroup:1"));
90+
when(_ownerTypesProvider.apply(any())).thenReturn(inputOwners);
91+
92+
DataLoader<Object, Object> userLoader = mock(DataLoader.class);
93+
DataLoader<Object, Object> groupLoader = mock(DataLoader.class);
94+
95+
when(_mockDataLoaderRegistry.getDataLoader("CorpUser")).thenReturn(userLoader);
96+
when(_mockDataLoaderRegistry.getDataLoader("CorpGroup")).thenReturn(groupLoader);
97+
98+
CorpUser mockUser = new CorpUser();
99+
mockUser.setUrn("urn:li:corpuser:1");
100+
101+
CorpGroup mockGroup = new CorpGroup();
102+
mockGroup.setUrn("urn:li:corpGroup:1");
103+
104+
when(userLoader.loadMany(any()))
105+
.thenReturn(CompletableFuture.completedFuture(ImmutableList.of(mockUser)));
106+
when(groupLoader.loadMany(any()))
107+
.thenReturn(CompletableFuture.completedFuture(ImmutableList.of(mockGroup)));
108+
109+
OwnerTypeBatchResolver resolver =
110+
new OwnerTypeBatchResolver(
111+
ImmutableList.of(
112+
new CorpUserType(_entityClient, null), new CorpGroupType(_entityClient)),
113+
_ownerTypesProvider);
114+
115+
List<OwnerType> response = resolver.get(_dataFetchingEnvironment).join();
116+
assertEquals(response.size(), 2);
117+
assertTrue(response.get(0) instanceof CorpUser);
118+
assertTrue(response.get(1) instanceof CorpGroup);
119+
assertEquals(((CorpUser) response.get(0)).getUrn(), "urn:li:corpuser:1");
120+
assertEquals(((CorpGroup) response.get(1)).getUrn(), "urn:li:corpGroup:1");
121+
}
122+
123+
@Test
124+
public void testDuplicateOwnerUrns() throws Exception {
125+
List<OwnerType> inputOwners =
126+
getRequestOwnerTypes(ImmutableList.of("urn:li:corpuser:1", "urn:li:corpuser:1"));
127+
when(_ownerTypesProvider.apply(any())).thenReturn(inputOwners);
128+
129+
DataLoader<Object, Object> userLoader = mock(DataLoader.class);
130+
when(_mockDataLoaderRegistry.getDataLoader("CorpUser")).thenReturn(userLoader);
131+
132+
CorpUser mockUser = new CorpUser();
133+
mockUser.setUrn("urn:li:corpuser:1");
134+
135+
when(userLoader.loadMany(any()))
136+
.thenReturn(CompletableFuture.completedFuture(ImmutableList.of(mockUser, mockUser)));
137+
138+
OwnerTypeBatchResolver resolver =
139+
new OwnerTypeBatchResolver(
140+
ImmutableList.of(new CorpUserType(_entityClient, null)), _ownerTypesProvider);
141+
142+
List<OwnerType> response = resolver.get(_dataFetchingEnvironment).join();
143+
assertEquals(response.size(), 2);
144+
assertTrue(response.get(0) instanceof CorpUser);
145+
assertTrue(response.get(1) instanceof CorpUser);
146+
assertEquals(((CorpUser) response.get(0)).getUrn(), "urn:li:corpuser:1");
147+
assertEquals(((CorpUser) response.get(1)).getUrn(), "urn:li:corpuser:1");
148+
}
149+
150+
@Test
151+
public void testUnknownOwnerType() throws Exception {
152+
// Create a custom owner type that isn't registered with the resolver
153+
class CustomOwnerType extends CorpUser {}
154+
155+
CustomOwnerType customOwner = new CustomOwnerType();
156+
customOwner.setUrn("urn:li:custom:1");
157+
158+
when(_ownerTypesProvider.apply(any())).thenReturn(ImmutableList.of(customOwner));
159+
160+
OwnerTypeBatchResolver resolver =
161+
new OwnerTypeBatchResolver(
162+
ImmutableList.of(new CorpUserType(_entityClient, null)), _ownerTypesProvider);
163+
164+
List<OwnerType> response = resolver.get(_dataFetchingEnvironment).join();
165+
assertEquals(response.size(), 0);
166+
}
167+
}

0 commit comments

Comments
 (0)