feat(ui): Batch set & unset Domain for assets via the UI

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

@@ -150,6 +150,7 @@
 import com.linkedin.datahub.graphql.resolvers.mutate.BatchRemoveOwnersResolver;
 import com.linkedin.datahub.graphql.resolvers.mutate.BatchRemoveTagsResolver;
 import com.linkedin.datahub.graphql.resolvers.mutate.BatchRemoveTermsResolver;
+import com.linkedin.datahub.graphql.resolvers.mutate.BatchSetDomainResolver;
 import com.linkedin.datahub.graphql.resolvers.mutate.MutableTypeResolver;
 import com.linkedin.datahub.graphql.resolvers.mutate.RemoveLinkResolver;
 import com.linkedin.datahub.graphql.resolvers.mutate.RemoveOwnerResolver;
@@ -720,6 +721,7 @@ private void configureMutationResolvers(final RuntimeWiring.Builder builder) {
             .dataFetcher("createDomain", new CreateDomainResolver(this.entityClient))
             .dataFetcher("deleteDomain", new DeleteDomainResolver(entityClient))
             .dataFetcher("setDomain", new SetDomainResolver(this.entityClient, this.entityService))
+            .dataFetcher("batchSetDomain", new BatchSetDomainResolver(this.entityService))
             .dataFetcher("updateDeprecation", new UpdateDeprecationResolver(this.entityClient, this.entityService))
             .dataFetcher("unsetDomain", new UnsetDomainResolver(this.entityClient, this.entityService))
             .dataFetcher("createSecret", new CreateSecretResolver(this.entityClient, this.secretService))

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

@@ -0,0 +1,88 @@
+package com.linkedin.datahub.graphql.resolvers.mutate;
+
+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.BatchSetDomainInput;
+import com.linkedin.datahub.graphql.generated.ResourceRefInput;
+import com.linkedin.datahub.graphql.resolvers.mutate.util.DomainUtils;
+import com.linkedin.datahub.graphql.resolvers.mutate.util.LabelUtils;
+import com.linkedin.metadata.entity.EntityService;
+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.Nullable;
+import lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
+
+import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.*;
+
+
+@Slf4j
+@RequiredArgsConstructor
+public class BatchSetDomainResolver implements DataFetcher<CompletableFuture<Boolean>> {
+
+  private final EntityService _entityService;
+
+  @Override
+  public CompletableFuture<Boolean> get(DataFetchingEnvironment environment) throws Exception {
+    final QueryContext context = environment.getContext();
+    final BatchSetDomainInput input = bindArgument(environment.getArgument("input"), BatchSetDomainInput.class);
+    final String maybeDomainUrn = input.getDomainUrn();
+    final List<ResourceRefInput> resources = input.getResources();
+
+    return CompletableFuture.supplyAsync(() -> {
+
+      // First, validate the domain
+      validateDomain(maybeDomainUrn);
+      validateInputResources(resources, context);
+
+      try {
+        // Then execute the bulk add
+        batchSetDomains(maybeDomainUrn, resources, context);
+        return true;
+      } catch (Exception e) {
+        log.error("Failed to perform update against input {}, {}", input.toString(), e.getMessage());
+        throw new RuntimeException(String.format("Failed to perform update against input %s", input.toString()), e);
+      }
+    });
+  }
+
+  private void validateDomain(@Nullable String maybeDomainUrn) {
+    if (maybeDomainUrn != null) {
+      DomainUtils.validateDomain(UrnUtils.getUrn(maybeDomainUrn), _entityService);
+    }
+  }
+
+  private void validateInputResources(List<ResourceRefInput> resources, QueryContext context) {
+    for (ResourceRefInput resource : resources) {
+      validateInputResource(resource, context);
+    }
+  }
+
+  private void validateInputResource(ResourceRefInput resource, QueryContext context) {
+    final Urn resourceUrn = UrnUtils.getUrn(resource.getResourceUrn());
+    if (!DomainUtils.isAuthorizedToUpdateDomainsForEntity(context, resourceUrn)) {
+      throw new AuthorizationException("Unauthorized to perform this action. Please contact your DataHub administrator.");
+    }
+    LabelUtils.validateResource(resourceUrn, resource.getSubResource(), resource.getSubResourceType(), _entityService);
+  }
+
+  private void batchSetDomains(String maybeDomainUrn, List<ResourceRefInput> resources, QueryContext context) {
+    log.debug("Batch adding Domains. domainUrn: {}, resources: {}", maybeDomainUrn, resources);
+    try {
+      DomainUtils.setDomainForResources(maybeDomainUrn == null ? null : UrnUtils.getUrn(maybeDomainUrn),
+          resources,
+          UrnUtils.getUrn(context.getActorUrn()),
+          _entityService);
+    } catch (Exception e) {
+      throw new RuntimeException(String.format("Failed to batch set Domain %s to resources with urns %s!",
+          maybeDomainUrn,
+          resources.stream().map(ResourceRefInput::getResourceUrn).collect(Collectors.toList())),
+          e);
+    }
+  }
+}

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

@@ -2,15 +2,27 @@
 
 import com.google.common.collect.ImmutableList;
 
+import com.linkedin.common.UrnArray;
 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.authorization.ConjunctivePrivilegeGroup;
 import com.linkedin.datahub.graphql.authorization.DisjunctivePrivilegeGroup;
+import com.linkedin.datahub.graphql.generated.ResourceRefInput;
+import com.linkedin.domain.Domains;
+import com.linkedin.metadata.Constants;
 import com.linkedin.metadata.authorization.PoliciesConfig;
+import com.linkedin.metadata.entity.EntityService;
+import com.linkedin.mxe.MetadataChangeProposal;
+import java.util.ArrayList;
+import java.util.List;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 import lombok.extern.slf4j.Slf4j;
 
+import static com.linkedin.datahub.graphql.resolvers.mutate.MutationUtils.*;
+
 
 @Slf4j
 public class DomainUtils {
@@ -33,4 +45,49 @@ public static boolean isAuthorizedToUpdateDomainsForEntity(@Nonnull QueryContext
         entityUrn.toString(),
         orPrivilegeGroups);
   }
+
+  public static void setDomainForResources(
+      @Nullable Urn domainUrn,
+      List<ResourceRefInput> resources,
+      Urn actor,
+      EntityService entityService
+  ) throws Exception {
+    final List<MetadataChangeProposal> changes = new ArrayList<>();
+    for (ResourceRefInput resource : resources) {
+      changes.add(buildSetDomainProposal(domainUrn, resource, actor, entityService));
+    }
+    ingestChangeProposals(changes, entityService, actor);
+  }
+
+  private static MetadataChangeProposal buildSetDomainProposal(
+      @Nullable Urn domainUrn,
+      ResourceRefInput resource,
+      Urn actor,
+      EntityService entityService
+  ) {
+    Domains domains = (Domains) getAspectFromEntity(
+        resource.getResourceUrn(),
+        Constants.DOMAINS_ASPECT_NAME,
+        entityService,
+        new Domains());
+    final UrnArray newDomains = new UrnArray();
+    if (domainUrn != null) {
+      newDomains.add(domainUrn);
+    }
+    domains.setDomains(newDomains);
+    return buildMetadataChangeProposal(UrnUtils.getUrn(resource.getResourceUrn()), Constants.DOMAINS_ASPECT_NAME, domains, actor, entityService);
+  }
+
+  public static void validateDomain(Urn domainUrn, EntityService entityService) {
+    if (!entityService.exists(domainUrn)) {
+      throw new IllegalArgumentException(String.format("Failed to validate Domain with urn %s. Urn does not exist.", domainUrn));
+    }
+  }
+
+  private static void ingestChangeProposals(List<MetadataChangeProposal> changes, EntityService entityService, Urn actor) {
+    // TODO: Replace this with a batch ingest proposals endpoint.
+    for (MetadataChangeProposal change : changes) {
+      entityService.ingestProposal(change, getAuditStamp(actor));
+    }
+  }
 }

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

@@ -249,6 +249,7 @@ public static void validateLabel(Urn labelUrn, String labelEntityType, EntitySer
     }
   }
 
+  // TODO: Move this out into a separate utilities class.
   public static void validateResource(Urn resourceUrn, String subResource, SubResourceType subResourceType, EntityService entityService) {
     if (!entityService.exists(resourceUrn)) {
       throw new IllegalArgumentException(String.format("Failed to update resource with urn %s. Entity does not exist.", resourceUrn));

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

@@ -360,7 +360,6 @@ type Mutation {
     """
     addGroupMembers(input: AddGroupMembersInput!): Boolean
 
-
     """
     Remove members from a group
     """
@@ -390,6 +389,11 @@ type Mutation {
     """
     setDomain(entityUrn: String!, domainUrn: String!): Boolean
 
+    """
+    Set domain for multiple Entities
+    """
+    batchSetDomain(input: BatchSetDomainInput!): Boolean
+
     """
     Sets the Domain for a Dataset, Chart, Dashboard, Data Flow (Pipeline), or Data Job (Task). Returns true if the Domain was successfully removed, or was already removed. Requires the Edit Domains privilege for an asset.
     """
@@ -6636,7 +6640,7 @@ input BatchAddTagsInput {
     """
     The target assets to attach the tags to
     """
-    resources: [ResourceRefInput]!
+    resources: [ResourceRefInput!]!
 }
 
 """
@@ -6909,6 +6913,21 @@ input UpdateDeprecationInput {
     note: String
 }
 
+"""
+Input provided when adding tags to a batch of assets
+"""
+input BatchSetDomainInput {
+    """
+    The primary key of the Domain, or null if the domain will be unset
+    """
+    domainUrn: String
+
+    """
+    The target assets to attach the Domain
+    """
+    resources: [ResourceRefInput!]!
+}
+
 """
 Input provided when creating or updating an Access Policy
 """