Skip to content

Commit

Permalink
add updateStructuredProperty graphql endpoint
Browse files Browse the repository at this point in the history
  • Loading branch information
chriscollins3456 committed Jul 1, 2024
1 parent 7569510 commit 992f442
Show file tree
Hide file tree
Showing 4 changed files with 320 additions and 0 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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",
Expand Down
Original file line number Diff line number Diff line change
@@ -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<CompletableFuture<StructuredPropertyEntity>> {

private final EntityClient _entityClient;

public UpdateStructuredPropertyResolver(@Nonnull final EntityClient entityClient) {
_entityClient = Objects.requireNonNull(entityClient, "entityClient must not be null");
}

@Override
public CompletableFuture<StructuredPropertyEntity> 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);
});
}
}
64 changes: 64 additions & 0 deletions datahub-graphql-core/src/main/resources/properties.graphql
Original file line number Diff line number Diff line change
Expand Up @@ -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!
}

"""
Expand Down Expand Up @@ -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!]
}
Original file line number Diff line number Diff line change
@@ -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;
}
}

0 comments on commit 992f442

Please sign in to comment.