datahub-project · jjoyce0510 · Aug 4, 2022 · Aug 2, 2022 · Aug 3, 2022 · Aug 3, 2022
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
@@ -144,8 +144,10 @@
 import com.linkedin.datahub.graphql.resolvers.mutate.AddTagsResolver;
 import com.linkedin.datahub.graphql.resolvers.mutate.AddTermResolver;
 import com.linkedin.datahub.graphql.resolvers.mutate.AddTermsResolver;
+import com.linkedin.datahub.graphql.resolvers.mutate.BatchAddOwnersResolver;
 import com.linkedin.datahub.graphql.resolvers.mutate.BatchAddTagsResolver;
 import com.linkedin.datahub.graphql.resolvers.mutate.BatchAddTermsResolver;
+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.MutableTypeResolver;
@@ -704,7 +706,9 @@ private void configureMutationResolvers(final RuntimeWiring.Builder builder) {
             .dataFetcher("updateDescription", new UpdateDescriptionResolver(entityService))
             .dataFetcher("addOwner", new AddOwnerResolver(entityService))
             .dataFetcher("addOwners", new AddOwnersResolver(entityService))
+            .dataFetcher("batchAddOwners", new BatchAddOwnersResolver(entityService))
             .dataFetcher("removeOwner", new RemoveOwnerResolver(entityService))
+            .dataFetcher("batchRemoveOwners", new BatchRemoveOwnersResolver(entityService))
             .dataFetcher("addLink", new AddLinkResolver(entityService))
             .dataFetcher("removeLink", new RemoveLinkResolver(entityService))
             .dataFetcher("addGroupMembers", new AddGroupMembersResolver(this.groupService))

diff --git a/...ql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/AddOwnerResolver.java b/...ql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/AddOwnerResolver.java
@@ -1,13 +1,16 @@
 package com.linkedin.datahub.graphql.resolvers.mutate;
 
+import com.google.common.collect.ImmutableList;
 import com.linkedin.common.urn.CorpuserUrn;
 
 import com.linkedin.common.urn.Urn;
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.exception.AuthorizationException;
 import com.linkedin.datahub.graphql.generated.AddOwnerInput;
 import com.linkedin.datahub.graphql.generated.OwnerEntityType;
+import com.linkedin.datahub.graphql.generated.OwnerInput;
 import com.linkedin.datahub.graphql.generated.OwnershipType;
+import com.linkedin.datahub.graphql.generated.ResourceRefInput;
 import com.linkedin.datahub.graphql.resolvers.mutate.util.OwnerUtils;
 import com.linkedin.metadata.entity.EntityService;
 import graphql.schema.DataFetcher;
@@ -50,11 +53,9 @@ public CompletableFuture<Boolean> get(DataFetchingEnvironment environment) throw
         log.debug("Adding Owner. input: {}", input.toString());
 
         Urn actor = CorpuserUrn.createFromString(((QueryContext) environment.getContext()).getActorUrn());
-        OwnerUtils.addOwner(
-            ownerUrn,
-            // Assumption Alert: Assumes that GraphQL ownership type === GMS ownership type
-            com.linkedin.common.OwnershipType.valueOf(type.name()),
-            targetUrn,
+        OwnerUtils.addOwnersToResources(
+            ImmutableList.of(new OwnerInput(input.getOwnerUrn(), ownerEntityType, type)),
+            ImmutableList.of(new ResourceRefInput(input.getResourceUrn(), null, null)),
             actor,
             _entityService
         );

diff --git a/...l-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/AddOwnersResolver.java b/...l-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/AddOwnersResolver.java
@@ -1,12 +1,14 @@
 package com.linkedin.datahub.graphql.resolvers.mutate;
 
+import com.google.common.collect.ImmutableList;
 import com.linkedin.common.urn.CorpuserUrn;
 
 import com.linkedin.common.urn.Urn;
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.exception.AuthorizationException;
 import com.linkedin.datahub.graphql.generated.AddOwnersInput;
 import com.linkedin.datahub.graphql.generated.OwnerInput;
+import com.linkedin.datahub.graphql.generated.ResourceRefInput;
 import com.linkedin.datahub.graphql.resolvers.mutate.util.OwnerUtils;
 import com.linkedin.metadata.entity.EntityService;
 import graphql.schema.DataFetcher;
@@ -47,9 +49,9 @@ public CompletableFuture<Boolean> get(DataFetchingEnvironment environment) throw
         log.debug("Adding Owners. input: {}", input.toString());
 
         Urn actor = CorpuserUrn.createFromString(((QueryContext) environment.getContext()).getActorUrn());
-        OwnerUtils.addOwners(
+        OwnerUtils.addOwnersToResources(
             owners,
-            targetUrn,
+            ImmutableList.of(new ResourceRefInput(input.getResourceUrn(), null, null)),
             actor,
             _entityService
         );

diff --git a/...e/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/BatchAddOwnersResolver.java b/...e/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/BatchAddOwnersResolver.java
@@ -0,0 +1,90 @@
+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.BatchAddOwnersInput;
+import com.linkedin.datahub.graphql.generated.OwnerInput;
+import com.linkedin.datahub.graphql.generated.ResourceRefInput;
+import com.linkedin.datahub.graphql.resolvers.mutate.util.LabelUtils;
+import com.linkedin.datahub.graphql.resolvers.mutate.util.OwnerUtils;
+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 lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
+
+import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.*;
+
+
+@Slf4j
+@RequiredArgsConstructor
+public class BatchAddOwnersResolver implements DataFetcher<CompletableFuture<Boolean>> {
+
+  private final EntityService _entityService;
+
+  @Override
+  public CompletableFuture<Boolean> get(DataFetchingEnvironment environment) throws Exception {
+    final BatchAddOwnersInput input = bindArgument(environment.getArgument("input"), BatchAddOwnersInput.class);
+    final List<OwnerInput> owners = input.getOwners();
+    final List<ResourceRefInput> resources = input.getResources();
+    final QueryContext context = environment.getContext();
+
+    return CompletableFuture.supplyAsync(() -> {
+
+      // First, validate the batch
+      validateOwners(owners);
+      validateInputResources(resources, context);
+
+      try {
+        // Then execute the bulk add
+        batchAddOwners(owners, 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 validateOwners(List<OwnerInput> owners) {
+    for (OwnerInput ownerInput : owners) {
+      OwnerUtils.validateOwner(UrnUtils.getUrn(ownerInput.getOwnerUrn()), ownerInput.getOwnerEntityType(), _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 (resource.getSubResource() != null) {
+      throw new IllegalArgumentException("Malformed input provided: owners cannot be applied to subresources.");
+    }
+
+    if (!OwnerUtils.isAuthorizedToUpdateOwners(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 batchAddOwners(List<OwnerInput> owners, List<ResourceRefInput> resources, QueryContext context) {
+    log.debug("Batch adding owners. owners: {}, resources: {}", owners, resources);
+    try {
+      OwnerUtils.addOwnersToResources(owners, resources, UrnUtils.getUrn(context.getActorUrn()), _entityService);
+    } catch (Exception e) {
+      throw new RuntimeException(String.format("Failed to batch add Owners %s to resources with urns %s!",
+          owners,
+          resources.stream().map(ResourceRefInput::getResourceUrn).collect(Collectors.toList())),
+          e);
+    }
+  }
+}
diff --git a/...rc/main/java/com/linkedin/datahub/graphql/resolvers/mutate/BatchRemoveOwnersResolver.java b/...rc/main/java/com/linkedin/datahub/graphql/resolvers/mutate/BatchRemoveOwnersResolver.java
@@ -0,0 +1,83 @@
+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.BatchRemoveOwnersInput;
+import com.linkedin.datahub.graphql.generated.ResourceRefInput;
+import com.linkedin.datahub.graphql.resolvers.mutate.util.LabelUtils;
+import com.linkedin.datahub.graphql.resolvers.mutate.util.OwnerUtils;
+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 lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
+
+import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.*;
+
+
+@Slf4j
+@RequiredArgsConstructor
+public class BatchRemoveOwnersResolver implements DataFetcher<CompletableFuture<Boolean>> {
+
+  private final EntityService _entityService;
+
+  @Override
+  public CompletableFuture<Boolean> get(DataFetchingEnvironment environment) throws Exception {
+    final BatchRemoveOwnersInput input = bindArgument(environment.getArgument("input"), BatchRemoveOwnersInput.class);
+    final List<String> owners = input.getOwnerUrns();
+    final List<ResourceRefInput> resources = input.getResources();
+    final QueryContext context = environment.getContext();
+
+    return CompletableFuture.supplyAsync(() -> {
+
+      // First, validate the batch
+      validateInputResources(resources, context);
+
+      try {
+        // Then execute the bulk remove
+        batchRemoveOwners(owners, 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 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 (resource.getSubResource() != null) {
+      throw new IllegalArgumentException("Malformed input provided: owners cannot be removed from subresources.");
+    }
+
+    if (!OwnerUtils.isAuthorizedToUpdateOwners(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 batchRemoveOwners(List<String> ownerUrns, List<ResourceRefInput> resources, QueryContext context) {
+    log.debug("Batch removing owners. owners: {}, resources: {}", ownerUrns, resources);
+    try {
+      OwnerUtils.removeOwnersFromResources(ownerUrns.stream().map(UrnUtils::getUrn).collect(
+          Collectors.toList()), resources, UrnUtils.getUrn(context.getActorUrn()), _entityService);
+    } catch (Exception e) {
+      throw new RuntimeException(String.format("Failed to batch remove Owners %s to resources with urns %s!",
+          ownerUrns,
+          resources.stream().map(ResourceRefInput::getResourceUrn).collect(Collectors.toList())),
+          e);
+    }
+  }
+}
diff --git a/...core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/RemoveOwnerResolver.java b/...core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/RemoveOwnerResolver.java
@@ -1,10 +1,12 @@
 package com.linkedin.datahub.graphql.resolvers.mutate;
 
+import com.google.common.collect.ImmutableList;
 import com.linkedin.common.urn.CorpuserUrn;
 import com.linkedin.common.urn.Urn;
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.exception.AuthorizationException;
 import com.linkedin.datahub.graphql.generated.RemoveOwnerInput;
+import com.linkedin.datahub.graphql.generated.ResourceRefInput;
 import com.linkedin.datahub.graphql.resolvers.mutate.util.OwnerUtils;
 import com.linkedin.metadata.entity.EntityService;
 import graphql.schema.DataFetcher;
@@ -39,12 +41,10 @@ public CompletableFuture<Boolean> get(DataFetchingEnvironment environment) throw
           _entityService
       );
       try {
-        log.debug("Removing Link input: {}", input);
-
         Urn actor = CorpuserUrn.createFromString(((QueryContext) environment.getContext()).getActorUrn());
-        OwnerUtils.removeOwner(
-            ownerUrn,
-            targetUrn,
+        OwnerUtils.removeOwnersFromResources(
+            ImmutableList.of(ownerUrn),
+            ImmutableList.of(new ResourceRefInput(input.getResourceUrn(), null, null)),
             actor,
             _entityService
         );