feat(graphql) Add new Revokable Token API

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

@@ -1,7 +1,7 @@
 package com.linkedin.datahub.graphql;
 
 import com.datahub.authentication.AuthenticationConfiguration;
-import com.datahub.authentication.token.TokenService;
+import com.datahub.authentication.token.StatefulTokenService;
 import com.datahub.authorization.AuthorizationConfiguration;
 import com.google.common.collect.ImmutableList;
 import com.linkedin.common.VersionedUrn;
@@ -12,6 +12,8 @@
 import com.linkedin.datahub.graphql.analytics.resolver.GetMetadataAnalyticsResolver;
 import com.linkedin.datahub.graphql.analytics.resolver.IsAnalyticsEnabledResolver;
 import com.linkedin.datahub.graphql.analytics.service.AnalyticsService;
+import com.linkedin.datahub.graphql.generated.AccessToken;
+import com.linkedin.datahub.graphql.generated.AccessTokenMetadata;
 import com.linkedin.datahub.graphql.generated.ActorFilter;
 import com.linkedin.datahub.graphql.generated.AggregationMetadata;
 import com.linkedin.datahub.graphql.generated.Assertion;
@@ -37,6 +39,7 @@
 import com.linkedin.datahub.graphql.generated.ForeignKeyConstraint;
 import com.linkedin.datahub.graphql.generated.InstitutionalMemoryMetadata;
 import com.linkedin.datahub.graphql.generated.LineageRelationship;
+import com.linkedin.datahub.graphql.generated.ListAccessTokenResult;
 import com.linkedin.datahub.graphql.generated.ListDomainsResult;
 import com.linkedin.datahub.graphql.generated.MLFeature;
 import com.linkedin.datahub.graphql.generated.MLFeatureProperties;
@@ -59,12 +62,15 @@
 import com.linkedin.datahub.graphql.resolvers.assertion.AssertionRunEventResolver;
 import com.linkedin.datahub.graphql.resolvers.assertion.DeleteAssertionResolver;
 import com.linkedin.datahub.graphql.resolvers.assertion.EntityAssertionsResolver;
+import com.linkedin.datahub.graphql.resolvers.auth.CreateAccessTokenResolver;
 import com.linkedin.datahub.graphql.resolvers.auth.GetAccessTokenResolver;
+import com.linkedin.datahub.graphql.resolvers.auth.ListAccessTokensResolver;
+import com.linkedin.datahub.graphql.resolvers.auth.RevokeAccessTokenResolver;
 import com.linkedin.datahub.graphql.resolvers.browse.BrowsePathsResolver;
 import com.linkedin.datahub.graphql.resolvers.browse.BrowseResolver;
 import com.linkedin.datahub.graphql.resolvers.config.AppConfigResolver;
-import com.linkedin.datahub.graphql.resolvers.container.ParentContainersResolver;
 import com.linkedin.datahub.graphql.resolvers.container.ContainerEntitiesResolver;
+import com.linkedin.datahub.graphql.resolvers.container.ParentContainersResolver;
 import com.linkedin.datahub.graphql.resolvers.dataset.DatasetHealthResolver;
 import com.linkedin.datahub.graphql.resolvers.deprecation.UpdateDeprecationResolver;
 import com.linkedin.datahub.graphql.resolvers.domain.CreateDomainResolver;
@@ -137,24 +143,25 @@
 import com.linkedin.datahub.graphql.resolvers.user.RemoveUserResolver;
 import com.linkedin.datahub.graphql.resolvers.user.UpdateUserStatusResolver;
 import com.linkedin.datahub.graphql.types.BrowsableEntityType;
-import com.linkedin.datahub.graphql.types.dataplatforminstance.DataPlatformInstanceType;
-import com.linkedin.datahub.graphql.types.dataprocessinst.mappers.DataProcessInstanceRunEventMapper;
 import com.linkedin.datahub.graphql.types.EntityType;
 import com.linkedin.datahub.graphql.types.LoadableType;
 import com.linkedin.datahub.graphql.types.SearchableEntityType;
 import com.linkedin.datahub.graphql.types.aspect.AspectType;
 import com.linkedin.datahub.graphql.types.assertion.AssertionType;
+import com.linkedin.datahub.graphql.types.auth.AccessTokenMetadataType;
 import com.linkedin.datahub.graphql.types.chart.ChartType;
 import com.linkedin.datahub.graphql.types.common.mappers.OperationMapper;
 import com.linkedin.datahub.graphql.types.container.ContainerType;
 import com.linkedin.datahub.graphql.types.corpgroup.CorpGroupType;
 import com.linkedin.datahub.graphql.types.corpuser.CorpUserType;
 import com.linkedin.datahub.graphql.types.dashboard.DashboardType;
-import com.linkedin.datahub.graphql.types.dataset.VersionedDatasetType;
 import com.linkedin.datahub.graphql.types.dataflow.DataFlowType;
 import com.linkedin.datahub.graphql.types.datajob.DataJobType;
 import com.linkedin.datahub.graphql.types.dataplatform.DataPlatformType;
+import com.linkedin.datahub.graphql.types.dataplatforminstance.DataPlatformInstanceType;
+import com.linkedin.datahub.graphql.types.dataprocessinst.mappers.DataProcessInstanceRunEventMapper;
 import com.linkedin.datahub.graphql.types.dataset.DatasetType;
+import com.linkedin.datahub.graphql.types.dataset.VersionedDatasetType;
 import com.linkedin.datahub.graphql.types.dataset.mappers.DatasetProfileMapper;
 import com.linkedin.datahub.graphql.types.domain.DomainType;
 import com.linkedin.datahub.graphql.types.glossary.GlossaryTermType;
@@ -183,6 +190,12 @@
 import graphql.schema.DataFetchingEnvironment;
 import graphql.schema.StaticDataFetcher;
 import graphql.schema.idl.RuntimeWiring;
+import lombok.extern.slf4j.Slf4j;
+import org.apache.commons.io.IOUtils;
+import org.dataloader.BatchLoaderContextProvider;
+import org.dataloader.DataLoader;
+import org.dataloader.DataLoaderOptions;
+
 import java.io.IOException;
 import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
@@ -195,15 +208,10 @@
 import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
-import lombok.extern.slf4j.Slf4j;
-import org.apache.commons.io.IOUtils;
-import org.dataloader.BatchLoaderContextProvider;
-import org.dataloader.DataLoader;
-import org.dataloader.DataLoaderOptions;
 
 import static com.linkedin.datahub.graphql.Constants.*;
-import static com.linkedin.metadata.Constants.*;
-import static graphql.Scalars.*;
+import static com.linkedin.metadata.Constants.DATA_PROCESS_INSTANCE_RUN_EVENT_ASPECT_NAME;
+import static graphql.Scalars.GraphQLLong;
 
 
 /**
@@ -220,7 +228,7 @@ public class GmsGraphQLEngine {
     private final AnalyticsService analyticsService;
     private final RecommendationsService recommendationsService;
     private final EntityRegistry entityRegistry;
-    private final TokenService tokenService;
+    private final StatefulTokenService statefulTokenService;
     private final SecretService secretService;
     private final GitVersion gitVersion;
     private final boolean supportsImpactAnalysis;
@@ -256,7 +264,7 @@ public class GmsGraphQLEngine {
     private final AssertionType assertionType;
     private final VersionedDatasetType versionedDatasetType;
     private final DataPlatformInstanceType dataPlatformInstanceType;
-
+    private final AccessTokenMetadataType accessTokenMetadataType;
 
     /**
      * Configures the graph objects that can be fetched primary key.
@@ -290,7 +298,7 @@ public GmsGraphQLEngine(
         final AnalyticsService analyticsService,
         final EntityService entityService,
         final RecommendationsService recommendationsService,
-        final TokenService tokenService,
+        final StatefulTokenService statefulTokenService,
         final TimeseriesAspectService timeseriesAspectService,
         final EntityRegistry entityRegistry,
         final SecretService secretService,
@@ -311,7 +319,7 @@ public GmsGraphQLEngine(
         this.analyticsService = analyticsService;
         this.entityService = entityService;
         this.recommendationsService = recommendationsService;
-        this.tokenService = tokenService;
+        this.statefulTokenService = statefulTokenService;
         this.secretService = secretService;
         this.entityRegistry = entityRegistry;
         this.gitVersion = gitVersion;
@@ -348,7 +356,7 @@ public GmsGraphQLEngine(
         this.assertionType = new AssertionType(entityClient);
         this.versionedDatasetType = new VersionedDatasetType(entityClient);
         this.dataPlatformInstanceType = new DataPlatformInstanceType(entityClient);
-
+        this.accessTokenMetadataType = new AccessTokenMetadataType(entityClient);
         // Init Lists
         this.entityTypes = ImmutableList.of(
             datasetType,
@@ -371,7 +379,8 @@ public GmsGraphQLEngine(
             domainType,
             assertionType,
             versionedDatasetType,
-            dataPlatformInstanceType
+            dataPlatformInstanceType,
+            accessTokenMetadataType
         );
         this.loadableTypes = new ArrayList<>(entityTypes);
         this.ownerTypes = ImmutableList.of(corpUserType, corpGroupType);
@@ -425,6 +434,7 @@ public void configureRuntimeWiring(final RuntimeWiring.Builder builder) {
         configurePolicyResolvers(builder);
         configureDataProcessInstanceResolvers(builder);
         configureVersionedDatasetResolvers(builder);
+        configureAccessAccessTokenMetadataResolvers(builder);
     }
 
     public GraphQLEngine.Builder builder() {
@@ -549,7 +559,8 @@ private void configureQueryResolvers(final RuntimeWiring.Builder builder) {
             .dataFetcher("listGroups", new ListGroupsResolver(this.entityClient))
             .dataFetcher("listRecommendations", new ListRecommendationsResolver(recommendationsService))
             .dataFetcher("getEntityCounts", new EntityCountsResolver(this.entityClient))
-            .dataFetcher("getAccessToken", new GetAccessTokenResolver(tokenService))
+            .dataFetcher("getAccessToken", new GetAccessTokenResolver(statefulTokenService))
+            .dataFetcher("listAccessTokens", new ListAccessTokensResolver(this.entityClient))
             .dataFetcher("container", getResolver(containerType))
             .dataFetcher("listDomains", new ListDomainsResolver(this.entityClient))
             .dataFetcher("listSecrets", new ListSecretsResolver(this.entityClient))
@@ -613,6 +624,8 @@ private void configureMutationResolvers(final RuntimeWiring.Builder builder) {
             .dataFetcher("unsetDomain", new UnsetDomainResolver(this.entityClient, this.entityService))
             .dataFetcher("createSecret", new CreateSecretResolver(this.entityClient, this.secretService))
             .dataFetcher("deleteSecret", new DeleteSecretResolver(this.entityClient))
+            .dataFetcher("createAccessToken", new CreateAccessTokenResolver(this.statefulTokenService))
+            .dataFetcher("revokeAccessToken", new RevokeAccessTokenResolver(this.entityClient, this.statefulTokenService))
             .dataFetcher("createIngestionSource", new UpsertIngestionSourceResolver(this.entityClient))
             .dataFetcher("updateIngestionSource", new UpsertIngestionSourceResolver(this.entityClient))
             .dataFetcher("deleteIngestionSource", new DeleteIngestionSourceResolver(this.entityClient))
@@ -763,6 +776,22 @@ private void configureVersionedDatasetResolvers(final RuntimeWiring.Builder buil
 
     }
 
+    /**
+     * Configures resolvers responsible for resolving the {@link com.linkedin.datahub.graphql.generated.AccessTokenMetadata} type.
+     */
+    private void configureAccessAccessTokenMetadataResolvers(final RuntimeWiring.Builder builder) {
+        builder.type("AccessToken", typeWiring -> typeWiring
+            .dataFetcher("metadata", new LoadableTypeResolver<>(accessTokenMetadataType,
+                (env) -> ((AccessToken) env.getSource()).getMetadata().getUrn()))
+        );
+        builder.type("ListAccessTokenResult", typeWiring -> typeWiring
+            .dataFetcher("tokens", new LoadableTypeBatchResolver<>(accessTokenMetadataType,
+                (env) -> ((ListAccessTokenResult) env.getSource()).getTokens().stream()
+                    .map(AccessTokenMetadata::getUrn)
+                    .collect(Collectors.toList())))
+        );
+    }
+
     private void configureGlossaryTermResolvers(final RuntimeWiring.Builder builder) {
         builder.type("GlossaryTerm", typeWiring -> typeWiring
             .dataFetcher("schemaMetadata", new AspectResolver())

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

@@ -21,6 +21,10 @@ public static boolean canGeneratePersonalAccessToken(@Nonnull QueryContext conte
     return isAuthorized(context, Optional.empty(), PoliciesConfig.GENERATE_PERSONAL_ACCESS_TOKENS_PRIVILEGE);
   }
 
+  public static boolean canManageTokens(@Nonnull QueryContext context) {
+    return isAuthorized(context, Optional.empty(), PoliciesConfig.MANAGE_ACCESS_TOKENS);
+  }
+
   public static boolean canManageDomains(@Nonnull QueryContext context) {
     return isAuthorized(context, Optional.empty(), PoliciesConfig.MANAGE_DOMAINS_PRIVILEGE);
   }

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

@@ -60,6 +60,7 @@ public CompletableFuture<AuthenticatedUser> get(DataFetchingEnvironment environm
         platformPrivileges.setManageDomains(canManageDomains(context));
         platformPrivileges.setManageIngestion(canManageIngestion(context));
         platformPrivileges.setManageSecrets(canManageSecrets(context));
+        platformPrivileges.setManageTokens(canManageTokens(context));
 
         // Construct and return authenticated user object.
         final AuthenticatedUser authUser = new AuthenticatedUser();
@@ -109,7 +110,14 @@ private boolean canManageDomains(final QueryContext context) {
   }
 
   /**
-   * Returns true if the the provided actor is authorized for a particular privilege, false otherwise.
+   * Returns true if the authenticated user has privileges to manage access tokens
+   */
+  private boolean canManageTokens(final QueryContext context) {
+    return isAuthorized(context.getAuthorizer(), context.getActorUrn(), PoliciesConfig.MANAGE_ACCESS_TOKENS);
+  }
+
+  /**
+   * Returns true if the provided actor is authorized for a particular privilege, false otherwise.
    */
   private boolean isAuthorized(final Authorizer authorizer, String actor, PoliciesConfig.Privilege privilege) {
     final AuthorizationRequest request = new AuthorizationRequest(actor, privilege.getType(), Optional.empty());

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/auth/CreateAccessTokenResolver.java

@@ -0,0 +1,105 @@
+package com.linkedin.datahub.graphql.resolvers.auth;
+
+import com.datahub.authentication.Actor;
+import com.datahub.authentication.ActorType;
+import com.datahub.authentication.token.StatefulTokenService;
+import com.datahub.authentication.token.TokenType;
+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.exception.AuthorizationException;
+import com.linkedin.datahub.graphql.generated.AccessToken;
+import com.linkedin.datahub.graphql.generated.AccessTokenType;
+import com.linkedin.datahub.graphql.generated.CreateAccessTokenInput;
+import com.linkedin.datahub.graphql.generated.EntityType;
+import com.linkedin.datahub.graphql.generated.AccessTokenMetadata;
+import com.linkedin.metadata.Constants;
+import graphql.schema.DataFetcher;
+import graphql.schema.DataFetchingEnvironment;
+import java.util.Date;
+import java.util.concurrent.CompletableFuture;
+import lombok.extern.slf4j.Slf4j;
+
+import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.*;
+
+
+/**
+ * Resolver for creating personal & service principal v2-type (stateful) access tokens.
+ */
+@Slf4j
+public class CreateAccessTokenResolver implements DataFetcher<CompletableFuture<AccessToken>> {
+
+  private final StatefulTokenService _statefulTokenService;
+
+  public CreateAccessTokenResolver(final StatefulTokenService statefulTokenService) {
+    _statefulTokenService = statefulTokenService;
+  }
+
+  @Override
+  public CompletableFuture<AccessToken> get(final DataFetchingEnvironment environment) throws Exception {
+    return CompletableFuture.supplyAsync(() -> {
+      final QueryContext context = environment.getContext();
+      final CreateAccessTokenInput input = bindArgument(environment.getArgument("input"), CreateAccessTokenInput.class);
+
+      log.info("User {} requesting new access token for user {} ", context.getActorUrn(), input.getActorUrn());
+
+      if (isAuthorizedToGenerateToken(context, input)) {
+        final TokenType type = TokenType.valueOf(
+            input.getType().toString()); // warn: if we are out of sync with AccessTokenType there are problems.
+        final String actorUrn = input.getActorUrn();
+        final Date date = new Date();
+        final long createdAtInMs = date.getTime();
+        final long expiresInMs = AccessTokenUtil.mapDurationToMs(input.getDuration());
+
+        final String tokenName = input.getName();
+        final String tokenDescription = input.getDescription();
+
+        final String accessToken =
+            _statefulTokenService.generateAccessToken(type, createActor(input.getType(), actorUrn), expiresInMs,
+                createdAtInMs, tokenName, tokenDescription, context.getActorUrn());
+        log.info("Generated access token for {} of type {} with duration {}", input.getActorUrn(), input.getType(),
+            input.getDuration());
+        try {
+          final String tokenHash = _statefulTokenService.hash(accessToken);
+
+          final AccessToken result = new AccessToken();
+          result.setAccessToken(accessToken);
+          final AccessTokenMetadata metadata = new AccessTokenMetadata();
+          metadata.setUrn(Urn.createFromTuple(Constants.ACCESS_TOKEN_ENTITY_NAME, tokenHash).toString());
+          metadata.setType(EntityType.ACCESS_TOKEN);
+          result.setMetadata(metadata);
+
+          return result;
+        } catch (Exception e) {
+          throw new RuntimeException(String.format("Failed to create new access token with name %s", input.getName()),
+              e);
+        }
+      }
+      throw new AuthorizationException(
+          "Unauthorized to perform this action. Please contact your DataHub administrator.");
+    });
+  }
+
+  private boolean isAuthorizedToGenerateToken(final QueryContext context, final CreateAccessTokenInput input) {
+    if (AccessTokenType.PERSONAL.equals(input.getType())) {
+      return isAuthorizedToGeneratePersonalAccessToken(context, input);
+    }
+    throw new UnsupportedOperationException(String.format("Unsupported AccessTokenType %s provided", input.getType()));
+  }
+
+  private boolean isAuthorizedToGeneratePersonalAccessToken(final QueryContext context,
+      final CreateAccessTokenInput input) {
+    return AuthorizationUtils.canManageTokens(context)
+        || input.getActorUrn().equals(context.getActorUrn()) && AuthorizationUtils.canGeneratePersonalAccessToken(
+        context);
+  }
+
+  private Actor createActor(AccessTokenType tokenType, String actorUrn) {
+    if (AccessTokenType.PERSONAL.equals(tokenType)) {
+      // If we are generating a personal access token, then the actor will be of "USER" type.
+      return new Actor(ActorType.USER, UrnUtils.getUrn(actorUrn).getId());
+    }
+    throw new IllegalArgumentException(String.format("Unsupported token type %s provided", tokenType));
+  }
+}