From 3f3744da4b0a63d68dfb03941e4d39e2017871e6 Mon Sep 17 00:00:00 2001 From: Chris Collins Date: Mon, 1 Jul 2024 18:41:13 -0400 Subject: [PATCH 1/3] add removeStructuredProperties graphql endpoint --- .../datahub/graphql/GmsGraphQLEngine.java | 4 + .../RemoveStructuredPropertiesResolver.java | 103 +++++++++++++++ .../src/main/resources/properties.graphql | 20 +++ ...emoveStructuredPropertiesResolverTest.java | 123 ++++++++++++++++++ 4 files changed, 250 insertions(+) create mode 100644 datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/structuredproperties/RemoveStructuredPropertiesResolver.java create mode 100644 datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/structuredproperties/RemoveStructuredPropertiesResolverTest.java 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 6f2e250c17c34..1897d76485661 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 @@ -285,6 +285,7 @@ import com.linkedin.datahub.graphql.resolvers.settings.view.UpdateGlobalViewsSettingsResolver; import com.linkedin.datahub.graphql.resolvers.step.BatchGetStepStatesResolver; import com.linkedin.datahub.graphql.resolvers.step.BatchUpdateStepStatesResolver; +import com.linkedin.datahub.graphql.resolvers.structuredproperties.RemoveStructuredPropertiesResolver; import com.linkedin.datahub.graphql.resolvers.structuredproperties.UpsertStructuredPropertiesResolver; import com.linkedin.datahub.graphql.resolvers.tag.CreateTagResolver; import com.linkedin.datahub.graphql.resolvers.tag.DeleteTagResolver; @@ -1316,6 +1317,9 @@ private void configureMutationResolvers(final RuntimeWiring.Builder builder) { .dataFetcher( "upsertStructuredProperties", new UpsertStructuredPropertiesResolver(this.entityClient)) + .dataFetcher( + "removeStructuredProperties", + new RemoveStructuredPropertiesResolver(this.entityClient)) .dataFetcher("raiseIncident", new RaiseIncidentResolver(this.entityClient)) .dataFetcher( "updateIncidentStatus", diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/structuredproperties/RemoveStructuredPropertiesResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/structuredproperties/RemoveStructuredPropertiesResolver.java new file mode 100644 index 0000000000000..ea8c6dac36a4a --- /dev/null +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/structuredproperties/RemoveStructuredPropertiesResolver.java @@ -0,0 +1,103 @@ +package com.linkedin.datahub.graphql.resolvers.structuredproperties; + +import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.bindArgument; + +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.authorization.AuthorizationUtils; +import com.linkedin.datahub.graphql.exception.AuthorizationException; +import com.linkedin.datahub.graphql.generated.RemoveStructuredPropertiesInput; +import com.linkedin.datahub.graphql.types.structuredproperty.StructuredPropertiesMapper; +import com.linkedin.entity.EntityResponse; +import com.linkedin.entity.client.EntityClient; +import com.linkedin.metadata.Constants; +import com.linkedin.metadata.aspect.patch.builder.StructuredPropertiesPatchBuilder; +import com.linkedin.mxe.MetadataChangeProposal; +import com.linkedin.structured.StructuredProperties; +import graphql.schema.DataFetcher; +import graphql.schema.DataFetchingEnvironment; +import java.util.Objects; +import java.util.concurrent.CompletableFuture; +import javax.annotation.Nonnull; + +public class RemoveStructuredPropertiesResolver + implements DataFetcher< + CompletableFuture> { + + private final EntityClient _entityClient; + + public RemoveStructuredPropertiesResolver(@Nonnull final EntityClient entityClient) { + _entityClient = Objects.requireNonNull(entityClient, "entityClient must not be null"); + } + + @Override + public CompletableFuture get( + final DataFetchingEnvironment environment) throws Exception { + final QueryContext context = environment.getContext(); + + final RemoveStructuredPropertiesInput input = + bindArgument(environment.getArgument("input"), RemoveStructuredPropertiesInput.class); + final Urn assetUrn = UrnUtils.getUrn(input.getAssetUrn()); + + return CompletableFuture.supplyAsync( + () -> { + try { + // check authorization first + if (!AuthorizationUtils.canEditProperties(assetUrn, context)) { + throw new AuthorizationException( + String.format( + "Not authorized to update properties on the gives urn %s", assetUrn)); + } + + if (!_entityClient.exists(context.getOperationContext(), assetUrn)) { + throw new RuntimeException( + String.format("Asset with provided urn %s does not exist", assetUrn)); + } + + StructuredPropertiesPatchBuilder patchBuilder = + new StructuredPropertiesPatchBuilder().urn(assetUrn); + + input + .getStructuredPropertyUrns() + .forEach( + propertyUrn -> { + patchBuilder.removeProperty(UrnUtils.getUrn(propertyUrn)); + }); + + // ingest change proposal + final MetadataChangeProposal structuredPropertiesProposal = patchBuilder.build(); + + _entityClient.ingestProposal( + context.getOperationContext(), structuredPropertiesProposal, false); + + EntityResponse response = + _entityClient.getV2( + context.getOperationContext(), + assetUrn.getEntityType(), + assetUrn, + ImmutableSet.of(Constants.STRUCTURED_PROPERTIES_ASPECT_NAME)); + + if (response == null + || response.getAspects().get(Constants.STRUCTURED_PROPERTIES_ASPECT_NAME) == null) { + throw new RuntimeException( + String.format("Failed to fetch structured properties from entity %s", assetUrn)); + } + + StructuredProperties structuredProperties = + new StructuredProperties( + response + .getAspects() + .get(Constants.STRUCTURED_PROPERTIES_ASPECT_NAME) + .getValue() + .data()); + + return StructuredPropertiesMapper.map(context, structuredProperties); + } catch (Exception e) { + throw new RuntimeException( + String.format("Failed to perform update against input %s", input), e); + } + }); + } +} diff --git a/datahub-graphql-core/src/main/resources/properties.graphql b/datahub-graphql-core/src/main/resources/properties.graphql index 120154e930d59..2e7a501116601 100644 --- a/datahub-graphql-core/src/main/resources/properties.graphql +++ b/datahub-graphql-core/src/main/resources/properties.graphql @@ -3,6 +3,11 @@ extend type Mutation { Upsert structured properties onto a given asset """ upsertStructuredProperties(input: UpsertStructuredPropertiesInput!): StructuredProperties! + + """ + Upsert structured properties onto a given asset + """ + removeStructuredProperties(input: RemoveStructuredPropertiesInput!): StructuredProperties! } """ @@ -184,6 +189,21 @@ input UpsertStructuredPropertiesInput { structuredPropertyInputParams: [StructuredPropertyInputParams!]! } +""" +Input for removing structured properties on a given asset +""" +input RemoveStructuredPropertiesInput { + """ + The urn of the asset that we are removing properties from + """ + assetUrn: String! + + """ + The list of structured properties you want to remove from this asset + """ + structuredPropertyUrns: [String!]! +} + """ A data type registered in DataHub """ diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/structuredproperties/RemoveStructuredPropertiesResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/structuredproperties/RemoveStructuredPropertiesResolverTest.java new file mode 100644 index 0000000000000..f7882bb2c93a8 --- /dev/null +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/structuredproperties/RemoveStructuredPropertiesResolverTest.java @@ -0,0 +1,123 @@ +package com.linkedin.datahub.graphql.resolvers.structuredproperties; + +import static com.linkedin.datahub.graphql.TestUtils.getMockAllowContext; +import static com.linkedin.datahub.graphql.TestUtils.getMockDenyContext; +import static com.linkedin.metadata.Constants.STRUCTURED_PROPERTIES_ASPECT_NAME; +import static org.mockito.ArgumentMatchers.any; +import static org.testng.Assert.assertThrows; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import com.linkedin.common.urn.UrnUtils; +import com.linkedin.datahub.graphql.QueryContext; +import com.linkedin.datahub.graphql.generated.RemoveStructuredPropertiesInput; +import com.linkedin.entity.Aspect; +import com.linkedin.entity.EntityResponse; +import com.linkedin.entity.EnvelopedAspect; +import com.linkedin.entity.EnvelopedAspectMap; +import com.linkedin.entity.client.EntityClient; +import com.linkedin.metadata.Constants; +import com.linkedin.mxe.MetadataChangeProposal; +import com.linkedin.structured.StructuredProperties; +import com.linkedin.structured.StructuredPropertyValueAssignmentArray; +import graphql.schema.DataFetchingEnvironment; +import java.util.concurrent.CompletionException; +import org.mockito.Mockito; +import org.testng.annotations.Test; + +public class RemoveStructuredPropertiesResolverTest { + + private static final String TEST_DATASET_URN = + "urn:li:dataset:(urn:li:dataPlatform:hive,name,PROD)"; + private static final String PROPERTY_URN_1 = "urn:li:structuredProperty:test1"; + private static final String PROPERTY_URN_2 = "urn:li:structuredProperty:test2"; + + private static final RemoveStructuredPropertiesInput TEST_INPUT = + new RemoveStructuredPropertiesInput( + TEST_DATASET_URN, ImmutableList.of(PROPERTY_URN_1, PROPERTY_URN_2)); + + @Test + public void testGetSuccess() throws Exception { + EntityClient mockEntityClient = initMockEntityClient(); + RemoveStructuredPropertiesResolver resolver = + new RemoveStructuredPropertiesResolver(mockEntityClient); + + // Execute resolver + QueryContext mockContext = getMockAllowContext(); + DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); + Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(TEST_INPUT); + Mockito.when(mockEnv.getContext()).thenReturn(mockContext); + + resolver.get(mockEnv).get(); + + // Validate that we called ingest + Mockito.verify(mockEntityClient, Mockito.times(1)) + .ingestProposal(any(), any(MetadataChangeProposal.class), Mockito.eq(false)); + } + + @Test + public void testGetUnauthorized() throws Exception { + EntityClient mockEntityClient = initMockEntityClient(); + RemoveStructuredPropertiesResolver resolver = + new RemoveStructuredPropertiesResolver(mockEntityClient); + + // Execute resolver + QueryContext mockContext = getMockDenyContext(); + DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); + Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(TEST_INPUT); + Mockito.when(mockEnv.getContext()).thenReturn(mockContext); + + assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); + + // Validate that we did NOT call ingest + Mockito.verify(mockEntityClient, Mockito.times(0)) + .ingestProposal(any(), any(MetadataChangeProposal.class), Mockito.eq(false)); + } + + @Test + public void testGetThrowsError() throws Exception { + // if the entity you are trying to remove properties from doesn't exist + EntityClient mockEntityClient = Mockito.mock(EntityClient.class); + Mockito.when(mockEntityClient.exists(any(), Mockito.eq(UrnUtils.getUrn(TEST_DATASET_URN)))) + .thenReturn(false); + RemoveStructuredPropertiesResolver resolver = + new RemoveStructuredPropertiesResolver(mockEntityClient); + + // Execute resolver + QueryContext mockContext = getMockDenyContext(); + DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); + Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(TEST_INPUT); + Mockito.when(mockEnv.getContext()).thenReturn(mockContext); + + assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); + + // Validate that we did NOT call ingest + Mockito.verify(mockEntityClient, Mockito.times(0)) + .ingestProposal(any(), any(MetadataChangeProposal.class), Mockito.eq(false)); + } + + private EntityClient initMockEntityClient() throws Exception { + EntityClient client = Mockito.mock(EntityClient.class); + EntityResponse response = new EntityResponse(); + response.setEntityName(Constants.DATASET_ENTITY_NAME); + response.setUrn(UrnUtils.getUrn(TEST_DATASET_URN)); + final EnvelopedAspectMap aspectMap = new EnvelopedAspectMap(); + StructuredProperties properties = new StructuredProperties(); + properties.setProperties(new StructuredPropertyValueAssignmentArray()); + aspectMap.put( + STRUCTURED_PROPERTIES_ASPECT_NAME, + new EnvelopedAspect().setValue(new Aspect(properties.data()))); + response.setAspects(aspectMap); + Mockito.when( + client.getV2( + any(), + Mockito.eq(Constants.DATASET_ENTITY_NAME), + Mockito.eq(UrnUtils.getUrn(TEST_DATASET_URN)), + Mockito.eq(ImmutableSet.of(Constants.STRUCTURED_PROPERTIES_ASPECT_NAME)))) + .thenReturn(response); + Mockito.when(client.exists(any(), Mockito.eq(UrnUtils.getUrn(TEST_DATASET_URN)))) + .thenReturn(true); + + return client; + } +} From c59deaf2c534ec3e960334c3c88378cf516fd8a4 Mon Sep 17 00:00:00 2001 From: Chris Collins Date: Mon, 1 Jul 2024 18:48:06 -0400 Subject: [PATCH 2/3] add createStructuredProperty graphql endpoint --- .../datahub/graphql/GmsGraphQLEngine.java | 4 + .../authorization/AuthorizationUtils.java | 7 + .../CreateStructuredPropertyResolver.java | 136 ++++++++++++++++++ .../src/main/resources/properties.graphql | 95 ++++++++++++ .../CreateStructuredPropertyResolverTest.java | 126 ++++++++++++++++ .../war/src/main/resources/boot/policies.json | 9 +- .../authorization/PoliciesConfig.java | 9 +- 7 files changed, 382 insertions(+), 4 deletions(-) create mode 100644 datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/structuredproperties/CreateStructuredPropertyResolver.java create mode 100644 datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/structuredproperties/CreateStructuredPropertyResolverTest.java 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 1897d76485661..ccd18aa56623d 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 @@ -285,6 +285,7 @@ import com.linkedin.datahub.graphql.resolvers.settings.view.UpdateGlobalViewsSettingsResolver; import com.linkedin.datahub.graphql.resolvers.step.BatchGetStepStatesResolver; import com.linkedin.datahub.graphql.resolvers.step.BatchUpdateStepStatesResolver; +import com.linkedin.datahub.graphql.resolvers.structuredproperties.CreateStructuredPropertyResolver; import com.linkedin.datahub.graphql.resolvers.structuredproperties.RemoveStructuredPropertiesResolver; import com.linkedin.datahub.graphql.resolvers.structuredproperties.UpsertStructuredPropertiesResolver; import com.linkedin.datahub.graphql.resolvers.tag.CreateTagResolver; @@ -1320,6 +1321,9 @@ private void configureMutationResolvers(final RuntimeWiring.Builder builder) { .dataFetcher( "removeStructuredProperties", new RemoveStructuredPropertiesResolver(this.entityClient)) + .dataFetcher( + "createStructuredProperty", + new CreateStructuredPropertyResolver(this.entityClient)) .dataFetcher("raiseIncident", new RaiseIncidentResolver(this.entityClient)) .dataFetcher( "updateIncidentStatus", diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/authorization/AuthorizationUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/authorization/AuthorizationUtils.java index c16e436a7805c..79995fc886aa7 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/authorization/AuthorizationUtils.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/authorization/AuthorizationUtils.java @@ -381,6 +381,13 @@ public static T restrictEntity(@Nonnull Object entity, Class clazz) { } } + public static boolean canManageStructuredProperties(@Nonnull QueryContext context) { + return AuthUtil.isAuthorized( + context.getAuthorizer(), + context.getActorUrn(), + PoliciesConfig.MANAGE_STRUCTURED_PROPERTIES_PRIVILEGE); + } + public static boolean isAuthorized( @Nonnull Authorizer authorizer, @Nonnull String actor, diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/structuredproperties/CreateStructuredPropertyResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/structuredproperties/CreateStructuredPropertyResolver.java new file mode 100644 index 0000000000000..3be7ea505abbf --- /dev/null +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/structuredproperties/CreateStructuredPropertyResolver.java @@ -0,0 +1,136 @@ +package com.linkedin.datahub.graphql.resolvers.structuredproperties; + +import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.bindArgument; +import static com.linkedin.metadata.Constants.STRUCTURED_PROPERTY_ENTITY_NAME; + +import com.linkedin.common.urn.Urn; +import com.linkedin.data.template.SetMode; +import com.linkedin.data.template.StringArray; +import com.linkedin.data.template.StringArrayMap; +import com.linkedin.datahub.graphql.QueryContext; +import com.linkedin.datahub.graphql.authorization.AuthorizationUtils; +import com.linkedin.datahub.graphql.exception.AuthorizationException; +import com.linkedin.datahub.graphql.generated.CreateStructuredPropertyInput; +import com.linkedin.datahub.graphql.generated.StructuredPropertyEntity; +import com.linkedin.datahub.graphql.types.structuredproperty.StructuredPropertyMapper; +import com.linkedin.entity.EntityResponse; +import com.linkedin.entity.client.EntityClient; +import com.linkedin.metadata.aspect.patch.builder.StructuredPropertyDefinitionPatchBuilder; +import com.linkedin.metadata.utils.EntityKeyUtils; +import com.linkedin.mxe.MetadataChangeProposal; +import com.linkedin.structured.PrimitivePropertyValue; +import com.linkedin.structured.PropertyCardinality; +import com.linkedin.structured.PropertyValue; +import com.linkedin.structured.StructuredPropertyKey; +import graphql.schema.DataFetcher; +import graphql.schema.DataFetchingEnvironment; +import java.util.Objects; +import java.util.UUID; +import java.util.concurrent.CompletableFuture; +import javax.annotation.Nonnull; + +public class CreateStructuredPropertyResolver + implements DataFetcher> { + + private final EntityClient _entityClient; + + public CreateStructuredPropertyResolver(@Nonnull final EntityClient entityClient) { + _entityClient = Objects.requireNonNull(entityClient, "entityClient must not be null"); + } + + @Override + public CompletableFuture get(final DataFetchingEnvironment environment) + throws Exception { + final QueryContext context = environment.getContext(); + + final CreateStructuredPropertyInput input = + bindArgument(environment.getArgument("input"), CreateStructuredPropertyInput.class); + + return CompletableFuture.supplyAsync( + () -> { + try { + if (!AuthorizationUtils.canManageStructuredProperties(context)) { + throw new AuthorizationException( + "Unable to create structured property. Please contact your admin."); + } + final StructuredPropertyKey key = new StructuredPropertyKey(); + final String id = input.getId() != null ? input.getId() : UUID.randomUUID().toString(); + key.setId(id); + final Urn propertyUrn = + EntityKeyUtils.convertEntityKeyToUrn(key, STRUCTURED_PROPERTY_ENTITY_NAME); + StructuredPropertyDefinitionPatchBuilder builder = + new StructuredPropertyDefinitionPatchBuilder().urn(propertyUrn); + + builder.setQualifiedName(input.getQualifiedName()); + builder.setValueType(input.getValueType()); + input.getEntityTypes().forEach(builder::addEntityType); + if (input.getDisplayName() != null) { + builder.setDisplayName(input.getDisplayName()); + } + if (input.getDescription() != null) { + builder.setDescription(input.getDescription()); + } + if (input.getImmutable() != null) { + builder.setImmutable(input.getImmutable()); + } + if (input.getTypeQualifier() != null) { + buildTypeQualifier(input, builder); + } + if (input.getAllowedValues() != null) { + buildAllowedValues(input, builder); + } + if (input.getCardinality() != null) { + builder.setCardinality( + PropertyCardinality.valueOf(input.getCardinality().toString())); + } + + MetadataChangeProposal mcp = builder.build(); + _entityClient.ingestProposal(context.getOperationContext(), mcp, false); + + EntityResponse response = + _entityClient.getV2( + context.getOperationContext(), + STRUCTURED_PROPERTY_ENTITY_NAME, + propertyUrn, + null); + return StructuredPropertyMapper.map(context, response); + } catch (Exception e) { + throw new RuntimeException( + String.format("Failed to perform update against input %s", input), e); + } + }); + } + + private void buildTypeQualifier( + @Nonnull final CreateStructuredPropertyInput input, + @Nonnull final StructuredPropertyDefinitionPatchBuilder builder) { + if (input.getTypeQualifier().getAllowedTypes() != null) { + final StringArrayMap typeQualifier = new StringArrayMap(); + StringArray allowedTypes = new StringArray(); + allowedTypes.addAll(input.getTypeQualifier().getAllowedTypes()); + typeQualifier.put("allowedTypes", allowedTypes); + builder.setTypeQualifier(typeQualifier); + } + } + + private void buildAllowedValues( + @Nonnull final CreateStructuredPropertyInput input, + @Nonnull final StructuredPropertyDefinitionPatchBuilder builder) { + input + .getAllowedValues() + .forEach( + allowedValueInput -> { + PropertyValue value = new PropertyValue(); + PrimitivePropertyValue primitiveValue = new PrimitivePropertyValue(); + if (allowedValueInput.getStringValue() != null) { + primitiveValue.setString(allowedValueInput.getStringValue()); + } + if (allowedValueInput.getNumberValue() != null) { + primitiveValue.setDouble(allowedValueInput.getNumberValue().doubleValue()); + } + value.setValue(primitiveValue); + value.setDescription(allowedValueInput.getDescription(), SetMode.IGNORE_NULL); + builder.addAllowedValue(value); + }); + } +} diff --git a/datahub-graphql-core/src/main/resources/properties.graphql b/datahub-graphql-core/src/main/resources/properties.graphql index 2e7a501116601..bebe7fed32e51 100644 --- a/datahub-graphql-core/src/main/resources/properties.graphql +++ b/datahub-graphql-core/src/main/resources/properties.graphql @@ -8,6 +8,11 @@ extend type Mutation { Upsert structured properties onto a given asset """ removeStructuredProperties(input: RemoveStructuredPropertiesInput!): StructuredProperties! + + """ + Create a new structured property + """ + createStructuredProperty(input: CreateStructuredPropertyInput!): StructuredPropertyEntity! } """ @@ -288,3 +293,93 @@ type DataTypeInfo { """ description: String } + +""" +Input for creating a new structured property entity +""" +input CreateStructuredPropertyInput { + """ + (Advanced) An optional unique ID to use when creating the urn of this entity + """ + id: String + + """ + The unique fully qualified name of this structured property, dot delimited. + """ + qualifiedName: String! + + """ + The optional display name for this property + """ + displayName: String + + """ + The optional description for this property + """ + description: String + + """ + Whether the property will be mutable once it is applied or not. Default is false. + """ + immutable: Boolean + + """ + The urn of the value type that this structured property accepts. + For example: urn:li:dataType:datahub.string or urn:li:dataType:datahub.date + """ + valueType: String! + + """ + The optional input for specifying specific entity types as values + """ + typeQualifier: TypeQualifierInput + + """ + The optional input for specifying a list of allowed values + """ + allowedValues: [AllowedValueInput!] + + """ + The optional input for specifying if one or multiple values can be applied. + Default is one value (single cardinality) + """ + cardinality: PropertyCardinality + + """ + The list of entity types that this property can be applied to. + For example: ["urn:li:entityType:datahub.dataset"] + """ + entityTypes: [String!]! +} + +""" +Input for specifying specific entity types as values +""" +input TypeQualifierInput { + """ + The list of allowed entity types as urns (ie. ["urn:li:entityType:datahub.corpuser"]) + """ + allowedTypes: [String!] +} + +""" +An input entry for an allowed value for a structured property +""" +input AllowedValueInput { + """ + The allowed string value if the value is of type string + Either this or numberValue is required. + """ + stringValue: String + + """ + The allowed number value if the value is of type number. + Either this or stringValue is required. + """ + numberValue: Float + + """ + The description of this allowed value + """ + description: String +} diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/structuredproperties/CreateStructuredPropertyResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/structuredproperties/CreateStructuredPropertyResolverTest.java new file mode 100644 index 0000000000000..72cdb78542e41 --- /dev/null +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/structuredproperties/CreateStructuredPropertyResolverTest.java @@ -0,0 +1,126 @@ +package com.linkedin.datahub.graphql.resolvers.structuredproperties; + +import static com.linkedin.datahub.graphql.TestUtils.getMockAllowContext; +import static com.linkedin.datahub.graphql.TestUtils.getMockDenyContext; +import static org.mockito.ArgumentMatchers.any; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertThrows; + +import com.linkedin.common.urn.UrnUtils; +import com.linkedin.datahub.graphql.QueryContext; +import com.linkedin.datahub.graphql.generated.CreateStructuredPropertyInput; +import com.linkedin.datahub.graphql.generated.StructuredPropertyEntity; +import com.linkedin.entity.EntityResponse; +import com.linkedin.entity.EnvelopedAspectMap; +import com.linkedin.entity.client.EntityClient; +import com.linkedin.metadata.Constants; +import com.linkedin.mxe.MetadataChangeProposal; +import com.linkedin.r2.RemoteInvocationException; +import graphql.schema.DataFetchingEnvironment; +import java.util.ArrayList; +import java.util.concurrent.CompletionException; +import org.mockito.Mockito; +import org.testng.annotations.Test; + +public class CreateStructuredPropertyResolverTest { + private static final String TEST_STRUCTURED_PROPERTY_URN = "urn:li:structuredProperty:1"; + + private static final CreateStructuredPropertyInput TEST_INPUT = + new CreateStructuredPropertyInput( + null, + "io.acryl.test", + "Display Name", + "description", + true, + null, + null, + null, + null, + new ArrayList<>()); + + @Test + public void testGetSuccess() throws Exception { + EntityClient mockEntityClient = initMockEntityClient(true); + CreateStructuredPropertyResolver resolver = + new CreateStructuredPropertyResolver(mockEntityClient); + + // Execute resolver + QueryContext mockContext = getMockAllowContext(); + DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); + Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(TEST_INPUT); + Mockito.when(mockEnv.getContext()).thenReturn(mockContext); + + StructuredPropertyEntity prop = resolver.get(mockEnv).get(); + + assertEquals(prop.getUrn(), TEST_STRUCTURED_PROPERTY_URN); + + // Validate that we called ingest + Mockito.verify(mockEntityClient, Mockito.times(1)) + .ingestProposal(any(), any(MetadataChangeProposal.class), Mockito.eq(false)); + } + + @Test + public void testGetUnauthorized() throws Exception { + EntityClient mockEntityClient = initMockEntityClient(true); + CreateStructuredPropertyResolver resolver = + new CreateStructuredPropertyResolver(mockEntityClient); + + // Execute resolver + QueryContext mockContext = getMockDenyContext(); + DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); + Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(TEST_INPUT); + Mockito.when(mockEnv.getContext()).thenReturn(mockContext); + + assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); + + // Validate that we did NOT call ingest + Mockito.verify(mockEntityClient, Mockito.times(0)) + .ingestProposal(any(), any(MetadataChangeProposal.class), Mockito.eq(false)); + } + + @Test + public void testGetFailure() throws Exception { + EntityClient mockEntityClient = initMockEntityClient(false); + CreateStructuredPropertyResolver resolver = + new CreateStructuredPropertyResolver(mockEntityClient); + + // Execute resolver + QueryContext mockContext = getMockAllowContext(); + DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); + Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(TEST_INPUT); + Mockito.when(mockEnv.getContext()).thenReturn(mockContext); + + assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); + + // Validate that ingest was called, but that caused a failure + Mockito.verify(mockEntityClient, Mockito.times(1)) + .ingestProposal(any(), any(MetadataChangeProposal.class), Mockito.eq(false)); + } + + private EntityClient initMockEntityClient(boolean shouldSucceed) throws Exception { + EntityClient client = Mockito.mock(EntityClient.class); + EntityResponse response = new EntityResponse(); + response.setEntityName(Constants.STRUCTURED_PROPERTY_ENTITY_NAME); + response.setUrn(UrnUtils.getUrn(TEST_STRUCTURED_PROPERTY_URN)); + response.setAspects(new EnvelopedAspectMap()); + if (shouldSucceed) { + Mockito.when( + client.getV2( + any(), + Mockito.eq(Constants.STRUCTURED_PROPERTY_ENTITY_NAME), + any(), + Mockito.eq(null))) + .thenReturn(response); + } else { + Mockito.when( + client.getV2( + any(), + Mockito.eq(Constants.STRUCTURED_PROPERTY_ENTITY_NAME), + any(), + Mockito.eq(null))) + .thenThrow(new RemoteInvocationException()); + } + + return client; + } +} diff --git a/metadata-service/war/src/main/resources/boot/policies.json b/metadata-service/war/src/main/resources/boot/policies.json index 454f0ba7d1163..b565c96930066 100644 --- a/metadata-service/war/src/main/resources/boot/policies.json +++ b/metadata-service/war/src/main/resources/boot/policies.json @@ -34,7 +34,8 @@ "MANAGE_GLOBAL_OWNERSHIP_TYPES", "GET_ANALYTICS_PRIVILEGE", "CREATE_BUSINESS_ATTRIBUTE", - "MANAGE_BUSINESS_ATTRIBUTE" + "MANAGE_BUSINESS_ATTRIBUTE", + "MANAGE_STRUCTURED_PROPERTIES" ], "displayName":"Root User - All Platform Privileges", "description":"Grants all platform privileges to root user.", @@ -179,7 +180,8 @@ "MANAGE_GLOBAL_OWNERSHIP_TYPES", "GET_ANALYTICS_PRIVILEGE", "CREATE_BUSINESS_ATTRIBUTE", - "MANAGE_BUSINESS_ATTRIBUTE" + "MANAGE_BUSINESS_ATTRIBUTE", + "MANAGE_STRUCTURED_PROPERTIES" ], "displayName":"Admins - Platform Policy", "description":"Admins have all platform privileges.", @@ -265,7 +267,8 @@ "MANAGE_GLOBAL_ANNOUNCEMENTS", "MANAGE_GLOSSARIES", "MANAGE_TAGS", - "MANAGE_BUSINESS_ATTRIBUTE" + "MANAGE_BUSINESS_ATTRIBUTE", + "MANAGE_STRUCTURED_PROPERTIES" ], "displayName":"Editors - Platform Policy", "description":"Editors can manage ingestion and view analytics.", diff --git a/metadata-utils/src/main/java/com/linkedin/metadata/authorization/PoliciesConfig.java b/metadata-utils/src/main/java/com/linkedin/metadata/authorization/PoliciesConfig.java index 24fa4ec080cfa..f85da06b84d22 100644 --- a/metadata-utils/src/main/java/com/linkedin/metadata/authorization/PoliciesConfig.java +++ b/metadata-utils/src/main/java/com/linkedin/metadata/authorization/PoliciesConfig.java @@ -151,6 +151,12 @@ public class PoliciesConfig { "Manage Connections", "Manage connections to external DataHub platforms."); + public static final Privilege MANAGE_STRUCTURED_PROPERTIES_PRIVILEGE = + Privilege.of( + "MANAGE_STRUCTURED_PROPERTIES", + "Manage Structured Properties", + "Manage structured properties in your instance."); + public static final List PLATFORM_PRIVILEGES = ImmutableList.of( MANAGE_POLICIES_PRIVILEGE, @@ -175,7 +181,8 @@ public class PoliciesConfig { MANAGE_GLOBAL_OWNERSHIP_TYPES, CREATE_BUSINESS_ATTRIBUTE_PRIVILEGE, MANAGE_BUSINESS_ATTRIBUTE_PRIVILEGE, - MANAGE_CONNECTIONS_PRIVILEGE); + MANAGE_CONNECTIONS_PRIVILEGE, + MANAGE_STRUCTURED_PROPERTIES_PRIVILEGE); // Resource Privileges // From d6a2d340f2688663bbb7257c08973d70b9442806 Mon Sep 17 00:00:00 2001 From: Chris Collins Date: Mon, 1 Jul 2024 18:55:11 -0400 Subject: [PATCH 3/3] add updateStructuredProperty graphql endpoint --- .../datahub/graphql/GmsGraphQLEngine.java | 4 + .../UpdateStructuredPropertyResolver.java | 129 ++++++++++++++++++ .../src/main/resources/properties.graphql | 64 +++++++++ .../UpdateStructuredPropertyResolverTest.java | 123 +++++++++++++++++ 4 files changed, 320 insertions(+) create mode 100644 datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/structuredproperties/UpdateStructuredPropertyResolver.java create mode 100644 datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/structuredproperties/UpdateStructuredPropertyResolverTest.java 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 ccd18aa56623d..06d27e43fe528 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 @@ -287,6 +287,7 @@ import com.linkedin.datahub.graphql.resolvers.step.BatchUpdateStepStatesResolver; import com.linkedin.datahub.graphql.resolvers.structuredproperties.CreateStructuredPropertyResolver; import com.linkedin.datahub.graphql.resolvers.structuredproperties.RemoveStructuredPropertiesResolver; +import com.linkedin.datahub.graphql.resolvers.structuredproperties.UpdateStructuredPropertyResolver; import com.linkedin.datahub.graphql.resolvers.structuredproperties.UpsertStructuredPropertiesResolver; import com.linkedin.datahub.graphql.resolvers.tag.CreateTagResolver; import com.linkedin.datahub.graphql.resolvers.tag.DeleteTagResolver; @@ -1324,6 +1325,9 @@ private void configureMutationResolvers(final RuntimeWiring.Builder builder) { .dataFetcher( "createStructuredProperty", new CreateStructuredPropertyResolver(this.entityClient)) + .dataFetcher( + "updateStructuredProperty", + new UpdateStructuredPropertyResolver(this.entityClient)) .dataFetcher("raiseIncident", new RaiseIncidentResolver(this.entityClient)) .dataFetcher( "updateIncidentStatus", diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/structuredproperties/UpdateStructuredPropertyResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/structuredproperties/UpdateStructuredPropertyResolver.java new file mode 100644 index 0000000000000..2549f303bacd9 --- /dev/null +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/structuredproperties/UpdateStructuredPropertyResolver.java @@ -0,0 +1,129 @@ +package com.linkedin.datahub.graphql.resolvers.structuredproperties; + +import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.bindArgument; +import static com.linkedin.metadata.Constants.STRUCTURED_PROPERTY_ENTITY_NAME; + +import com.linkedin.common.urn.Urn; +import com.linkedin.common.urn.UrnUtils; +import com.linkedin.data.template.SetMode; +import com.linkedin.data.template.StringArray; +import com.linkedin.data.template.StringArrayMap; +import com.linkedin.datahub.graphql.QueryContext; +import com.linkedin.datahub.graphql.authorization.AuthorizationUtils; +import com.linkedin.datahub.graphql.exception.AuthorizationException; +import com.linkedin.datahub.graphql.generated.StructuredPropertyEntity; +import com.linkedin.datahub.graphql.generated.UpdateStructuredPropertyInput; +import com.linkedin.datahub.graphql.types.structuredproperty.StructuredPropertyMapper; +import com.linkedin.entity.EntityResponse; +import com.linkedin.entity.client.EntityClient; +import com.linkedin.metadata.aspect.patch.builder.StructuredPropertyDefinitionPatchBuilder; +import com.linkedin.mxe.MetadataChangeProposal; +import com.linkedin.structured.PrimitivePropertyValue; +import com.linkedin.structured.PropertyCardinality; +import com.linkedin.structured.PropertyValue; +import graphql.schema.DataFetcher; +import graphql.schema.DataFetchingEnvironment; +import java.util.Objects; +import java.util.concurrent.CompletableFuture; +import javax.annotation.Nonnull; + +public class UpdateStructuredPropertyResolver + implements DataFetcher> { + + private final EntityClient _entityClient; + + public UpdateStructuredPropertyResolver(@Nonnull final EntityClient entityClient) { + _entityClient = Objects.requireNonNull(entityClient, "entityClient must not be null"); + } + + @Override + public CompletableFuture get(final DataFetchingEnvironment environment) + throws Exception { + final QueryContext context = environment.getContext(); + + final UpdateStructuredPropertyInput input = + bindArgument(environment.getArgument("input"), UpdateStructuredPropertyInput.class); + + return CompletableFuture.supplyAsync( + () -> { + try { + if (!AuthorizationUtils.canManageStructuredProperties(context)) { + throw new AuthorizationException( + "Unable to update structured property. Please contact your admin."); + } + final Urn propertyUrn = UrnUtils.getUrn(input.getUrn()); + StructuredPropertyDefinitionPatchBuilder builder = + new StructuredPropertyDefinitionPatchBuilder().urn(propertyUrn); + + if (input.getDisplayName() != null) { + builder.setDisplayName(input.getDisplayName()); + } + if (input.getDescription() != null) { + builder.setDescription(input.getDescription()); + } + if (input.getImmutable() != null) { + builder.setImmutable(input.getImmutable()); + } + if (input.getTypeQualifier() != null) { + buildTypeQualifier(input, builder); + } + if (input.getNewAllowedValues() != null) { + buildAllowedValues(input, builder); + } + if (input.getSetCardinalityAsMultiple() != null) { + builder.setCardinality(PropertyCardinality.MULTIPLE); + } + if (input.getNewEntityTypes() != null) { + input.getNewEntityTypes().forEach(builder::addEntityType); + } + + MetadataChangeProposal mcp = builder.build(); + _entityClient.ingestProposal(context.getOperationContext(), mcp, false); + + EntityResponse response = + _entityClient.getV2( + context.getOperationContext(), + STRUCTURED_PROPERTY_ENTITY_NAME, + propertyUrn, + null); + return StructuredPropertyMapper.map(context, response); + } catch (Exception e) { + throw new RuntimeException( + String.format("Failed to perform update against input %s", input), e); + } + }); + } + + private void buildTypeQualifier( + @Nonnull final UpdateStructuredPropertyInput input, + @Nonnull final StructuredPropertyDefinitionPatchBuilder builder) { + if (input.getTypeQualifier().getNewAllowedTypes() != null) { + final StringArrayMap typeQualifier = new StringArrayMap(); + StringArray allowedTypes = new StringArray(); + allowedTypes.addAll(input.getTypeQualifier().getNewAllowedTypes()); + typeQualifier.put("allowedTypes", allowedTypes); + builder.setTypeQualifier(typeQualifier); + } + } + + private void buildAllowedValues( + @Nonnull final UpdateStructuredPropertyInput input, + @Nonnull final StructuredPropertyDefinitionPatchBuilder builder) { + input + .getNewAllowedValues() + .forEach( + allowedValueInput -> { + PropertyValue value = new PropertyValue(); + PrimitivePropertyValue primitiveValue = new PrimitivePropertyValue(); + if (allowedValueInput.getStringValue() != null) { + primitiveValue.setString(allowedValueInput.getStringValue()); + } + if (allowedValueInput.getNumberValue() != null) { + primitiveValue.setDouble(allowedValueInput.getNumberValue().doubleValue()); + } + value.setValue(primitiveValue); + value.setDescription(allowedValueInput.getDescription(), SetMode.IGNORE_NULL); + builder.addAllowedValue(value); + }); + } +} diff --git a/datahub-graphql-core/src/main/resources/properties.graphql b/datahub-graphql-core/src/main/resources/properties.graphql index bebe7fed32e51..dfe8468645681 100644 --- a/datahub-graphql-core/src/main/resources/properties.graphql +++ b/datahub-graphql-core/src/main/resources/properties.graphql @@ -13,6 +13,11 @@ extend type Mutation { Create a new structured property """ createStructuredProperty(input: CreateStructuredPropertyInput!): StructuredPropertyEntity! + + """ + Update an existing structured property + """ + updateStructuredProperty(input: UpdateStructuredPropertyInput!): StructuredPropertyEntity! } """ @@ -383,3 +388,62 @@ input AllowedValueInput { """ description: String } + +""" +Input for updating an existing structured property entity +""" +input UpdateStructuredPropertyInput { + """ + The urn of the structured property being updated + """ + urn: String! + + """ + The optional display name for this property + """ + displayName: String + + """ + The optional description for this property + """ + description: String + + """ + Whether the property will be mutable once it is applied or not. Default is false. + """ + immutable: Boolean + + """ + The optional input for specifying specific entity types as values + """ + typeQualifier: UpdateTypeQualifierInput + + """ + Append to the list of allowed values for this property. + For backwards compatibility, this is append only. + """ + newAllowedValues: [AllowedValueInput!] + + """ + Set to true if you want to change the cardinality of this structured property + to multiple. Cannot change from multiple to single for backwards compatibility reasons. + """ + setCardinalityAsMultiple: Boolean + + """ + Append to the list of entity types that this property can be applied to. + For backwards compatibility, this is append only. + """ + newEntityTypes: [String!] +} + +""" +Input for updating specifying specific entity types as values +""" +input UpdateTypeQualifierInput { + """ + Append to the list of allowed entity types as urns for this property (ie. ["urn:li:entityType:datahub.corpuser"]) + For backwards compatibility, this is append only. + """ + newAllowedTypes: [String!] +} diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/structuredproperties/UpdateStructuredPropertyResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/structuredproperties/UpdateStructuredPropertyResolverTest.java new file mode 100644 index 0000000000000..971a53de9473b --- /dev/null +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/structuredproperties/UpdateStructuredPropertyResolverTest.java @@ -0,0 +1,123 @@ +package com.linkedin.datahub.graphql.resolvers.structuredproperties; + +import static com.linkedin.datahub.graphql.TestUtils.getMockAllowContext; +import static com.linkedin.datahub.graphql.TestUtils.getMockDenyContext; +import static org.mockito.ArgumentMatchers.any; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertThrows; + +import com.linkedin.common.urn.UrnUtils; +import com.linkedin.datahub.graphql.QueryContext; +import com.linkedin.datahub.graphql.generated.StructuredPropertyEntity; +import com.linkedin.datahub.graphql.generated.UpdateStructuredPropertyInput; +import com.linkedin.entity.EntityResponse; +import com.linkedin.entity.EnvelopedAspectMap; +import com.linkedin.entity.client.EntityClient; +import com.linkedin.metadata.Constants; +import com.linkedin.mxe.MetadataChangeProposal; +import com.linkedin.r2.RemoteInvocationException; +import graphql.schema.DataFetchingEnvironment; +import java.util.concurrent.CompletionException; +import org.mockito.Mockito; +import org.testng.annotations.Test; + +public class UpdateStructuredPropertyResolverTest { + private static final String TEST_STRUCTURED_PROPERTY_URN = "urn:li:structuredProperty:1"; + + private static final UpdateStructuredPropertyInput TEST_INPUT = + new UpdateStructuredPropertyInput( + TEST_STRUCTURED_PROPERTY_URN, + "New Display Name", + "new description", + true, + null, + null, + null, + null); + + @Test + public void testGetSuccess() throws Exception { + EntityClient mockEntityClient = initMockEntityClient(true); + UpdateStructuredPropertyResolver resolver = + new UpdateStructuredPropertyResolver(mockEntityClient); + + // Execute resolver + QueryContext mockContext = getMockAllowContext(); + DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); + Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(TEST_INPUT); + Mockito.when(mockEnv.getContext()).thenReturn(mockContext); + + StructuredPropertyEntity prop = resolver.get(mockEnv).get(); + + assertEquals(prop.getUrn(), TEST_STRUCTURED_PROPERTY_URN); + + // Validate that we called ingest + Mockito.verify(mockEntityClient, Mockito.times(1)) + .ingestProposal(any(), any(MetadataChangeProposal.class), Mockito.eq(false)); + } + + @Test + public void testGetUnauthorized() throws Exception { + EntityClient mockEntityClient = initMockEntityClient(true); + UpdateStructuredPropertyResolver resolver = + new UpdateStructuredPropertyResolver(mockEntityClient); + + // Execute resolver + QueryContext mockContext = getMockDenyContext(); + DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); + Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(TEST_INPUT); + Mockito.when(mockEnv.getContext()).thenReturn(mockContext); + + assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); + + // Validate that we did NOT call ingest + Mockito.verify(mockEntityClient, Mockito.times(0)) + .ingestProposal(any(), any(MetadataChangeProposal.class), Mockito.eq(false)); + } + + @Test + public void testGetFailure() throws Exception { + EntityClient mockEntityClient = initMockEntityClient(false); + UpdateStructuredPropertyResolver resolver = + new UpdateStructuredPropertyResolver(mockEntityClient); + + // Execute resolver + QueryContext mockContext = getMockAllowContext(); + DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); + Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(TEST_INPUT); + Mockito.when(mockEnv.getContext()).thenReturn(mockContext); + + assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); + + // Validate that ingest was called, but that caused a failure + Mockito.verify(mockEntityClient, Mockito.times(1)) + .ingestProposal(any(), any(MetadataChangeProposal.class), Mockito.eq(false)); + } + + private EntityClient initMockEntityClient(boolean shouldSucceed) throws Exception { + EntityClient client = Mockito.mock(EntityClient.class); + EntityResponse response = new EntityResponse(); + response.setEntityName(Constants.STRUCTURED_PROPERTY_ENTITY_NAME); + response.setUrn(UrnUtils.getUrn(TEST_STRUCTURED_PROPERTY_URN)); + response.setAspects(new EnvelopedAspectMap()); + if (shouldSucceed) { + Mockito.when( + client.getV2( + any(), + Mockito.eq(Constants.STRUCTURED_PROPERTY_ENTITY_NAME), + any(), + Mockito.eq(null))) + .thenReturn(response); + } else { + Mockito.when( + client.getV2( + any(), + Mockito.eq(Constants.STRUCTURED_PROPERTY_ENTITY_NAME), + any(), + Mockito.eq(null))) + .thenThrow(new RemoteInvocationException()); + } + + return client; + } +}