diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java index 49c28e3335e7d1..03bbc8f6d6f015 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java @@ -28,6 +28,7 @@ import com.linkedin.datahub.graphql.resolvers.MeResolver; import com.linkedin.datahub.graphql.resolvers.ResolverUtils; import com.linkedin.datahub.graphql.resolvers.application.BatchSetApplicationResolver; +import com.linkedin.datahub.graphql.resolvers.application.BatchUnsetApplicationResolver; import com.linkedin.datahub.graphql.resolvers.application.CreateApplicationResolver; import com.linkedin.datahub.graphql.resolvers.application.DeleteApplicationResolver; import com.linkedin.datahub.graphql.resolvers.assertion.AssertionRunEventResolver; @@ -1394,6 +1395,9 @@ private void configureMutationResolvers(final RuntimeWiring.Builder builder) { new DeleteApplicationResolver(this.entityClient, this.applicationService)) .dataFetcher( "batchSetApplication", new BatchSetApplicationResolver(this.applicationService)) + .dataFetcher( + "batchUnsetApplication", + new BatchUnsetApplicationResolver(this.applicationService)) .dataFetcher( "createOwnershipType", new CreateOwnershipTypeResolver(this.ownershipTypeService)) .dataFetcher( diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/application/ApplicationAuthorizationUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/application/ApplicationAuthorizationUtils.java index 45a1b6dba36b1d..cc70b6cd83f798 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/application/ApplicationAuthorizationUtils.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/application/ApplicationAuthorizationUtils.java @@ -4,7 +4,16 @@ import static com.linkedin.metadata.authorization.ApiOperation.MANAGE; import com.datahub.authorization.AuthUtil; +import com.datahub.authorization.ConjunctivePrivilegeGroup; +import com.datahub.authorization.DisjunctivePrivilegeGroup; +import com.google.common.collect.ImmutableList; +import com.linkedin.common.urn.Urn; +import com.linkedin.common.urn.UrnUtils; import com.linkedin.datahub.graphql.QueryContext; +import com.linkedin.datahub.graphql.authorization.AuthorizationUtils; +import com.linkedin.datahub.graphql.exception.AuthorizationException; +import com.linkedin.metadata.authorization.PoliciesConfig; +import com.linkedin.metadata.service.ApplicationService; import java.util.List; import javax.annotation.Nonnull; import lombok.extern.slf4j.Slf4j; @@ -22,4 +31,53 @@ public static boolean canManageApplications(@Nonnull QueryContext context) { return AuthUtil.isAuthorizedEntityType( context.getOperationContext(), MANAGE, List.of(APPLICATION_ENTITY_NAME)); } + + /** + * Verifies that the current user is authorized to edit applications on a specific resource + * entity. + * + * @throws AuthorizationException if the user is not authorized + */ + public static void verifyEditApplicationsAuthorization( + @Nonnull Urn resourceUrn, @Nonnull QueryContext context) { + if (!AuthorizationUtils.isAuthorized( + context, + resourceUrn.getEntityType(), + resourceUrn.toString(), + new DisjunctivePrivilegeGroup( + ImmutableList.of( + new ConjunctivePrivilegeGroup( + ImmutableList.of(PoliciesConfig.EDIT_ENTITY_APPLICATIONS_PRIVILEGE.getType())), + new ConjunctivePrivilegeGroup( + ImmutableList.of(PoliciesConfig.EDIT_ENTITY_PRIVILEGE.getType())))))) { + throw new AuthorizationException( + "Unauthorized to perform this action. Please contact your DataHub administrator."); + } + } + + /** + * Verifies that all resources exist and that the current user is authorized to edit applications + * on them. + * + * @param resources List of resource URN strings to verify + * @param applicationService Service to verify entity existence + * @param context Query context with operation context and authorization info + * @param operationName Name of the operation being performed (for error messages) + * @throws RuntimeException if any resource does not exist + * @throws AuthorizationException if the user is not authorized for any resource + */ + public static void verifyResourcesExistAndAuthorized( + @Nonnull List resources, + @Nonnull ApplicationService applicationService, + @Nonnull QueryContext context, + @Nonnull String operationName) { + for (String resource : resources) { + Urn resourceUrn = UrnUtils.getUrn(resource); + if (!applicationService.verifyEntityExists(context.getOperationContext(), resourceUrn)) { + throw new RuntimeException( + String.format("Failed to %s, %s in resources does not exist", operationName, resource)); + } + verifyEditApplicationsAuthorization(resourceUrn, context); + } + } } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/application/BatchSetApplicationResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/application/BatchSetApplicationResolver.java index e05272c0a7aae1..037267da672fc6 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/application/BatchSetApplicationResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/application/BatchSetApplicationResolver.java @@ -1,18 +1,13 @@ package com.linkedin.datahub.graphql.resolvers.application; import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.bindArgument; +import static com.linkedin.datahub.graphql.resolvers.application.ApplicationAuthorizationUtils.verifyResourcesExistAndAuthorized; -import com.datahub.authorization.ConjunctivePrivilegeGroup; -import com.datahub.authorization.DisjunctivePrivilegeGroup; -import com.google.common.collect.ImmutableList; import com.linkedin.common.urn.Urn; import com.linkedin.common.urn.UrnUtils; import com.linkedin.datahub.graphql.QueryContext; -import com.linkedin.datahub.graphql.authorization.AuthorizationUtils; import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils; -import com.linkedin.datahub.graphql.exception.AuthorizationException; import com.linkedin.datahub.graphql.generated.BatchSetApplicationInput; -import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.metadata.service.ApplicationService; import graphql.schema.DataFetcher; import graphql.schema.DataFetchingEnvironment; @@ -63,29 +58,7 @@ public CompletableFuture get(DataFetchingEnvironment environment) throw } private void verifyResources(List resources, QueryContext context) { - for (String resource : resources) { - if (!applicationService.verifyEntityExists( - context.getOperationContext(), UrnUtils.getUrn(resource))) { - throw new RuntimeException( - String.format( - "Failed to batch set Application, %s in resources does not exist", resource)); - } - Urn resourceUrn = UrnUtils.getUrn(resource); - if (!AuthorizationUtils.isAuthorized( - context, - resourceUrn.getEntityType(), - resourceUrn.toString(), - new DisjunctivePrivilegeGroup( - ImmutableList.of( - new ConjunctivePrivilegeGroup( - ImmutableList.of( - PoliciesConfig.EDIT_ENTITY_APPLICATIONS_PRIVILEGE.getType())), - new ConjunctivePrivilegeGroup( - ImmutableList.of(PoliciesConfig.EDIT_ENTITY_PRIVILEGE.getType())))))) { - throw new AuthorizationException( - "Unauthorized to perform this action. Please contact your DataHub administrator."); - } - } + verifyResourcesExistAndAuthorized(resources, applicationService, context, "set_application"); } private void verifyApplication(String maybeApplicationUrn, QueryContext context) { diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/application/BatchUnsetApplicationResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/application/BatchUnsetApplicationResolver.java new file mode 100644 index 00000000000000..d147e12e7e435a --- /dev/null +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/application/BatchUnsetApplicationResolver.java @@ -0,0 +1,85 @@ +package com.linkedin.datahub.graphql.resolvers.application; + +import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.bindArgument; +import static com.linkedin.datahub.graphql.resolvers.application.ApplicationAuthorizationUtils.verifyResourcesExistAndAuthorized; + +import com.linkedin.common.urn.Urn; +import com.linkedin.common.urn.UrnUtils; +import com.linkedin.datahub.graphql.QueryContext; +import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils; +import com.linkedin.datahub.graphql.generated.BatchUnsetApplicationInput; +import com.linkedin.metadata.service.ApplicationService; +import graphql.schema.DataFetcher; +import graphql.schema.DataFetchingEnvironment; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.stream.Collectors; +import javax.annotation.Nonnull; +import lombok.RequiredArgsConstructor; +import lombok.extern.slf4j.Slf4j; + +@Slf4j +@RequiredArgsConstructor +public class BatchUnsetApplicationResolver implements DataFetcher> { + + private final ApplicationService applicationService; + + @Override + public CompletableFuture get(DataFetchingEnvironment environment) throws Exception { + final QueryContext context = environment.getContext(); + final BatchUnsetApplicationInput input = + bindArgument(environment.getArgument("input"), BatchUnsetApplicationInput.class); + final String applicationUrn = input.getApplicationUrn(); + final List resources = input.getResourceUrns(); + + return GraphQLConcurrencyUtils.supplyAsync( + () -> { + verifyResources(resources, context); + verifyApplication(applicationUrn, context); + + try { + List resourceUrns = + resources.stream().map(UrnUtils::getUrn).collect(Collectors.toList()); + batchUnsetApplication(applicationUrn, resourceUrns, context); + return true; + } catch (Exception e) { + log.error("Failed to perform update against input {}, {}", input, e.getMessage()); + throw new RuntimeException( + String.format("Failed to perform update against input %s", input), e); + } + }, + this.getClass().getSimpleName(), + "get"); + } + + private void verifyResources(List resources, QueryContext context) { + verifyResourcesExistAndAuthorized(resources, applicationService, context, "unset_application"); + } + + private void verifyApplication(String applicationUrn, QueryContext context) { + if (!applicationService.verifyEntityExists( + context.getOperationContext(), UrnUtils.getUrn(applicationUrn))) { + throw new RuntimeException( + String.format( + "Failed to batch unset Application, Application urn %s does not exist", + applicationUrn)); + } + } + + private void batchUnsetApplication( + @Nonnull String applicationUrn, List resources, QueryContext context) { + try { + applicationService.batchUnsetApplication( + context.getOperationContext(), + UrnUtils.getUrn(applicationUrn), + resources, + UrnUtils.getUrn(context.getActorUrn())); + } catch (Exception e) { + throw new RuntimeException( + String.format( + "Failed to batch unset Application %s from resources with urns %s!", + applicationUrn, resources), + e); + } + } +} diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/application/ApplicationAssociationMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/application/ApplicationAssociationMapper.java index f7f9c8f9e7d601..ad5c3a223c189f 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/application/ApplicationAssociationMapper.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/application/ApplicationAssociationMapper.java @@ -6,6 +6,8 @@ import com.linkedin.datahub.graphql.generated.Application; import com.linkedin.datahub.graphql.generated.ApplicationAssociation; import com.linkedin.datahub.graphql.generated.EntityType; +import java.util.List; +import java.util.stream.Collectors; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -25,6 +27,13 @@ public static ApplicationAssociation map( return INSTANCE.apply(context, applications, entityUrn); } + public static List mapList( + @Nullable final QueryContext context, + @Nonnull final com.linkedin.application.Applications applications, + @Nonnull final String entityUrn) { + return INSTANCE.applyList(context, applications, entityUrn); + } + public ApplicationAssociation apply( @Nullable final QueryContext context, @Nonnull final com.linkedin.application.Applications applications, @@ -43,4 +52,26 @@ public ApplicationAssociation apply( } return null; } + + public List applyList( + @Nullable final QueryContext context, + @Nonnull final com.linkedin.application.Applications applications, + @Nonnull final String entityUrn) { + return applications.getApplications().stream() + .filter( + applicationUrn -> + context == null || canView(context.getOperationContext(), applicationUrn)) + .map( + applicationUrn -> { + ApplicationAssociation association = new ApplicationAssociation(); + association.setApplication( + Application.builder() + .setType(EntityType.APPLICATION) + .setUrn(applicationUrn.toString()) + .build()); + association.setAssociatedUrn(entityUrn); + return association; + }) + .collect(Collectors.toList()); + } } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/chart/mappers/ChartMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/chart/mappers/ChartMapper.java index 968db3bb547704..d7f3c1a0a53798 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/chart/mappers/ChartMapper.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/chart/mappers/ChartMapper.java @@ -314,6 +314,12 @@ private static void mapDomains( private static void mapApplicationAssociation( @Nullable final QueryContext context, @Nonnull Chart chart, @Nonnull DataMap dataMap) { final Applications applications = new Applications(dataMap); - chart.setApplication(ApplicationAssociationMapper.map(context, applications, chart.getUrn())); + final java.util.List + applicationAssociations = + ApplicationAssociationMapper.mapList(context, applications, chart.getUrn()); + chart.setApplications(applicationAssociations); + if (applicationAssociations != null && !applicationAssociations.isEmpty()) { + chart.setApplication(applicationAssociations.get(0)); + } } } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dashboard/mappers/DashboardMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dashboard/mappers/DashboardMapper.java index b7afee72968452..3418f7b902dcd6 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dashboard/mappers/DashboardMapper.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dashboard/mappers/DashboardMapper.java @@ -308,7 +308,12 @@ private static void mapApplicationAssociation( @Nonnull Dashboard dashboard, @Nonnull DataMap dataMap) { final Applications applications = new Applications(dataMap); - dashboard.setApplication( - ApplicationAssociationMapper.map(context, applications, dashboard.getUrn())); + final java.util.List + applicationAssociations = + ApplicationAssociationMapper.mapList(context, applications, dashboard.getUrn()); + dashboard.setApplications(applicationAssociations); + if (applicationAssociations != null && !applicationAssociations.isEmpty()) { + dashboard.setApplication(applicationAssociations.get(0)); + } } } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataflow/mappers/DataFlowMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataflow/mappers/DataFlowMapper.java index d7a22670c5cd51..610ebdc40d4f6f 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataflow/mappers/DataFlowMapper.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataflow/mappers/DataFlowMapper.java @@ -240,7 +240,12 @@ private static void mapDomains( private static void mapApplicationAssociation( @Nullable final QueryContext context, @Nonnull DataFlow dataFlow, @Nonnull DataMap dataMap) { final Applications applications = new Applications(dataMap); - dataFlow.setApplication( - ApplicationAssociationMapper.map(context, applications, dataFlow.getUrn())); + final java.util.List + applicationAssociations = + ApplicationAssociationMapper.mapList(context, applications, dataFlow.getUrn()); + dataFlow.setApplications(applicationAssociations); + if (applicationAssociations != null && !applicationAssociations.isEmpty()) { + dataFlow.setApplication(applicationAssociations.get(0)); + } } } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/datajob/mappers/DataJobMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/datajob/mappers/DataJobMapper.java index 2e440dd3040b7f..714f902e251c89 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/datajob/mappers/DataJobMapper.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/datajob/mappers/DataJobMapper.java @@ -236,8 +236,13 @@ private void mapApplicationAssociation( if (aspectMap.containsKey(APPLICATION_MEMBERSHIP_ASPECT_NAME)) { final Applications applications = new Applications(aspectMap.get(APPLICATION_MEMBERSHIP_ASPECT_NAME).getValue().data()); - dataJob.setApplication( - ApplicationAssociationMapper.map(context, applications, dataJob.getUrn())); + final java.util.List + applicationAssociations = + ApplicationAssociationMapper.mapList(context, applications, dataJob.getUrn()); + dataJob.setApplications(applicationAssociations); + if (applicationAssociations != null && !applicationAssociations.isEmpty()) { + dataJob.setApplication(applicationAssociations.get(0)); + } } } } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataproduct/mappers/DataProductMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataproduct/mappers/DataProductMapper.java index 5d40d40db2a212..6541afc540bf0e 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataproduct/mappers/DataProductMapper.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataproduct/mappers/DataProductMapper.java @@ -164,7 +164,12 @@ private static void mapApplicationAssociation( @Nonnull DataProduct dataProduct, @Nonnull DataMap dataMap) { final Applications applications = new Applications(dataMap); - dataProduct.setApplication( - ApplicationAssociationMapper.map(context, applications, dataProduct.getUrn())); + final java.util.List + applicationAssociations = + ApplicationAssociationMapper.mapList(context, applications, dataProduct.getUrn()); + dataProduct.setApplications(applicationAssociations); + if (applicationAssociations != null && !applicationAssociations.isEmpty()) { + dataProduct.setApplication(applicationAssociations.get(0)); + } } } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataset/mappers/DatasetMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataset/mappers/DatasetMapper.java index 44fb645a961cff..b2240202168fc5 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataset/mappers/DatasetMapper.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataset/mappers/DatasetMapper.java @@ -159,8 +159,6 @@ public Dataset apply( (dataset, dataMap) -> dataset.setDataPlatformInstance( DataPlatformInstanceAspectMapper.map(context, new DataPlatformInstance(dataMap)))); - mappingHelper.mapToResult( - "applications", (dataset, dataMap) -> mapApplicationAssociation(context, dataset, dataMap)); mappingHelper.mapToResult( SIBLINGS_ASPECT_NAME, (dataset, dataMap) -> @@ -324,7 +322,12 @@ private static void mapDomains( private static void mapApplicationAssociation( @Nullable final QueryContext context, @Nonnull Dataset dataset, @Nonnull DataMap dataMap) { final Applications applications = new Applications(dataMap); - dataset.setApplication( - ApplicationAssociationMapper.map(context, applications, dataset.getUrn())); + final java.util.List + applicationAssociations = + ApplicationAssociationMapper.mapList(context, applications, dataset.getUrn()); + dataset.setApplications(applicationAssociations); + if (applicationAssociations != null && !applicationAssociations.isEmpty()) { + dataset.setApplication(applicationAssociations.get(0)); + } } } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/glossary/mappers/GlossaryTermMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/glossary/mappers/GlossaryTermMapper.java index e1be7647fde57e..59710595a8ef67 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/glossary/mappers/GlossaryTermMapper.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/glossary/mappers/GlossaryTermMapper.java @@ -149,7 +149,12 @@ private static void mapApplicationAssociation( @Nonnull GlossaryTerm glossaryTerm, @Nonnull DataMap dataMap) { final Applications applications = new Applications(dataMap); - glossaryTerm.setApplication( - ApplicationAssociationMapper.map(context, applications, glossaryTerm.getUrn())); + final java.util.List + applicationAssociations = + ApplicationAssociationMapper.mapList(context, applications, glossaryTerm.getUrn()); + glossaryTerm.setApplications(applicationAssociations); + if (applicationAssociations != null && !applicationAssociations.isEmpty()) { + glossaryTerm.setApplication(applicationAssociations.get(0)); + } } } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLFeatureMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLFeatureMapper.java index c8dbe629840cad..49b71311bea6fb 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLFeatureMapper.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLFeatureMapper.java @@ -186,7 +186,12 @@ private static void mapApplicationAssociation( @Nonnull MLFeature mlFeature, @Nonnull DataMap dataMap) { final Applications applications = new Applications(dataMap); - mlFeature.setApplication( - ApplicationAssociationMapper.map(context, applications, mlFeature.getUrn())); + final java.util.List + applicationAssociations = + ApplicationAssociationMapper.mapList(context, applications, mlFeature.getUrn()); + mlFeature.setApplications(applicationAssociations); + if (applicationAssociations != null && !applicationAssociations.isEmpty()) { + mlFeature.setApplication(applicationAssociations.get(0)); + } } } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLFeatureTableMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLFeatureTableMapper.java index efcb61f4d112ed..f630a62e63d2d5 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLFeatureTableMapper.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLFeatureTableMapper.java @@ -189,7 +189,12 @@ private static void mapApplicationAssociation( @Nonnull MLFeatureTable mlFeatureTable, @Nonnull DataMap dataMap) { final Applications applications = new Applications(dataMap); - mlFeatureTable.setApplication( - ApplicationAssociationMapper.map(context, applications, mlFeatureTable.getUrn())); + final java.util.List + applicationAssociations = + ApplicationAssociationMapper.mapList(context, applications, mlFeatureTable.getUrn()); + mlFeatureTable.setApplications(applicationAssociations); + if (applicationAssociations != null && !applicationAssociations.isEmpty()) { + mlFeatureTable.setApplication(applicationAssociations.get(0)); + } } } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLModelGroupMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLModelGroupMapper.java index 09260c43ed1bd5..05b9ae29a635fb 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLModelGroupMapper.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLModelGroupMapper.java @@ -186,7 +186,12 @@ private static void mapApplicationAssociation( @Nonnull MLModelGroup mlModelGroup, @Nonnull DataMap dataMap) { final Applications applications = new Applications(dataMap); - mlModelGroup.setApplication( - ApplicationAssociationMapper.map(context, applications, mlModelGroup.getUrn())); + final java.util.List + applicationAssociations = + ApplicationAssociationMapper.mapList(context, applications, mlModelGroup.getUrn()); + mlModelGroup.setApplications(applicationAssociations); + if (applicationAssociations != null && !applicationAssociations.isEmpty()) { + mlModelGroup.setApplication(applicationAssociations.get(0)); + } } } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLModelMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLModelMapper.java index 1bbe9998aaf226..38182fc850ac07 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLModelMapper.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLModelMapper.java @@ -258,7 +258,12 @@ private static void mapEditableProperties(MLModel entity, DataMap dataMap) { private static void mapApplicationAssociation( @Nullable final QueryContext context, @Nonnull MLModel mlModel, @Nonnull DataMap dataMap) { final Applications applications = new Applications(dataMap); - mlModel.setApplication( - ApplicationAssociationMapper.map(context, applications, mlModel.getUrn())); + final java.util.List + applicationAssociations = + ApplicationAssociationMapper.mapList(context, applications, mlModel.getUrn()); + mlModel.setApplications(applicationAssociations); + if (applicationAssociations != null && !applicationAssociations.isEmpty()) { + mlModel.setApplication(applicationAssociations.get(0)); + } } } diff --git a/datahub-graphql-core/src/main/resources/entity.graphql b/datahub-graphql-core/src/main/resources/entity.graphql index 6dc7a279701bec..bdf7b2e2898f30 100644 --- a/datahub-graphql-core/src/main/resources/entity.graphql +++ b/datahub-graphql-core/src/main/resources/entity.graphql @@ -970,6 +970,14 @@ type Mutation { input: BatchSetApplicationInput! ): Boolean + """ + Batch unset a specific Application from a list of entities + """ + batchUnsetApplication( + "Input for batch unsetting application" + input: BatchUnsetApplicationInput! + ): Boolean + """ Create a Custom Ownership Type. This requires the 'Manage Ownership Types' Metadata Privilege. """ @@ -1754,9 +1762,16 @@ type Dataset implements EntityWithRelationships & Entity & BrowsableEntity { domain: DomainAssociation """ + The applications associated with the dataset + """ + applications: [ApplicationAssociation!] + + """ + Deprecated, use applications instead The application associated with the dataset """ application: ApplicationAssociation + @deprecated(reason: "Use applications instead") """ The forms associated with the Dataset @@ -2229,9 +2244,16 @@ type VersionedDataset implements Entity { domain: DomainAssociation """ + The applications associated with the entity + """ + applications: [ApplicationAssociation!] + + """ + Deprecated, use applications instead The application associated with the entity """ application: ApplicationAssociation + @deprecated(reason: "Use applications instead") """ Experimental! The resolved health status of the asset @@ -2451,9 +2473,16 @@ type GlossaryTerm implements Entity { domain: DomainAssociation """ + The applications associated with the glossary term + """ + applications: [ApplicationAssociation!] + + """ + Deprecated, use applications instead The application associated with the glossary term """ application: ApplicationAssociation + @deprecated(reason: "Use applications instead") """ References to internal resources related to the Glossary Term @@ -3115,9 +3144,16 @@ type Container implements Entity { domain: DomainAssociation """ + The applications associated with the entity + """ + applications: [ApplicationAssociation!] + + """ + Deprecated, use applications instead The application associated with the entity """ application: ApplicationAssociation + @deprecated(reason: "Use applications instead") """ The deprecation status of the container @@ -5680,9 +5716,16 @@ type Notebook implements Entity & BrowsableEntity { domain: DomainAssociation """ + The applications associated with the entity + """ + applications: [ApplicationAssociation!] + + """ + Deprecated, use applications instead The application associated with the entity """ application: ApplicationAssociation + @deprecated(reason: "Use applications instead") """ The specific instance of the data platform that this entity belongs to @@ -5972,9 +6015,16 @@ type Dashboard implements EntityWithRelationships & Entity & BrowsableEntity { domain: DomainAssociation """ + The applications associated with the entity + """ + applications: [ApplicationAssociation!] + + """ + Deprecated, use applications instead The application associated with the entity """ application: ApplicationAssociation + @deprecated(reason: "Use applications instead") """ The specific instance of the data platform that this entity belongs to @@ -6315,9 +6365,16 @@ type Chart implements EntityWithRelationships & Entity & BrowsableEntity { domain: DomainAssociation """ + The applications associated with the entity + """ + applications: [ApplicationAssociation!] + + """ + Deprecated, use applications instead The application associated with the entity """ application: ApplicationAssociation + @deprecated(reason: "Use applications instead") """ The specific instance of the data platform that this entity belongs to @@ -6719,9 +6776,16 @@ type DataFlow implements EntityWithRelationships & Entity & BrowsableEntity { domain: DomainAssociation """ + The applications associated with the entity + """ + applications: [ApplicationAssociation!] + + """ + Deprecated, use applications instead The application associated with the entity """ application: ApplicationAssociation + @deprecated(reason: "Use applications instead") """ The specific instance of the data platform that this entity belongs to @@ -6981,9 +7045,16 @@ type DataJob implements EntityWithRelationships & Entity & BrowsableEntity { domain: DomainAssociation """ + The applications associated with the entity + """ + applications: [ApplicationAssociation!] + + """ + Deprecated, use applications instead The application associated with the entity """ application: ApplicationAssociation + @deprecated(reason: "Use applications instead") """ Granular API for querying edges extending from this entity @@ -10371,9 +10442,16 @@ type MLModel implements EntityWithRelationships & Entity & BrowsableEntity { domain: DomainAssociation """ + The applications associated with the entity + """ + applications: [ApplicationAssociation!] + + """ + Deprecated, use applications instead The application associated with the entity """ application: ApplicationAssociation + @deprecated(reason: "Use applications instead") """ An additional set of of read write properties @@ -10513,9 +10591,16 @@ type MLModelGroup implements EntityWithRelationships & Entity & BrowsableEntity domain: DomainAssociation """ + The applications associated with the entity + """ + applications: [ApplicationAssociation!] + + """ + Deprecated, use applications instead The application associated with the entity """ application: ApplicationAssociation + @deprecated(reason: "Use applications instead") """ An additional set of of read write properties @@ -10700,9 +10785,16 @@ type MLFeature implements EntityWithRelationships & Entity { domain: DomainAssociation """ + The applications associated with the entity + """ + applications: [ApplicationAssociation!] + + """ + Deprecated, use applications instead The application associated with the entity """ application: ApplicationAssociation + @deprecated(reason: "Use applications instead") """ An additional set of of read write properties @@ -10959,9 +11051,16 @@ type MLPrimaryKey implements EntityWithRelationships & Entity { domain: DomainAssociation """ + The applications associated with the entity + """ + applications: [ApplicationAssociation!] + + """ + Deprecated, use applications instead The application associated with the entity """ application: ApplicationAssociation + @deprecated(reason: "Use applications instead") """ An additional set of of read write properties @@ -11113,9 +11212,16 @@ type MLFeatureTable implements EntityWithRelationships & Entity & BrowsableEntit domain: DomainAssociation """ + The applications associated with the entity + """ + applications: [ApplicationAssociation!] + + """ + Deprecated, use applications instead The application associated with the entity """ application: ApplicationAssociation + @deprecated(reason: "Use applications instead") """ An additional set of of read write properties @@ -13120,9 +13226,16 @@ type DataProduct implements Entity { domain: DomainAssociation """ + The applications associated with the data product + """ + applications: [ApplicationAssociation!] + + """ + Deprecated, use applications instead The application associated with the data product """ application: ApplicationAssociation + @deprecated(reason: "Use applications instead") """ Tags used for searching Data Product @@ -13470,6 +13583,21 @@ input BatchSetApplicationInput { resourceUrns: [String!]! } +""" +Input properties required for batch unsetting a specific Application from entities +""" +input BatchUnsetApplicationInput { + """ + The urn of the application to be removed from the resources + """ + applicationUrn: String! + + """ + The urns of the entities the given application should be removed from + """ + resourceUrns: [String!]! +} + """ Properties about an individual Custom Ownership Type. """ diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/application/ApplicationAuthorizationUtilsTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/application/ApplicationAuthorizationUtilsTest.java new file mode 100644 index 00000000000000..621b17419123c8 --- /dev/null +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/application/ApplicationAuthorizationUtilsTest.java @@ -0,0 +1,110 @@ +package com.linkedin.datahub.graphql.resolvers.application; + +import static com.linkedin.datahub.graphql.TestUtils.*; +import static org.testng.Assert.*; + +import com.google.common.collect.ImmutableList; +import com.linkedin.common.urn.Urn; +import com.linkedin.common.urn.UrnUtils; +import com.linkedin.datahub.graphql.QueryContext; +import com.linkedin.datahub.graphql.exception.AuthorizationException; +import com.linkedin.metadata.service.ApplicationService; +import org.mockito.Mockito; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +public class ApplicationAuthorizationUtilsTest { + + private static final String TEST_ENTITY_URN_1 = + "urn:li:dataset:(urn:li:dataPlatform:mysql,my-test,PROD)"; + private static final String TEST_ENTITY_URN_2 = + "urn:li:dataset:(urn:li:dataPlatform:mysql,my-test-2,PROD)"; + private static final String TEST_ACTOR_URN = "urn:li:corpuser:test"; + + private ApplicationService mockApplicationService; + private QueryContext mockAllowContext; + private QueryContext mockDenyContext; + + @BeforeMethod + public void setupTest() { + mockApplicationService = Mockito.mock(ApplicationService.class); + mockAllowContext = getMockAllowContext(TEST_ACTOR_URN); + mockDenyContext = getMockDenyContext(TEST_ACTOR_URN); + } + + private void mockExists(Urn urn, boolean exists) { + Mockito.when(mockApplicationService.verifyEntityExists(Mockito.any(), Mockito.eq(urn))) + .thenReturn(exists); + } + + @Test + public void testVerifyEditApplicationsAuthorizationFailure() { + Urn resourceUrn = UrnUtils.getUrn(TEST_ENTITY_URN_1); + + assertThrows( + AuthorizationException.class, + () -> + ApplicationAuthorizationUtils.verifyEditApplicationsAuthorization( + resourceUrn, mockDenyContext)); + } + + @Test + public void testVerifyResourcesExistAndAuthorizedSuccess() { + mockExists(UrnUtils.getUrn(TEST_ENTITY_URN_1), true); + mockExists(UrnUtils.getUrn(TEST_ENTITY_URN_2), true); + + ApplicationAuthorizationUtils.verifyResourcesExistAndAuthorized( + ImmutableList.of(TEST_ENTITY_URN_1, TEST_ENTITY_URN_2), + mockApplicationService, + mockAllowContext, + "test operation"); + + Mockito.verify(mockApplicationService, Mockito.times(1)) + .verifyEntityExists(Mockito.any(), Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1))); + Mockito.verify(mockApplicationService, Mockito.times(1)) + .verifyEntityExists(Mockito.any(), Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_2))); + } + + @Test + public void testVerifyResourcesExistAndAuthorizedResourceDoesNotExist() { + mockExists(UrnUtils.getUrn(TEST_ENTITY_URN_1), false); + + assertThrows( + RuntimeException.class, + () -> + ApplicationAuthorizationUtils.verifyResourcesExistAndAuthorized( + ImmutableList.of(TEST_ENTITY_URN_1), + mockApplicationService, + mockAllowContext, + "test operation")); + } + + @Test + public void testVerifyResourcesExistAndAuthorizedMultipleResourcesOneDoesNotExist() { + mockExists(UrnUtils.getUrn(TEST_ENTITY_URN_1), true); + mockExists(UrnUtils.getUrn(TEST_ENTITY_URN_2), false); + + assertThrows( + RuntimeException.class, + () -> + ApplicationAuthorizationUtils.verifyResourcesExistAndAuthorized( + ImmutableList.of(TEST_ENTITY_URN_1, TEST_ENTITY_URN_2), + mockApplicationService, + mockAllowContext, + "test operation")); + } + + @Test + public void testVerifyResourcesExistAndAuthorizedMultipleResourcesUnauthorized() { + mockExists(UrnUtils.getUrn(TEST_ENTITY_URN_1), true); + + assertThrows( + AuthorizationException.class, + () -> + ApplicationAuthorizationUtils.verifyResourcesExistAndAuthorized( + ImmutableList.of(TEST_ENTITY_URN_1, TEST_ENTITY_URN_2), + mockApplicationService, + mockDenyContext, + "test operation")); + } +} diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/application/BatchUnsetApplicationResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/application/BatchUnsetApplicationResolverTest.java new file mode 100644 index 00000000000000..95972aa390b63b --- /dev/null +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/application/BatchUnsetApplicationResolverTest.java @@ -0,0 +1,180 @@ +package com.linkedin.datahub.graphql.resolvers.application; + +import static com.linkedin.datahub.graphql.TestUtils.*; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.testng.Assert.*; + +import com.google.common.collect.ImmutableList; +import com.linkedin.common.urn.Urn; +import com.linkedin.common.urn.UrnUtils; +import com.linkedin.datahub.graphql.QueryContext; +import com.linkedin.datahub.graphql.generated.BatchUnsetApplicationInput; +import com.linkedin.metadata.service.ApplicationService; +import graphql.schema.DataFetchingEnvironment; +import java.util.concurrent.CompletionException; +import org.mockito.Mockito; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +public class BatchUnsetApplicationResolverTest { + + private static final String TEST_ENTITY_URN_1 = + "urn:li:dataset:(urn:li:dataPlatform:mysql,my-test,PROD)"; + private static final String TEST_ENTITY_URN_2 = + "urn:li:dataset:(urn:li:dataPlatform:mysql,my-test-2,PROD)"; + private static final String TEST_APPLICATION_URN = "urn:li:application:test-app-id"; + private static final String TEST_ACTOR_URN = "urn:li:corpuser:test"; + + private ApplicationService mockApplicationService; + private BatchUnsetApplicationResolver resolver; + private QueryContext mockContext; + private DataFetchingEnvironment mockEnv; + + @BeforeMethod + public void setupTest() { + mockApplicationService = Mockito.mock(ApplicationService.class); + resolver = new BatchUnsetApplicationResolver(mockApplicationService); + mockContext = getMockAllowContext(TEST_ACTOR_URN); + mockEnv = Mockito.mock(DataFetchingEnvironment.class); + Mockito.when(mockEnv.getContext()).thenReturn(mockContext); + } + + private void mockExists(Urn urn, boolean exists) { + Mockito.when(mockApplicationService.verifyEntityExists(any(), eq(urn))).thenReturn(exists); + } + + @Test + public void testGetSuccessUnsetApplication() throws Exception { + mockExists(UrnUtils.getUrn(TEST_ENTITY_URN_1), true); + mockExists(UrnUtils.getUrn(TEST_ENTITY_URN_2), true); + mockExists(UrnUtils.getUrn(TEST_APPLICATION_URN), true); + + BatchUnsetApplicationInput input = + new BatchUnsetApplicationInput( + TEST_APPLICATION_URN, ImmutableList.of(TEST_ENTITY_URN_1, TEST_ENTITY_URN_2)); + Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(input); + + assertTrue(resolver.get(mockEnv).get()); + + Mockito.verify(mockApplicationService, Mockito.times(1)) + .batchUnsetApplication( + any(), + eq(UrnUtils.getUrn(TEST_APPLICATION_URN)), + eq( + ImmutableList.of( + UrnUtils.getUrn(TEST_ENTITY_URN_1), UrnUtils.getUrn(TEST_ENTITY_URN_2))), + eq(UrnUtils.getUrn(TEST_ACTOR_URN))); + } + + @Test + public void testGetSuccessUnsetApplicationSingleResource() throws Exception { + mockExists(UrnUtils.getUrn(TEST_ENTITY_URN_1), true); + mockExists(UrnUtils.getUrn(TEST_APPLICATION_URN), true); + + BatchUnsetApplicationInput input = + new BatchUnsetApplicationInput(TEST_APPLICATION_URN, ImmutableList.of(TEST_ENTITY_URN_1)); + Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(input); + + assertTrue(resolver.get(mockEnv).get()); + + Mockito.verify(mockApplicationService, Mockito.times(1)) + .batchUnsetApplication( + any(), + eq(UrnUtils.getUrn(TEST_APPLICATION_URN)), + eq(ImmutableList.of(UrnUtils.getUrn(TEST_ENTITY_URN_1))), + eq(UrnUtils.getUrn(TEST_ACTOR_URN))); + } + + @Test + public void testGetFailureApplicationDoesNotExist() { + mockExists(UrnUtils.getUrn(TEST_ENTITY_URN_1), true); + mockExists(UrnUtils.getUrn(TEST_APPLICATION_URN), false); // Application does not exist + + BatchUnsetApplicationInput input = + new BatchUnsetApplicationInput(TEST_APPLICATION_URN, ImmutableList.of(TEST_ENTITY_URN_1)); + Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(input); + + assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); + Mockito.verify(mockApplicationService, Mockito.never()) + .batchUnsetApplication(any(), any(), any(), any()); + } + + @Test + public void testGetFailureResourceDoesNotExist() { + mockExists(UrnUtils.getUrn(TEST_ENTITY_URN_1), false); // Resource does not exist + mockExists(UrnUtils.getUrn(TEST_APPLICATION_URN), true); + + BatchUnsetApplicationInput input = + new BatchUnsetApplicationInput(TEST_APPLICATION_URN, ImmutableList.of(TEST_ENTITY_URN_1)); + Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(input); + + assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); + Mockito.verify(mockApplicationService, Mockito.never()) + .batchUnsetApplication(any(), any(), any(), any()); + } + + @Test + public void testGetUnauthorized() { + QueryContext mockDenyContext = getMockDenyContext(); + Mockito.when(mockEnv.getContext()).thenReturn(mockDenyContext); + + BatchUnsetApplicationInput input = + new BatchUnsetApplicationInput(TEST_APPLICATION_URN, ImmutableList.of(TEST_ENTITY_URN_1)); + Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(input); + + assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); + Mockito.verify(mockApplicationService, Mockito.never()) + .batchUnsetApplication(any(), any(), any(), any()); + } + + @Test + public void testGetVerifiesAllResourcesBeforeCallingService() throws Exception { + mockExists(UrnUtils.getUrn(TEST_ENTITY_URN_1), true); + mockExists(UrnUtils.getUrn(TEST_ENTITY_URN_2), true); + mockExists(UrnUtils.getUrn(TEST_APPLICATION_URN), true); + + BatchUnsetApplicationInput input = + new BatchUnsetApplicationInput( + TEST_APPLICATION_URN, ImmutableList.of(TEST_ENTITY_URN_1, TEST_ENTITY_URN_2)); + Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(input); + + assertTrue(resolver.get(mockEnv).get()); + + // Verify that all resources were checked for existence + Mockito.verify(mockApplicationService, Mockito.times(1)) + .verifyEntityExists(any(), eq(UrnUtils.getUrn(TEST_ENTITY_URN_1))); + Mockito.verify(mockApplicationService, Mockito.times(1)) + .verifyEntityExists(any(), eq(UrnUtils.getUrn(TEST_ENTITY_URN_2))); + Mockito.verify(mockApplicationService, Mockito.times(1)) + .verifyEntityExists(any(), eq(UrnUtils.getUrn(TEST_APPLICATION_URN))); + + // Verify batchUnsetApplication was called with correct parameters + Mockito.verify(mockApplicationService, Mockito.times(1)) + .batchUnsetApplication( + any(), + eq(UrnUtils.getUrn(TEST_APPLICATION_URN)), + eq( + ImmutableList.of( + UrnUtils.getUrn(TEST_ENTITY_URN_1), UrnUtils.getUrn(TEST_ENTITY_URN_2))), + eq(UrnUtils.getUrn(TEST_ACTOR_URN))); + } + + @Test + public void testGetFailureServiceThrowsException() { + mockExists(UrnUtils.getUrn(TEST_ENTITY_URN_1), true); + mockExists(UrnUtils.getUrn(TEST_APPLICATION_URN), true); + + Mockito.doThrow(new RuntimeException("Service error")) + .when(mockApplicationService) + .batchUnsetApplication(any(), any(), any(), any()); + + BatchUnsetApplicationInput input = + new BatchUnsetApplicationInput(TEST_APPLICATION_URN, ImmutableList.of(TEST_ENTITY_URN_1)); + Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(input); + + assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); + Mockito.verify(mockApplicationService, Mockito.times(1)) + .batchUnsetApplication(any(), any(), any(), any()); + } +} diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/types/dataproduct/mappers/DataProductMapperTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/types/dataproduct/mappers/DataProductMapperTest.java index 32bcd6607e3407..9cdce5150633e7 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/types/dataproduct/mappers/DataProductMapperTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/types/dataproduct/mappers/DataProductMapperTest.java @@ -159,8 +159,8 @@ public void testMapDataProductWithAllAspects() throws URISyntaxException { assertNotNull(result.getGlossaryTerms()); // Domain association might be null if DomainAssociationMapper returns null // assertNotNull(result.getDomain()); - // Application association might be null if ApplicationAssociationMapper returns null - // assertNotNull(result.getApplication()); + // Applications list might be null if ApplicationsMapper returns null + // assertNotNull(result.getApplications()); } } @@ -192,7 +192,7 @@ public void testMapDataProductWithMinimalAspects() { assertNull(result.getTags()); assertNull(result.getGlossaryTerms()); assertNull(result.getDomain()); - assertNull(result.getApplication()); + assertNull(result.getApplications()); } } diff --git a/datahub-web-react/src/Mocks.tsx b/datahub-web-react/src/Mocks.tsx index 774c6cc79eabf0..b7688dadcda9b4 100644 --- a/datahub-web-react/src/Mocks.tsx +++ b/datahub-web-react/src/Mocks.tsx @@ -347,7 +347,7 @@ export const dataset1 = { }, ], domain: null, - application: null, + applications: null, container: null, health: [], assertions: null, @@ -446,7 +446,7 @@ export const dataset2 = { }, ], domain: null, - application: null, + applications: null, container: null, health: [], assertions: null, @@ -700,7 +700,7 @@ export const dataset3 = { }, ], domain: null, - application: null, + applications: null, container: null, lineage: null, relationships: null, @@ -1446,7 +1446,7 @@ export const dataFlow1 = { }, }, domain: null, - application: null, + applications: null, deprecation: null, autoRenderAspects: [], activeIncidents: null, @@ -1537,7 +1537,7 @@ export const dataJob1 = { ], }, domain: null, - application: null, + applications: null, status: null, deprecation: null, autoRenderAspects: [], @@ -1712,7 +1712,7 @@ export const dataJob2 = { ], }, domain: null, - application: null, + applications: null, upstream: null, downstream: null, deprecation: null, @@ -1789,7 +1789,7 @@ export const dataJob3 = { ], }, domain: null, - application: null, + applications: null, upstream: null, downstream: null, status: null, diff --git a/datahub-web-react/src/app/domain/CreateDomainModal.tsx b/datahub-web-react/src/app/domain/CreateDomainModal.tsx index da49f25058bf30..5d832bda067937 100644 --- a/datahub-web-react/src/app/domain/CreateDomainModal.tsx +++ b/datahub-web-react/src/app/domain/CreateDomainModal.tsx @@ -105,8 +105,10 @@ export default function CreateDomainModal({ onClose, onCreate }: Props) { .catch((e) => { message.destroy(); message.error({ content: `Failed to create Domain!: \n ${e.message || ''}`, duration: 3 }); + }) + .finally(() => { + onClose(); }); - onClose(); }; // Handle the Enter press diff --git a/datahub-web-react/src/app/domain/DomainsList.tsx b/datahub-web-react/src/app/domain/DomainsList.tsx index 8787a67d68c322..4d786530ec3aae 100644 --- a/datahub-web-react/src/app/domain/DomainsList.tsx +++ b/datahub-web-react/src/app/domain/DomainsList.tsx @@ -198,6 +198,7 @@ export const DomainsList = () => { }, ownership: null, entities: null, + applicationsInDomain: null, displayProperties: null, institutionalMemory: null, }, diff --git a/datahub-web-react/src/app/domain/utils.ts b/datahub-web-react/src/app/domain/utils.ts index 9bb6ad994e6e3e..2676cb43ca539c 100644 --- a/datahub-web-react/src/app/domain/utils.ts +++ b/datahub-web-react/src/app/domain/utils.ts @@ -74,6 +74,7 @@ export const updateListDomainsCache = ( entities: null, children: null, dataProducts: null, + applicationsInDomain: null, parentDomains: null, displayProperties: null, institutionalMemory: null, diff --git a/datahub-web-react/src/app/domainV2/CreateDomainModal.tsx b/datahub-web-react/src/app/domainV2/CreateDomainModal.tsx index 5dae929b3ca896..64343516423f23 100644 --- a/datahub-web-react/src/app/domainV2/CreateDomainModal.tsx +++ b/datahub-web-react/src/app/domainV2/CreateDomainModal.tsx @@ -130,8 +130,10 @@ export default function CreateDomainModal({ onClose, onCreate }: Props) { .catch((e) => { message.destroy(); message.error({ content: `Failed to create Domain!: \n ${e.message || ''}`, duration: 3 }); + }) + .finally(() => { + onClose(); }); - onClose(); }; // Handle the Enter press diff --git a/datahub-web-react/src/app/entity/shared/types.ts b/datahub-web-react/src/app/entity/shared/types.ts index eda056f17efa54..b576f10fba8482 100644 --- a/datahub-web-react/src/app/entity/shared/types.ts +++ b/datahub-web-react/src/app/entity/shared/types.ts @@ -96,7 +96,7 @@ export type GenericEntityProperties = { glossaryTerms?: Maybe; ownership?: Maybe; domain?: Maybe; - application?: Maybe; + applications?: Maybe; dataProduct?: Maybe; platform?: Maybe; dataPlatformInstance?: Maybe; diff --git a/datahub-web-react/src/app/entityV2/shared/containers/profile/sidebar/Applications/SidebarApplicationSection.tsx b/datahub-web-react/src/app/entityV2/shared/containers/profile/sidebar/Applications/SidebarApplicationSection.tsx index e929bde174ba54..50832c47439d07 100644 --- a/datahub-web-react/src/app/entityV2/shared/containers/profile/sidebar/Applications/SidebarApplicationSection.tsx +++ b/datahub-web-react/src/app/entityV2/shared/containers/profile/sidebar/Applications/SidebarApplicationSection.tsx @@ -13,7 +13,8 @@ import { SidebarSection } from '@app/entityV2/shared/containers/profile/sidebar/ import { ApplicationLink } from '@app/shared/tags/ApplicationLink'; import { useAppConfig } from '@app/useAppConfig'; -import { useBatchSetApplicationMutation } from '@graphql/application.generated'; +import { useBatchUnsetApplicationMutation } from '@graphql/application.generated'; +import { ApplicationAssociation } from '@types'; const Content = styled.div` display: flex; @@ -24,7 +25,7 @@ const Content = styled.div` `; const ApplicationLinkWrapper = styled.div` - margin-right: 12px; + margin-right: 6px; display: flex; align-items: center; `; @@ -46,20 +47,20 @@ export const SidebarApplicationSection = ({ readOnly, properties }: Props) => { const { entityData } = useEntityData(); const refetch = useRefetch(); const urn = useMutationUrn(); - const [batchSetApplicationMutation] = useBatchSetApplicationMutation(); + const [batchUnsetApplicationMutation] = useBatchUnsetApplicationMutation(); const [showModal, setShowModal] = useState(false); - const application = entityData?.application?.application; + const applications = entityData?.applications || []; const canEditApplication = !!entityData?.privileges?.canEditProperties; - if (!application && !visualConfig.application?.showSidebarSectionWhenEmpty) { + if (applications.length === 0 && !visualConfig.application?.showSidebarSectionWhenEmpty) { return null; } - const removeApplication = () => { - batchSetApplicationMutation({ + const removeApplication = (applicationUrn: string) => { + batchUnsetApplicationMutation({ variables: { input: { - applicationUrn: null, + applicationUrn, resourceUrns: [urn], }, }, @@ -76,12 +77,12 @@ export const SidebarApplicationSection = ({ readOnly, properties }: Props) => { }); }; - const onRemoveApplication = () => { + const onRemoveApplication = (applicationUrn: string) => { Modal.confirm({ title: `Confirm Application Removal`, content: `Are you sure you want to remove this application?`, onOk() { - removeApplication(); + removeApplication(applicationUrn); }, onCancel() {}, okText: 'Yes', @@ -93,25 +94,32 @@ export const SidebarApplicationSection = ({ readOnly, properties }: Props) => { return (
- {application && ( - - { - e.preventDefault(); - onRemoveApplication(); - }} - fontSize={12} - /> - - )} - {(!application || !!updateOnly) && ( - <>{!application && } + {applications.length > 0 && + applications.map((appAssociation: ApplicationAssociation) => ( + + { + e.preventDefault(); + if (appAssociation.application?.urn) { + onRemoveApplication(appAssociation.application.urn); + } + }} + fontSize={12} + /> + + ))} + {(applications.length === 0 || !!updateOnly) && ( + <> + {applications.length === 0 && ( + + )} + )} } @@ -119,7 +127,7 @@ export const SidebarApplicationSection = ({ readOnly, properties }: Props) => { !readOnly && ( : } + button={applications.length > 0 ? : } onClick={(event) => { setShowModal(true); event.stopPropagation(); diff --git a/datahub-web-react/src/graphql/application.graphql b/datahub-web-react/src/graphql/application.graphql index 6b9920890ddd87..64181f5e29cd66 100644 --- a/datahub-web-react/src/graphql/application.graphql +++ b/datahub-web-react/src/graphql/application.graphql @@ -93,3 +93,7 @@ mutation deleteApplication($urn: String!) { mutation batchSetApplication($input: BatchSetApplicationInput!) { batchSetApplication(input: $input) } + +mutation batchUnsetApplication($input: BatchUnsetApplicationInput!) { + batchUnsetApplication(input: $input) +} diff --git a/datahub-web-react/src/graphql/chart.graphql b/datahub-web-react/src/graphql/chart.graphql index 80abe7026d3bc7..d46474825dbdf3 100644 --- a/datahub-web-react/src/graphql/chart.graphql +++ b/datahub-web-react/src/graphql/chart.graphql @@ -49,7 +49,7 @@ query getChart($urn: String!) { domain { ...entityDomain } - application { + applications { ...entityApplication } ...entityDataProduct diff --git a/datahub-web-react/src/graphql/fragments.graphql b/datahub-web-react/src/graphql/fragments.graphql index 743744d2036f90..4427724a4c38df 100644 --- a/datahub-web-react/src/graphql/fragments.graphql +++ b/datahub-web-react/src/graphql/fragments.graphql @@ -468,7 +468,7 @@ fragment nonRecursiveDatasetFields on Dataset { domain { ...entityDomain } - application { + applications { ...entityApplication } container { @@ -513,7 +513,7 @@ fragment nonRecursiveDataFlowFields on DataFlow { domain { ...entityDomain } - application { + applications { ...entityApplication } ...entityDataProduct @@ -542,7 +542,7 @@ fragment nonRecursiveDataJobFields on DataJob { domain { ...entityDomain } - application { + applications { ...entityApplication } ...entityDataProduct @@ -593,7 +593,7 @@ fragment dataJobFields on DataJob { domain { ...entityDomain } - application { + applications { ...entityApplication } ...entityDataProduct @@ -677,7 +677,7 @@ fragment dashboardFields on Dashboard { domain { ...entityDomain } - application { + applications { ...entityApplication } ...entityDataProduct @@ -776,7 +776,7 @@ fragment nonRecursiveMLFeature on MLFeature { domain { ...entityDomain } - application { + applications { ...entityApplication } ...entityDataProduct @@ -930,7 +930,7 @@ fragment nonRecursiveMLFeatureTable on MLFeatureTable { domain { ...entityDomain } - application { + applications { ...entityApplication } ...entityDataProduct @@ -1134,7 +1134,7 @@ fragment nonRecursiveMLModel on MLModel { domain { ...entityDomain } - application { + applications { ...entityApplication } ...entityDataProduct @@ -1181,7 +1181,7 @@ fragment nonRecursiveMLModelGroupFields on MLModelGroup { domain { ...entityDomain } - application { + applications { ...entityApplication } ...entityDataProduct @@ -1303,7 +1303,9 @@ fragment domainEntitiesFields on Domain { dataProducts: entities(input: { start: 0, count: 0, filters: [{ field: "_entityType", values: "DATA_PRODUCT" }] }) { total } - applications: entities(input: { start: 0, count: 0, filters: [{ field: "_entityType", values: "APPLICATION" }] }) { + applicationsInDomain: entities( + input: { start: 0, count: 0, filters: [{ field: "_entityType", values: "APPLICATION" }] } + ) { total } children: relationships(input: { types: ["IsPartOf"], direction: INCOMING, start: 0, count: 0 }) { diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/service/ApplicationService.java b/metadata-service/services/src/main/java/com/linkedin/metadata/service/ApplicationService.java index 4a9fd00bf9de98..0a5158c7ae58c0 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/service/ApplicationService.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/service/ApplicationService.java @@ -187,9 +187,17 @@ public void batchSetApplicationAssets( Objects.requireNonNull(applicationUrn, "applicationUrn must not be null"); Objects.requireNonNull(resourceUrns, "resourceUrns must not be null"); + log.info( + "Batch setting application {} to {} resource(s): {}", + applicationUrn, + resourceUrns.size(), + resourceUrns); + final List proposals = resourceUrns.stream() - .map(resourceUrn -> buildSetApplicationAssetsProposal(applicationUrn, resourceUrn)) + .map( + resourceUrn -> + buildAddApplicationAssetsProposal(opContext, applicationUrn, resourceUrn)) .collect(Collectors.toList()); try { @@ -202,6 +210,23 @@ public void batchSetApplicationAssets( } } + private MetadataChangeProposal buildAddApplicationAssetsProposal( + @Nonnull OperationContext opContext, @Nonnull Urn applicationUrn, Urn resourceUrn) { + Applications applications = getExistingApplications(opContext, resourceUrn); + + if (!applications.getApplications().contains(applicationUrn)) { + applications.getApplications().add(applicationUrn); + } else { + log.info( + "Application {} already exists on resource {}. Skipping duplicate.", + applicationUrn, + resourceUrn); + } + + return AspectUtils.buildMetadataChangeProposal( + resourceUrn, Constants.APPLICATION_MEMBERSHIP_ASPECT_NAME, applications); + } + private MetadataChangeProposal buildSetApplicationAssetsProposal( @Nullable Urn applicationUrn, Urn resourceUrn) { Applications applications = new Applications(new DataMap()); @@ -214,6 +239,42 @@ private MetadataChangeProposal buildSetApplicationAssetsProposal( resourceUrn, Constants.APPLICATION_MEMBERSHIP_ASPECT_NAME, applications); } + private Applications getExistingApplications( + @Nonnull OperationContext opContext, @Nonnull Urn resourceUrn) { + try { + final EntityResponse response = + this.entityClient.getV2( + opContext, + resourceUrn.getEntityType(), + resourceUrn, + ImmutableSet.of(Constants.APPLICATION_MEMBERSHIP_ASPECT_NAME)); + + if (response != null + && response.getAspects().containsKey(Constants.APPLICATION_MEMBERSHIP_ASPECT_NAME)) { + return new Applications( + response + .getAspects() + .get(Constants.APPLICATION_MEMBERSHIP_ASPECT_NAME) + .getValue() + .data()); + } else { + log.info( + "No existing applications aspect found for resource {} (response: {})", + resourceUrn, + response != null ? "present but no aspect" : "null"); + } + } catch (Exception e) { + log.warn( + "Failed to retrieve existing Applications for resource {}, will create new aspect", + resourceUrn, + e); + } + + Applications applications = new Applications(new DataMap()); + applications.setApplications(new UrnArray()); + return applications; + } + public void batchRemoveApplicationAssets( @Nonnull OperationContext opContext, @Nonnull Urn applicationUrn, @@ -255,6 +316,46 @@ public void unsetApplication( } } + public void batchUnsetApplication( + @Nonnull OperationContext opContext, + @Nonnull Urn applicationUrn, + @Nonnull List resourceUrns, + @Nonnull Urn actorUrn) { + Objects.requireNonNull(applicationUrn, "applicationUrn must not be null"); + Objects.requireNonNull(resourceUrns, "resourceUrns must not be null"); + Objects.requireNonNull(actorUrn, "actorUrn must not be null"); + Objects.requireNonNull(opContext.getSessionAuthentication(), "authentication must not be null"); + + log.info( + "Batch unsetting application {} from {} resources", applicationUrn, resourceUrns.size()); + + final List proposals = + resourceUrns.stream() + .map( + resourceUrn -> + buildUnsetApplicationProposal(opContext, applicationUrn, resourceUrn)) + .collect(Collectors.toList()); + + try { + this.entityClient.batchIngestProposals(opContext, proposals, false); + } catch (Exception e) { + log.error("Failed to batch unset application {} from resources", applicationUrn, e); + throw new RuntimeException( + String.format("Failed to batch unset application %s from resources", applicationUrn), e); + } + } + + private MetadataChangeProposal buildUnsetApplicationProposal( + @Nonnull OperationContext opContext, @Nonnull Urn applicationUrn, Urn resourceUrn) { + Applications applications = getExistingApplications(opContext, resourceUrn); + + // Remove the specific application from the list + applications.getApplications().remove(applicationUrn); + + return AspectUtils.buildMetadataChangeProposal( + resourceUrn, Constants.APPLICATION_MEMBERSHIP_ASPECT_NAME, applications); + } + public boolean verifyEntityExists(@Nonnull OperationContext opContext, @Nonnull Urn entityUrn) { try { return this.entityClient.exists(opContext, entityUrn); diff --git a/metadata-service/services/src/test/java/com/linkedin/metadata/service/ApplicationServiceTest.java b/metadata-service/services/src/test/java/com/linkedin/metadata/service/ApplicationServiceTest.java index d3c85260cd7d25..cf2509a155514b 100644 --- a/metadata-service/services/src/test/java/com/linkedin/metadata/service/ApplicationServiceTest.java +++ b/metadata-service/services/src/test/java/com/linkedin/metadata/service/ApplicationServiceTest.java @@ -13,6 +13,7 @@ import com.linkedin.application.ApplicationKey; import com.linkedin.application.ApplicationProperties; import com.linkedin.application.Applications; +import com.linkedin.common.UrnArray; import com.linkedin.common.urn.Urn; import com.linkedin.common.urn.UrnUtils; import com.linkedin.domain.Domains; @@ -313,4 +314,193 @@ public void testVerifyEntityExists() throws Exception { RuntimeException.class, () -> _applicationService.verifyEntityExists(_opContext, TEST_DOMAIN_URN)); } + + @Test + public void testBatchUnsetApplicationFromEmptyList() throws Exception { + // Mock: Resource has no existing applications + when(_entityClient.getV2( + any(OperationContext.class), + eq(TEST_ASSET_URN.getEntityType()), + eq(TEST_ASSET_URN), + eq(ImmutableSet.of(Constants.APPLICATION_MEMBERSHIP_ASPECT_NAME)))) + .thenReturn(null); + + List resourceUrns = ImmutableList.of(TEST_ASSET_URN); + _applicationService.batchUnsetApplication( + _opContext, TEST_APPLICATION_URN, resourceUrns, TEST_USER_URN); + + ArgumentCaptor> mcpListCaptor = + ArgumentCaptor.forClass(List.class); + verify(_entityClient, times(1)) + .batchIngestProposals(eq(_opContext), mcpListCaptor.capture(), eq(false)); + + List mcps = mcpListCaptor.getValue(); + assertEquals(mcps.size(), 1); + + MetadataChangeProposal mcp = mcps.get(0); + assertEquals(mcp.getEntityUrn(), TEST_ASSET_URN); + assertEquals(mcp.getAspectName(), Constants.APPLICATION_MEMBERSHIP_ASPECT_NAME); + + Applications apps = + GenericRecordUtils.deserializeAspect( + mcp.getAspect().getValue(), mcp.getAspect().getContentType(), Applications.class); + assertTrue(apps.getApplications().isEmpty()); + } + + @Test + public void testBatchUnsetApplicationRemovesSpecificApplication() throws Exception { + // Mock: Resource has two applications, we want to remove one + Applications existingApps = new Applications(); + existingApps.setApplications( + new UrnArray(ImmutableList.of(TEST_APPLICATION_URN, TEST_APPLICATION_URN_2))); + + EntityResponse response = new EntityResponse(); + EnvelopedAspect envelopedAspect = new EnvelopedAspect(); + envelopedAspect.setValue(new Aspect(existingApps.data())); + response.setAspects( + new EnvelopedAspectMap( + Collections.singletonMap( + Constants.APPLICATION_MEMBERSHIP_ASPECT_NAME, envelopedAspect))); + response.setEntityName(TEST_ASSET_URN.getEntityType()); + response.setUrn(TEST_ASSET_URN); + + when(_entityClient.getV2( + any(OperationContext.class), + eq(TEST_ASSET_URN.getEntityType()), + eq(TEST_ASSET_URN), + eq(ImmutableSet.of(Constants.APPLICATION_MEMBERSHIP_ASPECT_NAME)))) + .thenReturn(response); + + List resourceUrns = ImmutableList.of(TEST_ASSET_URN); + _applicationService.batchUnsetApplication( + _opContext, TEST_APPLICATION_URN, resourceUrns, TEST_USER_URN); + + ArgumentCaptor> mcpListCaptor = + ArgumentCaptor.forClass(List.class); + verify(_entityClient, times(1)) + .batchIngestProposals(eq(_opContext), mcpListCaptor.capture(), eq(false)); + + List mcps = mcpListCaptor.getValue(); + assertEquals(mcps.size(), 1); + + MetadataChangeProposal mcp = mcps.get(0); + Applications apps = + GenericRecordUtils.deserializeAspect( + mcp.getAspect().getValue(), mcp.getAspect().getContentType(), Applications.class); + + // Should only have TEST_APPLICATION_URN_2 left + assertEquals(apps.getApplications().size(), 1); + assertEquals(apps.getApplications().get(0), TEST_APPLICATION_URN_2); + assertFalse(apps.getApplications().contains(TEST_APPLICATION_URN)); + } + + @Test + public void testBatchUnsetApplicationHandlesException() throws Exception { + List resourceUrns = ImmutableList.of(TEST_ASSET_URN); + + when(_entityClient.getV2( + any(OperationContext.class), + eq(TEST_ASSET_URN.getEntityType()), + eq(TEST_ASSET_URN), + eq(ImmutableSet.of(Constants.APPLICATION_MEMBERSHIP_ASPECT_NAME)))) + .thenReturn(null); + + when(_entityClient.batchIngestProposals(eq(_opContext), any(List.class), eq(false))) + .thenThrow(new RuntimeException("Batch ingest failed")); + + assertThrows( + RuntimeException.class, + () -> + _applicationService.batchUnsetApplication( + _opContext, TEST_APPLICATION_URN, resourceUrns, TEST_USER_URN)); + } + + @Test + public void testBatchSetApplicationAssetsWithExistingApplications() throws Exception { + // Mock: Resource already has TEST_APPLICATION_URN_2, we're adding TEST_APPLICATION_URN + Applications existingApps = new Applications(); + existingApps.setApplications(new UrnArray(ImmutableList.of(TEST_APPLICATION_URN_2))); + + EntityResponse response = new EntityResponse(); + EnvelopedAspect envelopedAspect = new EnvelopedAspect(); + envelopedAspect.setValue(new Aspect(existingApps.data())); + response.setAspects( + new EnvelopedAspectMap( + Collections.singletonMap( + Constants.APPLICATION_MEMBERSHIP_ASPECT_NAME, envelopedAspect))); + response.setUrn(TEST_ASSET_URN); + + when(_entityClient.getV2( + any(OperationContext.class), + eq(TEST_ASSET_URN.getEntityType()), + eq(TEST_ASSET_URN), + eq(ImmutableSet.of(Constants.APPLICATION_MEMBERSHIP_ASPECT_NAME)))) + .thenReturn(response); + + List assetUrns = ImmutableList.of(TEST_ASSET_URN); + _applicationService.batchSetApplicationAssets( + _opContext, TEST_APPLICATION_URN, assetUrns, TEST_USER_URN); + + ArgumentCaptor> mcpListCaptor = + ArgumentCaptor.forClass(List.class); + verify(_entityClient, times(1)) + .batchIngestProposals(eq(_opContext), mcpListCaptor.capture(), eq(false)); + + List mcps = mcpListCaptor.getValue(); + assertEquals(mcps.size(), 1); + + MetadataChangeProposal mcp = mcps.get(0); + Applications apps = + GenericRecordUtils.deserializeAspect( + mcp.getAspect().getValue(), mcp.getAspect().getContentType(), Applications.class); + + // Should now have both applications + assertEquals(apps.getApplications().size(), 2); + assertTrue(apps.getApplications().contains(TEST_APPLICATION_URN)); + assertTrue(apps.getApplications().contains(TEST_APPLICATION_URN_2)); + } + + @Test + public void testBatchSetApplicationAssetsSkipsDuplicates() throws Exception { + // Mock: Resource already has TEST_APPLICATION_URN + Applications existingApps = new Applications(); + existingApps.setApplications(new UrnArray(ImmutableList.of(TEST_APPLICATION_URN))); + + EntityResponse response = new EntityResponse(); + EnvelopedAspect envelopedAspect = new EnvelopedAspect(); + envelopedAspect.setValue(new Aspect(existingApps.data())); + response.setAspects( + new EnvelopedAspectMap( + Collections.singletonMap( + Constants.APPLICATION_MEMBERSHIP_ASPECT_NAME, envelopedAspect))); + response.setUrn(TEST_ASSET_URN); + + when(_entityClient.getV2( + any(OperationContext.class), + eq(TEST_ASSET_URN.getEntityType()), + eq(TEST_ASSET_URN), + eq(ImmutableSet.of(Constants.APPLICATION_MEMBERSHIP_ASPECT_NAME)))) + .thenReturn(response); + + List assetUrns = ImmutableList.of(TEST_ASSET_URN); + _applicationService.batchSetApplicationAssets( + _opContext, TEST_APPLICATION_URN, assetUrns, TEST_USER_URN); + + ArgumentCaptor> mcpListCaptor = + ArgumentCaptor.forClass(List.class); + verify(_entityClient, times(1)) + .batchIngestProposals(eq(_opContext), mcpListCaptor.capture(), eq(false)); + + List mcps = mcpListCaptor.getValue(); + assertEquals(mcps.size(), 1); + + MetadataChangeProposal mcp = mcps.get(0); + Applications apps = + GenericRecordUtils.deserializeAspect( + mcp.getAspect().getValue(), mcp.getAspect().getContentType(), Applications.class); + + // Should still only have one application (no duplicate) + assertEquals(apps.getApplications().size(), 1); + assertEquals(apps.getApplications().get(0), TEST_APPLICATION_URN); + } }