From d6a2d340f2688663bbb7257c08973d70b9442806 Mon Sep 17 00:00:00 2001 From: Chris Collins Date: Mon, 1 Jul 2024 18:55:11 -0400 Subject: [PATCH] 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 ccd18aa56623dd..06d27e43fe5281 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 00000000000000..2549f303bacd95 --- /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 bebe7fed32e51d..dfe84686456814 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 00000000000000..971a53de9473b5 --- /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; + } +}