fix(ownership): Corrects validation of ownership type and makes it co…

…nsistent across graphQL calls

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

@@ -32,30 +32,33 @@ public class AddOwnerResolver implements DataFetcher<CompletableFuture<Boolean>>
   @Override
   public CompletableFuture<Boolean> get(DataFetchingEnvironment environment) throws Exception {
     final AddOwnerInput input = bindArgument(environment.getArgument("input"), AddOwnerInput.class);
-
     Urn ownerUrn = Urn.createFromString(input.getOwnerUrn());
-    OwnerEntityType ownerEntityType = input.getOwnerEntityType();
-    OwnershipType type = input.getType() == null ? OwnershipType.NONE : input.getType();
-    String ownershipUrn = input.getOwnershipTypeUrn() == null ? mapOwnershipTypeToEntity(type.name()) : input.getOwnershipTypeUrn();
     Urn targetUrn = Urn.createFromString(input.getResourceUrn());
+    OwnerInput.Builder ownerInputBuilder = OwnerInput.builder();
+    ownerInputBuilder.setOwnerUrn(input.getOwnerUrn());
+    ownerInputBuilder.setOwnerEntityType(input.getOwnerEntityType());
+    if(input.getType() != null) {
+      ownerInputBuilder.setType(input.getType());
+    }
+    if(input.getOwnershipTypeUrn() != null) {
+      ownerInputBuilder.setOwnershipTypeUrn(input.getOwnershipTypeUrn());
+    }
 
+    OwnerInput ownerInput = ownerInputBuilder.build();
     if (!OwnerUtils.isAuthorizedToUpdateOwners(environment.getContext(), targetUrn)) {
       throw new AuthorizationException("Unauthorized to perform this action. Please contact your DataHub administrator.");
     }
 
     return CompletableFuture.supplyAsync(() -> {
-      OwnerUtils.validateAddInput(
-          ownerUrn, input.getOwnershipTypeUrn(), ownerEntityType,
-          targetUrn,
-          _entityService
-      );
+      OwnerUtils.validateAddOwnerInput(ownerInput, ownerUrn, _entityService);
+
       try {
 
         log.debug("Adding Owner. input: {}", input);
 
         Urn actor = CorpuserUrn.createFromString(((QueryContext) environment.getContext()).getActorUrn());
         OwnerUtils.addOwnersToResources(
-            ImmutableList.of(new OwnerInput(input.getOwnerUrn(), ownerEntityType, type, ownershipUrn)),
+            ImmutableList.of(ownerInput),
             ImmutableList.of(new ResourceRefInput(input.getResourceUrn(), null, null)),
             actor,
             _entityService

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

@@ -39,7 +39,7 @@ public CompletableFuture<Boolean> get(DataFetchingEnvironment environment) throw
         throw new AuthorizationException("Unauthorized to perform this action. Please contact your DataHub administrator.");
       }
 
-      OwnerUtils.validateAddInput(
+      OwnerUtils.validateAddOwnerInput(
           owners,
           targetUrn,
           _entityService

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

@@ -53,8 +53,7 @@ public CompletableFuture<Boolean> get(DataFetchingEnvironment environment) throw
 
   private void validateOwners(List<OwnerInput> owners) {
     for (OwnerInput ownerInput : owners) {
-      OwnerUtils.validateOwner(UrnUtils.getUrn(ownerInput.getOwnerUrn()), ownerInput.getOwnerEntityType(),
-          UrnUtils.getUrn(ownerInput.getOwnershipTypeUrn()), _entityService);
+      OwnerUtils.validateOwner(ownerInput, _entityService);
     }
   }
 

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

@@ -50,7 +50,7 @@ public static void addOwnersToResources(
   ) {
     final List<MetadataChangeProposal> changes = new ArrayList<>();
     for (ResourceRefInput resource : resources) {
-      changes.add(buildAddOwnersProposal(owners, UrnUtils.getUrn(resource.getResourceUrn()), actor, entityService));
+      changes.add(buildAddOwnersProposal(owners, UrnUtils.getUrn(resource.getResourceUrn()), entityService));
     }
     EntityUtils.ingestChangeProposals(changes, entityService, actor, false);
   }
@@ -69,13 +69,13 @@ public static void removeOwnersFromResources(
   }
 
 
-  private static MetadataChangeProposal buildAddOwnersProposal(List<OwnerInput> owners, Urn resourceUrn, Urn actor, EntityService entityService) {
+  static MetadataChangeProposal buildAddOwnersProposal(List<OwnerInput> owners, Urn resourceUrn, EntityService entityService) {
     Ownership ownershipAspect = (Ownership) EntityUtils.getAspectFromEntity(
         resourceUrn.toString(),
         Constants.OWNERSHIP_ASPECT_NAME, entityService,
         new Ownership());
     for (OwnerInput input : owners) {
-      addOwner(ownershipAspect, UrnUtils.getUrn(input.getOwnerUrn()), input.getType(), UrnUtils.getUrn(input.getOwnershipTypeUrn()));
+      addOwner(ownershipAspect, UrnUtils.getUrn(input.getOwnerUrn()), ownershipType, UrnUtils.getUrn(input.getOwnershipTypeUrn()));
     }
     return buildMetadataChangeProposalWithUrn(resourceUrn, Constants.OWNERSHIP_ASPECT_NAME, ownershipAspect);
   }
@@ -181,63 +181,43 @@ public static boolean isAuthorizedToUpdateOwners(@Nonnull QueryContext context,
         orPrivilegeGroups);
   }
 
-  public static Boolean validateAddInput(
+  public static Boolean validateAddOwnerInput(
       List<OwnerInput> owners,
       Urn resourceUrn,
       EntityService entityService
   ) {
     for (OwnerInput owner : owners) {
-      boolean result = validateAddInput(
-          UrnUtils.getUrn(owner.getOwnerUrn()),
-          owner.getOwnershipTypeUrn(),
-          owner.getOwnerEntityType(),
-          resourceUrn,
-          entityService);
+      boolean result = validateAddOwnerInput(owner, resourceUrn, entityService);
       if (!result) {
         return false;
       }
     }
     return true;
   }
 
-  public static Boolean validateAddInput(
-      Urn ownerUrn,
-      String ownershipEntityUrn,
-      OwnerEntityType ownerEntityType,
+  public static Boolean validateAddOwnerInput(
+      OwnerInput owner,
       Urn resourceUrn,
       EntityService entityService
   ) {
 
-    if (OwnerEntityType.CORP_GROUP.equals(ownerEntityType) && !Constants.CORP_GROUP_ENTITY_NAME.equals(ownerUrn.getEntityType())) {
-      throw new IllegalArgumentException(String.format("Failed to change ownership for resource %s. Expected a corp group urn.", resourceUrn));
-    }
-
-    if (OwnerEntityType.CORP_USER.equals(ownerEntityType) && !Constants.CORP_USER_ENTITY_NAME.equals(ownerUrn.getEntityType())) {
-      throw new IllegalArgumentException(String.format("Failed to change ownership for resource %s. Expected a corp user urn.", resourceUrn));
-    }
-
     if (!entityService.exists(resourceUrn)) {
       throw new IllegalArgumentException(String.format("Failed to change ownership for resource %s. Resource does not exist.", resourceUrn));
     }
 
-    if (!entityService.exists(ownerUrn)) {
-      throw new IllegalArgumentException(String.format("Failed to change ownership for resource %s. Owner %s does not exist.", resourceUrn, ownerUrn));
-    }
-
-    if (ownershipEntityUrn != null && !entityService.exists(UrnUtils.getUrn(ownershipEntityUrn))) {
-      throw new IllegalArgumentException(String.format("Failed to change ownership type for resource %s. Ownership Type "
-          + "%s does not exist.", resourceUrn, ownershipEntityUrn));
-    }
+    validateOwner(owner, entityService);
 
     return true;
   }
 
   public static void validateOwner(
-      Urn ownerUrn,
-      OwnerEntityType ownerEntityType,
-      Urn ownershipEntityUrn,
+      OwnerInput owner,
       EntityService entityService
   ) {
+
+    OwnerEntityType ownerEntityType = owner.getOwnerEntityType();
+    Urn ownerUrn = UrnUtils.getUrn(owner.getOwnerUrn());
+
     if (OwnerEntityType.CORP_GROUP.equals(ownerEntityType) && !Constants.CORP_GROUP_ENTITY_NAME.equals(ownerUrn.getEntityType())) {
       throw new IllegalArgumentException(
           String.format("Failed to change ownership for resource(s). Expected a corp group urn, found %s", ownerUrn));
@@ -252,9 +232,14 @@ public static void validateOwner(
       throw new IllegalArgumentException(String.format("Failed to change ownership for resource(s). Owner with urn %s does not exist.", ownerUrn));
     }
 
-    if (!entityService.exists(ownershipEntityUrn)) {
-      throw new IllegalArgumentException(String.format("Failed to change ownership for resource(s). Ownership type with "
-          + "urn %s does not exist.", ownershipEntityUrn));
+    if (owner.getOwnershipTypeUrn() != null && !entityService.exists(UrnUtils.getUrn(owner.getOwnershipTypeUrn()))) {
+      throw new IllegalArgumentException(String.format("Failed to change ownership for resource(s). Custom Ownership type with "
+          + "urn %s does not exist.", owner.getOwnershipTypeUrn()));
+    }
+
+    if (owner.getType() == null && owner.getOwnershipTypeUrn() == null) {
+      throw new IllegalArgumentException("Failed to change ownership for resource(s). Expected either "
+          + "type or ownershipTypeUrn to be specified.");
     }
   }
 
@@ -269,11 +254,11 @@ public static Boolean validateRemoveInput(
   }
 
   public static void addCreatorAsOwner(
-    QueryContext context,
-    String urn,
-    OwnerEntityType ownerEntityType,
-    OwnershipType ownershipType,
-    EntityService entityService) {
+      QueryContext context,
+      String urn,
+      OwnerEntityType ownerEntityType,
+      OwnershipType ownershipType,
+      EntityService entityService) {
     try {
       Urn actorUrn = CorpuserUrn.createFromString(context.getActorUrn());
       String ownershipTypeUrn = mapOwnershipTypeToEntity(ownershipType.name());