feat(roles): add roles feature to DataHub

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

@@ -147,16 +147,16 @@
 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.MutableTypeBatchResolver;
 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.BatchUpdateDeprecationResolver;
 import com.linkedin.datahub.graphql.resolvers.mutate.BatchSetDomainResolver;
+import com.linkedin.datahub.graphql.resolvers.mutate.BatchUpdateDeprecationResolver;
 import com.linkedin.datahub.graphql.resolvers.mutate.BatchUpdateSoftDeletedResolver;
+import com.linkedin.datahub.graphql.resolvers.mutate.MutableTypeBatchResolver;
 import com.linkedin.datahub.graphql.resolvers.mutate.MutableTypeResolver;
 import com.linkedin.datahub.graphql.resolvers.mutate.RemoveLinkResolver;
 import com.linkedin.datahub.graphql.resolvers.mutate.RemoveOwnerResolver;
@@ -172,6 +172,8 @@
 import com.linkedin.datahub.graphql.resolvers.policy.ListPoliciesResolver;
 import com.linkedin.datahub.graphql.resolvers.policy.UpsertPolicyResolver;
 import com.linkedin.datahub.graphql.resolvers.recommendation.ListRecommendationsResolver;
+import com.linkedin.datahub.graphql.resolvers.role.BatchAssignRoleResolver;
+import com.linkedin.datahub.graphql.resolvers.role.ListRolesResolver;
 import com.linkedin.datahub.graphql.resolvers.search.AutoCompleteForMultipleResolver;
 import com.linkedin.datahub.graphql.resolvers.search.AutoCompleteResolver;
 import com.linkedin.datahub.graphql.resolvers.search.SearchAcrossEntitiesResolver;
@@ -230,6 +232,8 @@
 import com.linkedin.datahub.graphql.types.mlmodel.MLModelType;
 import com.linkedin.datahub.graphql.types.mlmodel.MLPrimaryKeyType;
 import com.linkedin.datahub.graphql.types.notebook.NotebookType;
+import com.linkedin.datahub.graphql.types.policy.DataHubPolicyType;
+import com.linkedin.datahub.graphql.types.role.DataHubRoleType;
 import com.linkedin.datahub.graphql.types.tag.TagType;
 import com.linkedin.datahub.graphql.types.test.TestType;
 import com.linkedin.entity.client.EntityClient;
@@ -335,6 +339,8 @@ public class GmsGraphQLEngine {
     private final DataPlatformInstanceType dataPlatformInstanceType;
     private final AccessTokenMetadataType accessTokenMetadataType;
     private final TestType testType;
+    private final DataHubPolicyType dataHubPolicyType;
+    private final DataHubRoleType dataHubRoleType;
 
     /**
      * Configures the graph objects that can be fetched primary key.
@@ -433,6 +439,8 @@ public GmsGraphQLEngine(
         this.dataPlatformInstanceType = new DataPlatformInstanceType(entityClient);
         this.accessTokenMetadataType = new AccessTokenMetadataType(entityClient);
         this.testType = new TestType(entityClient);
+        this.dataHubPolicyType = new DataHubPolicyType(entityClient);
+        this.dataHubRoleType = new DataHubRoleType(entityClient);
         // Init Lists
         this.entityTypes = ImmutableList.of(
             datasetType,
@@ -458,7 +466,9 @@ public GmsGraphQLEngine(
             versionedDatasetType,
             dataPlatformInstanceType,
             accessTokenMetadataType,
-            testType
+            testType,
+            dataHubPolicyType,
+            dataHubRoleType
         );
         this.loadableTypes = new ArrayList<>(entityTypes);
         this.ownerTypes = ImmutableList.of(corpUserType, corpGroupType);
@@ -516,6 +526,7 @@ public void configureRuntimeWiring(final RuntimeWiring.Builder builder) {
         configureVersionedDatasetResolvers(builder);
         configureAccessAccessTokenMetadataResolvers(builder);
         configureTestResultResolvers(builder);
+        configureRoleResolvers(builder);
     }
 
     public GraphQLEngine.Builder builder() {
@@ -660,6 +671,7 @@ private void configureQueryResolvers(final RuntimeWiring.Builder builder) {
             .dataFetcher("entityExists", new EntityExistsResolver(this.entityService))
             .dataFetcher("getNativeUserInviteToken", new GetNativeUserInviteTokenResolver(this.nativeUserService))
             .dataFetcher("entity", getEntityResolver())
+            .dataFetcher("listRoles", new ListRolesResolver(this.entityClient))
         );
     }
 
@@ -765,6 +777,8 @@ private void configureMutationResolvers(final RuntimeWiring.Builder builder) {
             .dataFetcher("batchUpdateSoftDeleted", new BatchUpdateSoftDeletedResolver(this.entityService))
             .dataFetcher("updateUserSetting", new UpdateUserSettingResolver(this.entityService))
             .dataFetcher("rollbackIngestion", new RollbackIngestionResolver(this.entityClient))
+            .dataFetcher("batchAssignRole", new BatchAssignRoleResolver(this.entityClient))
+
         );
     }
 
@@ -1393,35 +1407,30 @@ private void configureAssertionResolvers(final RuntimeWiring.Builder builder) {
 
     private void configurePolicyResolvers(final RuntimeWiring.Builder builder) {
         // Register resolvers for "resolvedUsers" and "resolvedGroups" field of the Policy type.
-        builder.type("ActorFilter", typeWiring -> typeWiring
-            .dataFetcher("resolvedUsers", new LoadableTypeBatchResolver<>(corpUserType,
-                (env) -> {
-                    final ActorFilter filter = env.getSource();
-                    return filter.getUsers();
-                }
-            ))
-            .dataFetcher("resolvedGroups", new LoadableTypeBatchResolver<>(corpGroupType,
-                    (env) -> {
-                        final ActorFilter filter = env.getSource();
-                        return filter.getGroups();
-                    }
-            ))
-        );
+        builder.type("ActorFilter", typeWiring -> typeWiring.dataFetcher("resolvedUsers",
+            new LoadableTypeBatchResolver<>(corpUserType, (env) -> {
+                final ActorFilter filter = env.getSource();
+                return filter.getUsers();
+            })).dataFetcher("resolvedGroups", new LoadableTypeBatchResolver<>(corpGroupType, (env) -> {
+            final ActorFilter filter = env.getSource();
+            return filter.getGroups();
+        })).dataFetcher("resolvedRoles", new LoadableTypeBatchResolver<>(dataHubRoleType, (env) -> {
+            final ActorFilter filter = env.getSource();
+            return filter.getRoles();
+        })));
+    }
+
+    private void configureRoleResolvers(final RuntimeWiring.Builder builder) {
+        builder.type("DataHubRole",
+            typeWiring -> typeWiring.dataFetcher("relationships", new EntityRelationshipsResultResolver(graphClient)));
     }
 
     private void configureDataProcessInstanceResolvers(final RuntimeWiring.Builder builder) {
-        builder.type("DataProcessInstance", typeWiring -> typeWiring
-            .dataFetcher("relationships", new EntityRelationshipsResultResolver(graphClient))
-            .dataFetcher("lineage", new EntityLineageResultResolver(siblingGraphService))
-            .dataFetcher("state",
-                new TimeSeriesAspectResolver(
-                    this.entityClient,
-                    "dataProcessInstance",
-                    DATA_PROCESS_INSTANCE_RUN_EVENT_ASPECT_NAME,
-                    DataProcessInstanceRunEventMapper::map
-                )
-            )
-        );
+        builder.type("DataProcessInstance",
+            typeWiring -> typeWiring.dataFetcher("relationships", new EntityRelationshipsResultResolver(graphClient))
+                .dataFetcher("lineage", new EntityLineageResultResolver(siblingGraphService))
+                .dataFetcher("state", new TimeSeriesAspectResolver(this.entityClient, "dataProcessInstance",
+                    DATA_PROCESS_INSTANCE_RUN_EVENT_ASPECT_NAME, DataProcessInstanceRunEventMapper::map)));
     }
 
     private void configureTestResultResolvers(final RuntimeWiring.Builder builder) {

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/authorization/AuthorizationUtils.java

@@ -30,6 +30,10 @@ public static boolean canManageUsersAndGroups(@Nonnull QueryContext context) {
     return isAuthorized(context, Optional.empty(), PoliciesConfig.MANAGE_USERS_AND_GROUPS_PRIVILEGE);
   }
 
+  public static boolean canManagePolicies(@Nonnull QueryContext context) {
+    return isAuthorized(context, Optional.empty(), PoliciesConfig.MANAGE_POLICIES_PRIVILEGE);
+  }
+
   public static boolean canGeneratePersonalAccessToken(@Nonnull QueryContext context) {
     return isAuthorized(context, Optional.empty(), PoliciesConfig.GENERATE_PERSONAL_ACCESS_TOKENS_PRIVILEGE);
   }

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/policy/mappers/PolicyInfoPolicyMapper.java

@@ -59,6 +59,9 @@ private ActorFilter mapActors(final DataHubActorFilter actorFilter) {
     if (actorFilter.hasUsers()) {
       result.setUsers(actorFilter.getUsers().stream().map(Urn::toString).collect(Collectors.toList()));
     }
+    if (actorFilter.hasRoles()) {
+      result.setRoles(actorFilter.getRoles().stream().map(Urn::toString).collect(Collectors.toList()));
+    }
     return result;
   }
 

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/role/BatchAssignRoleResolver.java

@@ -0,0 +1,89 @@
+package com.linkedin.datahub.graphql.resolvers.role;
+
+import com.datahub.authentication.Authentication;
+import com.linkedin.common.UrnArray;
+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.BatchAssignRoleInput;
+import com.linkedin.entity.client.EntityClient;
+import com.linkedin.events.metadata.ChangeType;
+import com.linkedin.identity.RoleMembership;
+import com.linkedin.metadata.utils.GenericRecordUtils;
+import com.linkedin.mxe.MetadataChangeProposal;
+import com.linkedin.r2.RemoteInvocationException;
+import graphql.schema.DataFetcher;
+import graphql.schema.DataFetchingEnvironment;
+import java.net.URISyntaxException;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
+
+import static com.linkedin.datahub.graphql.authorization.AuthorizationUtils.*;
+import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.*;
+import static com.linkedin.metadata.Constants.*;
+
+
+@Slf4j
+@RequiredArgsConstructor
+public class BatchAssignRoleResolver implements DataFetcher<CompletableFuture<Boolean>> {
+  private final EntityClient _entityClient;
+
+  @Override
+  public CompletableFuture<Boolean> get(DataFetchingEnvironment environment) throws Exception {
+    final QueryContext context = environment.getContext();
+    if (!canManagePolicies(context)) {
+      throw new AuthorizationException(
+          "Unauthorized to assign roles. Please contact your DataHub administrator if this needs corrective action.");
+    }
+
+    final BatchAssignRoleInput input = bindArgument(environment.getArgument("input"), BatchAssignRoleInput.class);
+    final String roleUrnStr = input.getRoleUrn();
+    final List<String> actors = input.getActors();
+    final Authentication authentication = context.getAuthentication();
+
+    return CompletableFuture.supplyAsync(() -> {
+      try {
+        Urn roleUrn = Urn.createFromString(roleUrnStr);
+        if (!_entityClient.exists(roleUrn, authentication)) {
+          throw new RuntimeException(String.format("Role %s does not exist", roleUrnStr));
+        }
+
+        actors.forEach(actor -> {
+          try {
+            assignRoleToActor(actor, roleUrn, authentication);
+          } catch (Exception e) {
+            log.warn(
+                String.format("Failed to assign role %s to actor %s. Skipping actor assignment", roleUrnStr, actor), e);
+          }
+        });
+        return true;
+      } catch (Exception e) {
+        throw new RuntimeException(String.format("Failed to perform update against input %s", input), e);
+      }
+    });
+  }
+
+  private void assignRoleToActor(String actor, Urn roleUrn, Authentication authentication)
+      throws URISyntaxException, RemoteInvocationException {
+    Urn actorUrn = Urn.createFromString(actor);
+    if (!_entityClient.exists(actorUrn, authentication)) {
+      log.warn(String.format("Failed to assign role %s to actor %s, actor does not exist. Skipping actor assignment",
+          roleUrn.toString(), actor));
+      return;
+    }
+
+    RoleMembership roleMembership = new RoleMembership();
+    roleMembership.setRoles(new UrnArray(roleUrn));
+
+    // Finally, create the MetadataChangeProposal.
+    final MetadataChangeProposal proposal = new MetadataChangeProposal();
+    proposal.setEntityUrn(actorUrn);
+    proposal.setEntityType(CORP_USER_ENTITY_NAME);
+    proposal.setAspectName(ROLE_MEMBERSHIP_ASPECT_NAME);
+    proposal.setAspect(GenericRecordUtils.serializeAspect(roleMembership));
+    proposal.setChangeType(ChangeType.UPSERT);
+    _entityClient.ingestProposal(proposal, authentication);
+  }
+}

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/role/ListRolesResolver.java

@@ -0,0 +1,80 @@
+package com.linkedin.datahub.graphql.resolvers.role;
+
+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.DataHubRole;
+import com.linkedin.datahub.graphql.generated.ListRolesInput;
+import com.linkedin.datahub.graphql.generated.ListRolesResult;
+import com.linkedin.datahub.graphql.types.role.mappers.DataHubRoleMapper;
+import com.linkedin.entity.EntityResponse;
+import com.linkedin.entity.client.EntityClient;
+import com.linkedin.metadata.search.SearchEntity;
+import com.linkedin.metadata.search.SearchResult;
+import graphql.schema.DataFetcher;
+import graphql.schema.DataFetchingEnvironment;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.stream.Collectors;
+import lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
+
+import static com.linkedin.datahub.graphql.authorization.AuthorizationUtils.*;
+import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.*;
+import static com.linkedin.metadata.Constants.*;
+
+
+@Slf4j
+@RequiredArgsConstructor
+public class ListRolesResolver implements DataFetcher<CompletableFuture<ListRolesResult>> {
+  private static final Integer DEFAULT_START = 0;
+  private static final Integer DEFAULT_COUNT = 20;
+  private static final String DEFAULT_QUERY = "";
+
+  private final EntityClient _entityClient;
+
+  @Override
+  public CompletableFuture<ListRolesResult> get(final DataFetchingEnvironment environment) throws Exception {
+    final QueryContext context = environment.getContext();
+    if (!canManagePolicies(context)) {
+      throw new AuthorizationException(
+          "Unauthorized to view roles. Please contact your DataHub administrator if this needs corrective action.");
+    }
+
+    final ListRolesInput input = bindArgument(environment.getArgument("input"), ListRolesInput.class);
+    final Integer start = input.getStart() == null ? DEFAULT_START : input.getStart();
+    final Integer count = input.getCount() == null ? DEFAULT_COUNT : input.getCount();
+    final String query = input.getQuery() == null ? DEFAULT_QUERY : input.getQuery();
+
+    return CompletableFuture.supplyAsync(() -> {
+      try {
+        // First, get all role Urns.
+        final SearchResult gmsResult =
+            _entityClient.search(DATAHUB_ROLE_ENTITY_NAME, query, Collections.emptyMap(), start, count,
+                context.getAuthentication());
+
+        // Then, get and hydrate all users.
+        final Map<Urn, EntityResponse> entities = _entityClient.batchGetV2(DATAHUB_ROLE_ENTITY_NAME,
+            new HashSet<>(gmsResult.getEntities().stream().map(SearchEntity::getEntity).collect(Collectors.toList())),
+            null, context.getAuthentication());
+
+        final ListRolesResult result = new ListRolesResult();
+        result.setStart(gmsResult.getFrom());
+        result.setCount(gmsResult.getPageSize());
+        result.setTotal(gmsResult.getNumEntities());
+        result.setRoles(mapEntitiesToRoles(entities.values()));
+        return result;
+      } catch (Exception e) {
+        throw new RuntimeException("Failed to list roles", e);
+      }
+    });
+  }
+
+  private List<DataHubRole> mapEntitiesToRoles(final Collection<EntityResponse> entities) {
+    return entities.stream().map(DataHubRoleMapper::map).collect(Collectors.toList());
+  }
+}

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/UrnToEntityMapper.java

@@ -8,6 +8,8 @@
 import com.linkedin.datahub.graphql.generated.CorpUser;
 import com.linkedin.datahub.graphql.generated.Dashboard;
 import com.linkedin.datahub.graphql.generated.DataFlow;
+import com.linkedin.datahub.graphql.generated.DataHubPolicy;
+import com.linkedin.datahub.graphql.generated.DataHubRole;
 import com.linkedin.datahub.graphql.generated.DataJob;
 import com.linkedin.datahub.graphql.generated.DataPlatform;
 import com.linkedin.datahub.graphql.generated.DataPlatformInstance;
@@ -28,6 +30,8 @@
 import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
 import javax.annotation.Nonnull;
 
+import static com.linkedin.metadata.Constants.*;
+
 
 public class UrnToEntityMapper implements ModelMapper<com.linkedin.common.urn.Urn, Entity>  {
   public static final UrnToEntityMapper INSTANCE = new UrnToEntityMapper();
@@ -149,6 +153,16 @@ public Entity apply(Urn input) {
       ((Assertion) partialEntity).setUrn(input.toString());
       ((Assertion) partialEntity).setType(EntityType.TEST);
     }
+    if (input.getEntityType().equals(DATAHUB_ROLE_ENTITY_NAME)) {
+      partialEntity = new DataHubRole();
+      ((DataHubRole) partialEntity).setUrn(input.toString());
+      ((DataHubRole) partialEntity).setType(EntityType.DATAHUB_ROLE);
+    }
+    if (input.getEntityType().equals(POLICY_ENTITY_NAME)) {
+      partialEntity = new DataHubPolicy();
+      ((DataHubPolicy) partialEntity).setUrn(input.toString());
+      ((DataHubPolicy) partialEntity).setType(EntityType.DATAHUB_POLICY);
+    }
     return partialEntity;
   }
 }