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

feat(graphql) Add new 'entities' endpoint to batch get entities by urn (#5915)

Browse Source
This commit is contained in:
Chris Collins 2022-09-12 21:11:22 -04:00 committed by GitHub
parent b6d3fe88e5
commit e2822d4add
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 116 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/Constants.java
View File

@ -8,6 +8,7 @@ public class Constants {
private Constants() { };
public static final String URN_FIELD_NAME = "urn";
public static final String URNS_FIELD_NAME = "urns";
public static final String GMS_SCHEMA_FILE = "entity.graphql";
public static final String SEARCH_SCHEMA_FILE = "search.graphql";
public static final String APP_SCHEMA_FILE = "app.graphql";

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

@ -132,6 +132,7 @@ import com.linkedin.datahub.graphql.resolvers.ingest.source.UpsertIngestionSourc
import com.linkedin.datahub.graphql.resolvers.jobs.DataJobRunsResolver;
import com.linkedin.datahub.graphql.resolvers.jobs.EntityRunsResolver;
import com.linkedin.datahub.graphql.resolvers.load.AspectResolver;
import com.linkedin.datahub.graphql.resolvers.load.BatchGetEntitiesResolver;
import com.linkedin.datahub.graphql.resolvers.load.EntityLineageResultResolver;
import com.linkedin.datahub.graphql.resolvers.load.EntityRelationshipsResultResolver;
import com.linkedin.datahub.graphql.resolvers.load.EntityTypeBatchResolver;
@ -671,10 +672,26 @@ public class GmsGraphQLEngine {
.dataFetcher("entityExists", new EntityExistsResolver(this.entityService))
.dataFetcher("getNativeUserInviteToken", new GetNativeUserInviteTokenResolver(this.nativeUserService))
.dataFetcher("entity", getEntityResolver())
.dataFetcher("entities", getEntitiesResolver())
.dataFetcher("listRoles", new ListRolesResolver(this.entityClient))
);
}
private DataFetcher getEntitiesResolver() {
return new BatchGetEntitiesResolver(entityTypes,
(env) -> {
List<String> urns = env.getArgument(URNS_FIELD_NAME);
return urns.stream().map((urn) -> {
try {
Urn entityUrn = Urn.createFromString(urn);
return UrnToEntityMapper.map(entityUrn);
} catch (Exception e) {
throw new RuntimeException("Failed to get entity", e);
}
}).collect(Collectors.toList());
});
}
private DataFetcher getEntityResolver() {
return new EntityTypeResolver(entityTypes,
(env) -> {

38
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/BatchLoadUtils.java Normal file
View File

@ -0,0 +1,38 @@
package com.linkedin.datahub.graphql.resolvers;
import com.google.common.collect.Iterables;
import com.linkedin.datahub.graphql.generated.Entity;
import org.dataloader.DataLoader;
import org.dataloader.DataLoaderRegistry;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.stream.Collectors;
public class BatchLoadUtils {
private BatchLoadUtils() { }
public static CompletableFuture<List<Entity>> batchLoadEntitiesOfSameType(
List<Entity> entities,
List<com.linkedin.datahub.graphql.types.EntityType<?, ?>> entityTypes,
DataLoaderRegistry dataLoaderRegistry) {
if (entities.isEmpty()) {
return CompletableFuture.completedFuture(Collections.emptyList());
}
// Assume all entities are of the same type
final com.linkedin.datahub.graphql.types.EntityType filteredEntity =
Iterables.getOnlyElement(entityTypes.stream()
.filter(entity -> entities.get(0).getClass().isAssignableFrom(entity.objectClass()))
.collect(Collectors.toList()));
final DataLoader loader = dataLoaderRegistry.getDataLoader(filteredEntity.name());
List keyList = new ArrayList();
for (Entity entity : entities) {
keyList.add(filteredEntity.getKeyProvider().apply(entity));
}
return loader.loadMany(keyList);
}
}

53
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/BatchGetEntitiesResolver.java Normal file
View File

@ -0,0 +1,53 @@
package com.linkedin.datahub.graphql.resolvers.load;
import com.linkedin.datahub.graphql.generated.Entity;
import com.linkedin.datahub.graphql.generated.EntityType;
import com.linkedin.datahub.graphql.resolvers.BatchLoadUtils;
import graphql.schema.DataFetcher;
import graphql.schema.DataFetchingEnvironment;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.function.Function;
import java.util.stream.Collectors;
public class BatchGetEntitiesResolver implements DataFetcher<CompletableFuture<List<Entity>>> {
private final List<com.linkedin.datahub.graphql.types.EntityType<?, ?>> _entityTypes;
private final Function<DataFetchingEnvironment, List<Entity>> _entitiesProvider;
public BatchGetEntitiesResolver(
final List<com.linkedin.datahub.graphql.types.EntityType<?, ?>> entityTypes,
final Function<DataFetchingEnvironment, List<Entity>> entitiesProvider
) {
_entityTypes = entityTypes;
_entitiesProvider = entitiesProvider;
}
@Override
public CompletableFuture<List<Entity>> get(DataFetchingEnvironment environment) {
final List<Entity> entities = _entitiesProvider.apply(environment);
Map<EntityType, List<Entity>> entityTypeToEntities = new HashMap<>();
entities.forEach((entity) -> {
EntityType type = entity.getType();
List<Entity> entitiesList = entityTypeToEntities.getOrDefault(type, new ArrayList<>());
entitiesList.add(entity);
entityTypeToEntities.put(type, entitiesList);
});
List<CompletableFuture<List<Entity>>> entitiesFutures = new ArrayList<>();
for (Map.Entry<EntityType, List<Entity>> entry : entityTypeToEntities.entrySet()) {
CompletableFuture<List<Entity>> entitiesFuture = BatchLoadUtils
.batchLoadEntitiesOfSameType(entry.getValue(), _entityTypes, environment.getDataLoaderRegistry());
entitiesFutures.add(entitiesFuture);
}
return CompletableFuture.allOf(entitiesFutures.toArray(new CompletableFuture[0]))
.thenApply(v -> entitiesFutures.stream().flatMap(future -> future.join().stream()).collect(Collectors.toList()));
}
}

22
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/EntityTypeBatchResolver.java
View File

@ -1,16 +1,12 @@
package com.linkedin.datahub.graphql.resolvers.load;
import com.google.common.collect.Iterables;
import com.linkedin.datahub.graphql.generated.Entity;
import com.linkedin.datahub.graphql.resolvers.BatchLoadUtils;
import graphql.schema.DataFetcher;
import graphql.schema.DataFetchingEnvironment;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.function.Function;
import java.util.stream.Collectors;
import org.dataloader.DataLoader;
/**
@ -37,20 +33,6 @@ public class EntityTypeBatchResolver implements DataFetcher<CompletableFuture<Li
@Override
public CompletableFuture<List<Entity>> get(DataFetchingEnvironment environment) {
final List<Entity> entities = _entitiesProvider.apply(environment);
if (entities.isEmpty()) {
return CompletableFuture.completedFuture(Collections.emptyList());
}
// Assume all entities are of the same type
final com.linkedin.datahub.graphql.types.EntityType filteredEntity =
Iterables.getOnlyElement(_entityTypes.stream()
.filter(entity -> entities.get(0).getClass().isAssignableFrom(entity.objectClass()))
.collect(Collectors.toList()));
final DataLoader loader = environment.getDataLoaderRegistry().getDataLoader(filteredEntity.name());
List keyList = new ArrayList();
for (Entity entity : entities) {
keyList.add(filteredEntity.getKeyProvider().apply(entity));
}
return loader.loadMany(keyList);
return BatchLoadUtils.batchLoadEntitiesOfSameType(entities, _entityTypes, environment.getDataLoaderRegistry());
}
}

5
datahub-graphql-core/src/main/resources/entity.graphql
View File

@ -174,6 +174,11 @@ type Query {
"""
entity(urn: String!): Entity
"""
Gets entities based on their urns
"""
entities(urns: [String!]!): [Entity]
"""
List all DataHub Roles
"""