Skip to content

Commit

Permalink
feat(ingest): datamodel to ingest organisation role metadata for a da…
Browse files Browse the repository at this point in the history
…taset
  • Loading branch information
sheeru authored and Sheerapthinath Sadagopan committed Jun 20, 2023
1 parent abf73b2 commit 07e3cc2
Show file tree
Hide file tree
Showing 21 changed files with 606 additions and 3 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -277,6 +277,7 @@
import com.linkedin.datahub.graphql.types.dataset.VersionedDatasetType;
import com.linkedin.datahub.graphql.types.dataset.mappers.DatasetProfileMapper;
import com.linkedin.datahub.graphql.types.domain.DomainType;
import com.linkedin.datahub.graphql.types.externalrolemetadata.ExternalRoleMetadataType;
import com.linkedin.datahub.graphql.types.glossary.GlossaryNodeType;
import com.linkedin.datahub.graphql.types.glossary.GlossaryTermType;
import com.linkedin.datahub.graphql.types.mlmodel.MLFeatureTableType;
Expand Down Expand Up @@ -391,6 +392,9 @@ public class GmsGraphQLEngine {
private final ViewsConfiguration viewsConfiguration;

private final DatasetType datasetType;

private final ExternalRoleMetadataType externalRoleMetadataType;

private final CorpUserType corpUserType;
private final CorpGroupType corpGroupType;
private final ChartType chartType;
Expand Down Expand Up @@ -488,6 +492,7 @@ public GmsGraphQLEngine(final GmsGraphQLEngineArgs args) {
this.featureFlags = args.featureFlags;

this.datasetType = new DatasetType(entityClient);
this.externalRoleMetadataType = new ExternalRoleMetadataType(entityClient);
this.corpUserType = new CorpUserType(entityClient, featureFlags);
this.corpGroupType = new CorpGroupType(entityClient);
this.chartType = new ChartType(entityClient);
Expand Down Expand Up @@ -523,6 +528,7 @@ public GmsGraphQLEngine(final GmsGraphQLEngineArgs args) {
// Init Lists
this.entityTypes = ImmutableList.of(
datasetType,
externalRoleMetadataType,
corpUserType,
corpGroupType,
dataPlatformType,
Expand Down Expand Up @@ -602,6 +608,7 @@ public void configureRuntimeWiring(final RuntimeWiring.Builder builder) {
configureContainerResolvers(builder);
configureDataPlatformInstanceResolvers(builder);
configureGlossaryTermResolvers(builder);
configureExtRoleResolvers(builder);
configureGlossaryNodeResolvers(builder);
configureDomainResolvers(builder);
configureAssertionResolvers(builder);
Expand All @@ -618,6 +625,23 @@ public void configureRuntimeWiring(final RuntimeWiring.Builder builder) {
configureOwnershipTypeResolver(builder);
}

private void configureExtRoleResolvers(RuntimeWiring.Builder builder) {
builder.type("ExternalRoleMetadata", typeWiring -> typeWiring
.dataFetcher("relationships", new EntityRelationshipsResultResolver(graphClient))
.dataFetcher("provisionedUsers", new LoadableTypeBatchResolver<>(corpUserType,
(env) -> ((com.linkedin.datahub.graphql.generated.ExternalRoleMetadata)
env.getSource()).getProvisionedUsers().stream()
.map(CorpUser::getUrn)
.collect(Collectors.toList())))
);
builder.type("ExternalRoleAssociation", typeWiring -> typeWiring
.dataFetcher("externalRoleMetadata",
new LoadableTypeResolver<>(externalRoleMetadataType,
(env) -> ((com.linkedin.datahub.graphql.generated.ExternalRoleAssociation)
env.getSource()).getExternalRoleMetadata().getUrn()))
);
}

public GraphQLEngine.Builder builder() {
return GraphQLEngine.builder()
.addSchema(fileBasedSchema(GMS_SCHEMA_FILE))
Expand Down Expand Up @@ -723,6 +747,7 @@ private void configureQueryResolvers(final RuntimeWiring.Builder builder) {
.dataFetcher("browse", new BrowseResolver(browsableTypes))
.dataFetcher("browsePaths", new BrowsePathsResolver(browsableTypes))
.dataFetcher("dataset", getResolver(datasetType))
.dataFetcher("externalRoleMetadata", getResolver(externalRoleMetadataType))
.dataFetcher("versionedDataset", getResolver(versionedDatasetType,
(env) -> new VersionedUrn().setUrn(UrnUtils.getUrn(env.getArgument(URN_FIELD_NAME)))
.setVersionStamp(env.getArgument(VERSION_STAMP_FIELD_NAME))))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ public class EntityTypeMapper {
static final Map<EntityType, String> ENTITY_TYPE_TO_NAME =
ImmutableMap.<EntityType, String>builder()
.put(EntityType.DATASET, "dataset")
.put(EntityType.EXTERNAL_ROLE_METADATA, "externalRoleMetadata")
.put(EntityType.CORP_USER, "corpuser")
.put(EntityType.CORP_GROUP, "corpGroup")
.put(EntityType.DATA_PLATFORM, "dataPlatform")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -70,8 +70,10 @@ private SearchUtils() {
EntityType.CORP_GROUP,
EntityType.CONTAINER,
EntityType.DOMAIN,
EntityType.NOTEBOOK,
EntityType.DATA_PRODUCT);
EntityType.DATA_PRODUCT,
EntityType.EXTERNAL_ROLE_METADATA,
EntityType.NOTEBOOK);


/**
* Entities that are part of autocomplete by default in Auto Complete Across Entities
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
import com.linkedin.datahub.graphql.generated.Domain;
import com.linkedin.datahub.graphql.generated.Entity;
import com.linkedin.datahub.graphql.generated.EntityType;
import com.linkedin.datahub.graphql.generated.ExternalRoleMetadata;
import com.linkedin.datahub.graphql.generated.GlossaryNode;
import com.linkedin.datahub.graphql.generated.GlossaryTerm;
import com.linkedin.datahub.graphql.generated.MLFeature;
Expand Down Expand Up @@ -52,6 +53,11 @@ public Entity apply(Urn input) {
((Dataset) partialEntity).setUrn(input.toString());
((Dataset) partialEntity).setType(EntityType.DATASET);
}
if (input.getEntityType().equals("externalRoleMetadata")) {
partialEntity = new ExternalRoleMetadata();
((ExternalRoleMetadata) partialEntity).setUrn(input.toString());
((ExternalRoleMetadata) partialEntity).setType(EntityType.EXTERNAL_ROLE_METADATA);
}
if (input.getEntityType().equals("glossaryTerm")) {
partialEntity = new GlossaryTerm();
((GlossaryTerm) partialEntity).setUrn(input.toString());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,8 @@ public class DatasetType implements SearchableEntityType<Dataset, String>, Brows
DATA_PLATFORM_INSTANCE_ASPECT_NAME,
SIBLINGS_ASPECT_NAME,
EMBED_ASPECT_NAME,
DATA_PRODUCTS_ASPECT_NAME
DATA_PRODUCTS_ASPECT_NAME,
EXTERNAL_ROLES_ASPECT_NAME
);

private static final Set<String> FACET_FIELDS = ImmutableSet.of("origin", "platform");
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import com.linkedin.common.DataPlatformInstance;
import com.linkedin.common.Deprecation;
import com.linkedin.common.Embed;
import com.linkedin.common.ExternalRolesMetadata;
import com.linkedin.common.GlobalTags;
import com.linkedin.common.GlossaryTerms;
import com.linkedin.common.InstitutionalMemory;
Expand Down Expand Up @@ -30,6 +31,7 @@
import com.linkedin.datahub.graphql.types.common.mappers.util.MappingHelper;
import com.linkedin.datahub.graphql.types.common.mappers.util.SystemMetadataUtils;
import com.linkedin.datahub.graphql.types.domain.DomainAssociationMapper;
import com.linkedin.datahub.graphql.types.externalrolemetadata.mappers.ExternalRolesMetadataMapper;
import com.linkedin.datahub.graphql.types.glossary.mappers.GlossaryTermsMapper;
import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
import com.linkedin.datahub.graphql.types.tag.mappers.GlobalTagsMapper;
Expand Down Expand Up @@ -106,6 +108,8 @@ public Dataset apply(@Nonnull final EntityResponse entityResponse) {
dataset.setFineGrainedLineages(UpstreamLineagesMapper.map(new UpstreamLineage(dataMap))));
mappingHelper.mapToResult(EMBED_ASPECT_NAME, (dataset, dataMap) ->
dataset.setEmbed(EmbedMapper.map(new Embed(dataMap))));
mappingHelper.mapToResult(EXTERNAL_ROLES_ASPECT_NAME, ((dataset, dataMap) ->
dataset.setExternalRolesMetadata(ExternalRolesMetadataMapper.map(new ExternalRolesMetadata(dataMap), entityUrn))));
return mappingHelper.getResult();
}

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,120 @@
package com.linkedin.datahub.graphql.types.externalrolemetadata;

import com.google.common.collect.ImmutableSet;
import com.linkedin.common.urn.Urn;
import com.linkedin.common.urn.UrnUtils;
import com.linkedin.datahub.graphql.QueryContext;
import com.linkedin.datahub.graphql.generated.AutoCompleteResults;
import com.linkedin.datahub.graphql.generated.Entity;
import com.linkedin.datahub.graphql.generated.EntityType;
import com.linkedin.datahub.graphql.generated.ExternalRoleMetadata;
import com.linkedin.datahub.graphql.generated.FacetFilterInput;
import com.linkedin.datahub.graphql.generated.SearchResults;
import com.linkedin.datahub.graphql.types.SearchableEntityType;
import com.linkedin.datahub.graphql.types.externalrolemetadata.mappers.ExternalRoleMetadataMapper;
import com.linkedin.datahub.graphql.types.mappers.AutoCompleteResultsMapper;
import com.linkedin.datahub.graphql.types.mappers.UrnSearchResultsMapper;
import com.linkedin.entity.EntityResponse;
import com.linkedin.entity.client.EntityClient;
import com.linkedin.metadata.Constants;
import com.linkedin.metadata.query.AutoCompleteResult;
import com.linkedin.metadata.query.SearchFlags;
import com.linkedin.metadata.query.filter.Filter;
import com.linkedin.metadata.search.SearchResult;
import graphql.execution.DataFetcherResult;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.Function;
import java.util.stream.Collectors;

public class ExternalRoleMetadataType implements SearchableEntityType<ExternalRoleMetadata, String>,
com.linkedin.datahub.graphql.types.EntityType<ExternalRoleMetadata, String> {

static final Set<String> ASPECTS_TO_FETCH = ImmutableSet.of(
Constants.EXTERNAL_ROLE_METADATA_KEY_ASPECT_NAME,
Constants.EXTERNAL_ROLE_PROPERTIES_ASPECT_NAME,
Constants.EXTERNAL_ROLE_PROVISIONEDUSERS_ASPECT_NAME
);

private final EntityClient _entityClient;

public ExternalRoleMetadataType(final EntityClient entityClient) {
_entityClient = entityClient;
}

@Override
public EntityType type() {
return EntityType.EXTERNAL_ROLE_METADATA;
}

@Override
public Function<Entity, String> getKeyProvider() {
return Entity::getUrn;
}

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

@Override
public List<DataFetcherResult<ExternalRoleMetadata>> batchLoad(@Nonnull List<String> urns,
@Nonnull QueryContext context) throws Exception {
final List<Urn> externalRolesUrns = urns.stream()
.map(UrnUtils::getUrn)
.collect(Collectors.toList());

try {
final Map<Urn, EntityResponse> entities = _entityClient.batchGetV2(
Constants.EXTERNAL_ROLE_METADATA_ENTITY_NAME,
new HashSet<>(externalRolesUrns),
ASPECTS_TO_FETCH,
context.getAuthentication());

final List<EntityResponse> gmsResults = new ArrayList<>();
for (Urn urn : externalRolesUrns) {
gmsResults.add(entities.getOrDefault(urn, null));
}
return gmsResults.stream()
.map(gmsResult ->
gmsResult == null ? null : DataFetcherResult.<ExternalRoleMetadata>newResult()
.data(ExternalRoleMetadataMapper.map(gmsResult))
.build()
)
.collect(Collectors.toList());
} catch (Exception e) {
throw new RuntimeException("Failed to batch load ExternalRoleMetadata", e);
}
}

@Override
public SearchResults search(@Nonnull String query,
@Nullable List<FacetFilterInput> filters,
int start,
int count,
@Nonnull final QueryContext context) throws Exception {
final SearchResult searchResult = _entityClient.search(Constants.EXTERNAL_ROLE_METADATA_ENTITY_NAME,
query, Collections.emptyMap(), start, count,
context.getAuthentication(), new SearchFlags().setFulltext(true));
return UrnSearchResultsMapper.map(searchResult);
}

@Override
public AutoCompleteResults autoComplete(@Nonnull String query,
@Nullable String field,
@Nullable Filter filters,
int limit,
@Nonnull final QueryContext context) throws Exception {
final AutoCompleteResult result = _entityClient.autoComplete(Constants.EXTERNAL_ROLE_METADATA_ENTITY_NAME,
query, filters, limit, context.getAuthentication());
return AutoCompleteResultsMapper.map(result);
}

}
Original file line number Diff line number Diff line change
@@ -0,0 +1,78 @@
package com.linkedin.datahub.graphql.types.externalrolemetadata.mappers;

import com.linkedin.common.urn.Urn;
import com.linkedin.datahub.graphql.generated.CorpUser;
import com.linkedin.datahub.graphql.generated.EntityType;
import com.linkedin.datahub.graphql.generated.ExternalRoleMetadata;
import com.linkedin.datahub.graphql.generated.ExternalRoleProperties;
import com.linkedin.datahub.graphql.generated.PrivilegeType;
import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
import com.linkedin.entity.EntityResponse;
import com.linkedin.entity.EnvelopedAspect;
import com.linkedin.entity.EnvelopedAspectMap;
import com.linkedin.metadata.Constants;
import com.linkedin.metadata.key.ExternalRoleMetadataKey;

import javax.annotation.Nonnull;
import java.util.stream.Collectors;

public class ExternalRoleMetadataMapper implements ModelMapper<EntityResponse, ExternalRoleMetadata> {

public static final ExternalRoleMetadataMapper INSTANCE = new ExternalRoleMetadataMapper();

public static ExternalRoleMetadata map(@Nonnull final EntityResponse entityResponse) {
return INSTANCE.apply(entityResponse);
}

private static ExternalRoleProperties mapExtRoleProperties(final com.linkedin.externalrolemetadata.ExternalRoleProperties e) {
final ExternalRoleProperties propertiesResult = new ExternalRoleProperties();
propertiesResult.setName(e.getName());
propertiesResult.setDescription(e.getDescription());
propertiesResult.setType(PrivilegeType.valueOf(e.getType().toString()));
propertiesResult.setRequestUrl(e.getRequesturl());

return propertiesResult;
}

private static CorpUser mapExtRoleProvisionedUsers(final com.linkedin.externalrolemetadata.ExternalRoleProvisionedUser provisionedUser) {
CorpUser result = new CorpUser();
result.setUrn(provisionedUser.getProvisionedUser().toString());
return result;
}

@Override
public ExternalRoleMetadata apply(EntityResponse input) {


final ExternalRoleMetadata result = new ExternalRoleMetadata();
final Urn entityUrn = input.getUrn();

result.setUrn(entityUrn.toString());
result.setType(EntityType.EXTERNAL_ROLE_METADATA);

final EnvelopedAspectMap aspects = input.getAspects();

final EnvelopedAspect externalRoleMetadataKeyAspect = aspects.get(Constants.EXTERNAL_ROLE_METADATA_KEY_ASPECT_NAME);
if (externalRoleMetadataKeyAspect != null) {
result.setId(new ExternalRoleMetadataKey(externalRoleMetadataKeyAspect.getValue().data()).getId());
}
final EnvelopedAspect envelopedPropertiesAspect = aspects.get(Constants.EXTERNAL_ROLE_PROPERTIES_ASPECT_NAME);
if (envelopedPropertiesAspect != null) {
result.setProperties(mapExtRoleProperties(
new com.linkedin.externalrolemetadata.ExternalRoleProperties(
envelopedPropertiesAspect.getValue().data()))
);
}

final EnvelopedAspect envelopedProvisionedUsers = aspects.get(Constants.EXTERNAL_ROLE_PROVISIONEDUSERS_ASPECT_NAME);
if (envelopedProvisionedUsers != null) {
result.setProvisionedUsers(
new com.linkedin.externalrolemetadata.ExternalRoleProvisionedUsers(envelopedProvisionedUsers.getValue().data())
.getProvisionedUsers()
.stream().map(x -> mapExtRoleProvisionedUsers(x))
.collect(Collectors.toList()));
}

return result;
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
package com.linkedin.datahub.graphql.types.externalrolemetadata.mappers;


import com.linkedin.common.urn.Urn;
import com.linkedin.datahub.graphql.generated.EntityType;
import com.linkedin.datahub.graphql.generated.ExternalRoleAssociation;
import com.linkedin.datahub.graphql.generated.ExternalRoleMetadata;
import javax.annotation.Nonnull;
import java.util.stream.Collectors;

public class ExternalRolesMetadataMapper {
public static final ExternalRolesMetadataMapper INSTANCE = new ExternalRolesMetadataMapper();

public static com.linkedin.datahub.graphql.generated.ExternalRolesMetadata map(
@Nonnull final com.linkedin.common.ExternalRolesMetadata externalRolesMetadata,
@Nonnull final Urn entityUrn) {
return INSTANCE.apply(externalRolesMetadata, entityUrn);
}

public com.linkedin.datahub.graphql.generated.ExternalRolesMetadata apply(
@Nonnull final com.linkedin.common.ExternalRolesMetadata externalRolesMetadata,
@Nonnull final Urn entityUrn) {
com.linkedin.datahub.graphql.generated.ExternalRolesMetadata result = new com.linkedin.datahub.graphql.generated.ExternalRolesMetadata();
result.setRoles(externalRolesMetadata.getRoles().stream().map(
association -> this.mapExtRoleAssociation(association, entityUrn)
).collect(Collectors.toList()));
return result;
}

private ExternalRoleAssociation mapExtRoleAssociation(com.linkedin.common.ExternalRoleAssociation association, Urn entityUrn) {
ExternalRoleAssociation extRoleAssociation = new ExternalRoleAssociation();
ExternalRoleMetadata resultExtRole = new ExternalRoleMetadata();
resultExtRole.setType(EntityType.EXTERNAL_ROLE_METADATA);
resultExtRole.setUrn(association.getUrn().toString());
extRoleAssociation.setExternalRoleMetadata(resultExtRole);
extRoleAssociation.setAssociatedUrn(entityUrn.toString());
return extRoleAssociation;
}

}
Loading

0 comments on commit 07e3cc2

Please sign in to comment.