feat(properties) Add upsertStructuredProperties graphql endpoint for assets

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java

@@ -256,6 +256,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.UpsertStructuredPropertiesResolver;
 import com.linkedin.datahub.graphql.resolvers.tag.CreateTagResolver;
 import com.linkedin.datahub.graphql.resolvers.tag.DeleteTagResolver;
 import com.linkedin.datahub.graphql.resolvers.tag.SetTagColorResolver;
@@ -1216,6 +1217,9 @@ private void configureMutationResolvers(final RuntimeWiring.Builder builder) {
                 .dataFetcher(
                     "verifyForm", new VerifyFormResolver(this.formService, this.groupService))
                 .dataFetcher("batchRemoveForm", new BatchRemoveFormResolver(this.formService))
+                .dataFetcher(
+                    "upsertStructuredProperties",
+                    new UpsertStructuredPropertiesResolver(this.entityClient))
                 .dataFetcher("raiseIncident", new RaiseIncidentResolver(this.entityClient))
                 .dataFetcher(
                     "updateIncidentStatus",

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/FormUtils.java

@@ -12,7 +12,6 @@
 import com.linkedin.metadata.query.filter.Criterion;
 import com.linkedin.metadata.query.filter.CriterionArray;
 import com.linkedin.metadata.query.filter.Filter;
-import com.linkedin.structured.PrimitivePropertyValue;
 import com.linkedin.structured.PrimitivePropertyValueArray;
 import java.util.Objects;
 import javax.annotation.Nonnull;
@@ -37,14 +36,7 @@ public static PrimitivePropertyValueArray getStructuredPropertyValuesFromInput(
     input
         .getStructuredPropertyParams()
         .getValues()
-        .forEach(
-            value -> {
-              if (value.getStringValue() != null) {
-                values.add(PrimitivePropertyValue.create(value.getStringValue()));
-              } else if (value.getNumberValue() != null) {
-                values.add(PrimitivePropertyValue.create(value.getNumberValue().doubleValue()));
-              }
-            });
+        .forEach(value -> values.add(StructuredPropertyUtils.mapPropertyValueInput(value)));
 
     return values;
   }

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/StructuredPropertyUtils.java

@@ -0,0 +1,51 @@
+package com.linkedin.datahub.graphql.resolvers.mutate.util;
+
+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.datahub.graphql.QueryContext;
+import com.linkedin.datahub.graphql.authorization.AuthorizationUtils;
+import com.linkedin.datahub.graphql.generated.PropertyValueInput;
+import com.linkedin.metadata.authorization.PoliciesConfig;
+import com.linkedin.structured.PrimitivePropertyValue;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+public class StructuredPropertyUtils {
+  private static final ConjunctivePrivilegeGroup ALL_PRIVILEGES_GROUP =
+      new ConjunctivePrivilegeGroup(
+          ImmutableList.of(PoliciesConfig.EDIT_ENTITY_PRIVILEGE.getType()));
+
+  private StructuredPropertyUtils() {}
+
+  @Nullable
+  public static PrimitivePropertyValue mapPropertyValueInput(
+      @Nonnull final PropertyValueInput valueInput) {
+    if (valueInput.getStringValue() != null) {
+      return PrimitivePropertyValue.create(valueInput.getStringValue());
+    } else if (valueInput.getNumberValue() != null) {
+      return PrimitivePropertyValue.create(valueInput.getNumberValue().doubleValue());
+    }
+    return null;
+  }
+
+  public static boolean isAuthorizedToUpdateProperties(
+      @Nonnull QueryContext context, @Nonnull Urn targetUrn) {
+    // If you either have all entity privileges, or have the specific privileges required, you are
+    // authorized.
+    final DisjunctivePrivilegeGroup orPrivilegeGroups =
+        new DisjunctivePrivilegeGroup(
+            ImmutableList.of(
+                ALL_PRIVILEGES_GROUP,
+                new ConjunctivePrivilegeGroup(
+                    ImmutableList.of(PoliciesConfig.EDIT_ENTITY_PROPERTIES_PRIVILEGE.getType()))));
+
+    return AuthorizationUtils.isAuthorized(
+        context.getAuthorizer(),
+        context.getActorUrn(),
+        targetUrn.getEntityType(),
+        targetUrn.toString(),
+        orPrivilegeGroups);
+  }
+}

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/structuredproperties/UpsertStructuredPropertiesResolver.java

@@ -0,0 +1,171 @@
+package com.linkedin.datahub.graphql.resolvers.structuredproperties;
+
+import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.bindArgument;
+import static com.linkedin.metadata.Constants.STRUCTURED_PROPERTIES_ASPECT_NAME;
+
+import com.datahub.authentication.Authentication;
+import com.linkedin.common.AuditStamp;
+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.datahub.graphql.generated.PropertyValueInput;
+import com.linkedin.datahub.graphql.generated.UpsertStructuredPropertiesInput;
+import com.linkedin.datahub.graphql.resolvers.mutate.util.StructuredPropertyUtils;
+import com.linkedin.datahub.graphql.types.structuredproperty.StructuredPropertiesMapper;
+import com.linkedin.entity.EntityResponse;
+import com.linkedin.entity.client.EntityClient;
+import com.linkedin.metadata.entity.AspectUtils;
+import com.linkedin.metadata.utils.AuditStampUtils;
+import com.linkedin.mxe.MetadataChangeProposal;
+import com.linkedin.structured.PrimitivePropertyValueArray;
+import com.linkedin.structured.StructuredProperties;
+import com.linkedin.structured.StructuredPropertyValueAssignment;
+import com.linkedin.structured.StructuredPropertyValueAssignmentArray;
+import graphql.com.google.common.collect.ImmutableSet;
+import graphql.schema.DataFetcher;
+import graphql.schema.DataFetchingEnvironment;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.CompletableFuture;
+import java.util.stream.Collectors;
+import javax.annotation.Nonnull;
+
+public class UpsertStructuredPropertiesResolver
+    implements DataFetcher<
+        CompletableFuture<com.linkedin.datahub.graphql.generated.StructuredProperties>> {
+
+  private final EntityClient _entityClient;
+
+  public UpsertStructuredPropertiesResolver(@Nonnull final EntityClient entityClient) {
+    _entityClient = Objects.requireNonNull(entityClient, "entityClient must not be null");
+  }
+
+  @Override
+  public CompletableFuture<com.linkedin.datahub.graphql.generated.StructuredProperties> get(
+      final DataFetchingEnvironment environment) throws Exception {
+    final QueryContext context = environment.getContext();
+    final Authentication authentication = context.getAuthentication();
+
+    final UpsertStructuredPropertiesInput input =
+        bindArgument(environment.getArgument("input"), UpsertStructuredPropertiesInput.class);
+    final Urn assetUrn = UrnUtils.getUrn(input.getAssetUrn());
+    Map<String, List<PropertyValueInput>> updateMap = new HashMap<>();
+    // create a map of updates from our input
+    input
+        .getStructuredPropertyInputParams()
+        .forEach(param -> updateMap.put(param.getStructuredPropertyUrn(), param.getValues()));
+
+    return CompletableFuture.supplyAsync(
+        () -> {
+          try {
+            // check authorization first
+            if (!StructuredPropertyUtils.isAuthorizedToUpdateProperties(context, assetUrn)) {
+              throw new AuthorizationException(
+                  String.format(
+                      "Not authorized to update properties on the gives urn %s", assetUrn));
+            }
+
+            final AuditStamp auditStamp =
+                AuditStampUtils.createAuditStamp(authentication.getActor().toUrnStr());
+
+            if (!_entityClient.exists(assetUrn, authentication)) {
+              throw new RuntimeException(
+                  String.format("Asset with provided urn %s does not exist", assetUrn));
+            }
+
+            // get or default the structured properties aspect
+            StructuredProperties structuredProperties =
+                getStructuredProperties(assetUrn, authentication);
+
+            // update the existing properties based on new value
+            StructuredPropertyValueAssignmentArray properties =
+                updateExistingProperties(structuredProperties, updateMap, auditStamp);
+
+            // append any new properties from our input
+            addNewProperties(properties, updateMap, auditStamp);
+
+            structuredProperties.setProperties(properties);
+
+            // ingest change proposal
+            final MetadataChangeProposal structuredPropertiesProposal =
+                AspectUtils.buildMetadataChangeProposal(
+                    assetUrn, STRUCTURED_PROPERTIES_ASPECT_NAME, structuredProperties);
+
+            _entityClient.ingestProposal(structuredPropertiesProposal, authentication, false);
+
+            return StructuredPropertiesMapper.map(structuredProperties);
+          } catch (Exception e) {
+            throw new RuntimeException(
+                String.format("Failed to perform update against input %s", input), e);
+          }
+        });
+  }
+
+  private StructuredProperties getStructuredProperties(Urn assetUrn, Authentication authentication)
+      throws Exception {
+    EntityResponse response =
+        _entityClient.getV2(
+            assetUrn.getEntityType(),
+            assetUrn,
+            ImmutableSet.of(STRUCTURED_PROPERTIES_ASPECT_NAME),
+            authentication);
+    StructuredProperties structuredProperties = new StructuredProperties();
+    structuredProperties.setProperties(new StructuredPropertyValueAssignmentArray());
+    if (response != null && response.getAspects().containsKey(STRUCTURED_PROPERTIES_ASPECT_NAME)) {
+      structuredProperties =
+          new StructuredProperties(
+              response.getAspects().get(STRUCTURED_PROPERTIES_ASPECT_NAME).getValue().data());
+    }
+    return structuredProperties;
+  }
+
+  private StructuredPropertyValueAssignmentArray updateExistingProperties(
+      StructuredProperties structuredProperties,
+      Map<String, List<PropertyValueInput>> updateMap,
+      AuditStamp auditStamp) {
+    return new StructuredPropertyValueAssignmentArray(
+        structuredProperties.getProperties().stream()
+            .map(
+                propAssignment -> {
+                  String propUrnString = propAssignment.getPropertyUrn().toString();
+                  if (updateMap.containsKey(propUrnString)) {
+                    List<PropertyValueInput> valueList = updateMap.get(propUrnString);
+                    PrimitivePropertyValueArray values =
+                        new PrimitivePropertyValueArray(
+                            valueList.stream()
+                                .map(StructuredPropertyUtils::mapPropertyValueInput)
+                                .collect(Collectors.toList()));
+                    propAssignment.setValues(values);
+                    propAssignment.setLastModified(auditStamp);
+                  }
+                  return propAssignment;
+                })
+            .collect(Collectors.toList()));
+  }
+
+  private void addNewProperties(
+      StructuredPropertyValueAssignmentArray properties,
+      Map<String, List<PropertyValueInput>> updateMap,
+      AuditStamp auditStamp) {
+    // first remove existing properties from updateMap so that we append only new properties
+    properties.forEach(prop -> updateMap.remove(prop.getPropertyUrn().toString()));
+
+    updateMap.forEach(
+        (structuredPropUrn, values) -> {
+          StructuredPropertyValueAssignment valueAssignment =
+              new StructuredPropertyValueAssignment();
+          valueAssignment.setPropertyUrn(UrnUtils.getUrn(structuredPropUrn));
+          valueAssignment.setValues(
+              new PrimitivePropertyValueArray(
+                  values.stream()
+                      .map(StructuredPropertyUtils::mapPropertyValueInput)
+                      .collect(Collectors.toList())));
+          valueAssignment.setCreated(auditStamp);
+          valueAssignment.setLastModified(auditStamp);
+          properties.add(valueAssignment);
+        });
+  }
+}

datahub-graphql-core/src/main/resources/properties.graphql

@@ -1,3 +1,10 @@
+extend type Mutation {
+    """
+    Upsert structured properties onto a given asset
+    """
+    upsertStructuredProperties(input: UpsertStructuredPropertiesInput!): StructuredProperties!
+}
+
 """
 A structured property that can be shared between different entities
 """
@@ -157,6 +164,21 @@ type StructuredPropertiesEntry {
     valueEntities: [Entity]
 }
 
+"""
+Input for upserting structured properties on a given asset
+"""
+input UpsertStructuredPropertiesInput {
+    """
+    The urn of the asset that we are updating
+    """
+    assetUrn: String!
+
+    """
+    The list of structured properties you want to upsert on this asset
+    """
+    structuredPropertyInputParams: [StructuredPropertyInputParams!]!
+}
+
 """
 A data type registered in DataHub
 """