Address review comments

datahub-project · Sep 22, 2022 · d12e162 · d12e162
1 parent 12aa5b1
commit d12e162
Show file tree

Hide file tree

Showing 21 changed files with 204 additions and 428 deletions.
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
@@ -674,7 +674,7 @@ private void configureQueryResolvers(final RuntimeWiring.Builder builder) {
             .dataFetcher("getRootGlossaryTerms", new GetRootGlossaryTermsResolver(this.entityClient))
             .dataFetcher("getRootGlossaryNodes", new GetRootGlossaryNodesResolver(this.entityClient))
             .dataFetcher("entityExists", new EntityExistsResolver(this.entityService))
-            .dataFetcher("getNativeUserInviteToken", new GetNativeUserInviteTokenResolver(this.nativeUserService))
+            .dataFetcher("getNativeUserInviteToken", new GetNativeUserInviteTokenResolver(this.inviteTokenService))
             .dataFetcher("entity", getEntityResolver())
             .dataFetcher("entities", getEntitiesResolver())
             .dataFetcher("listRoles", new ListRolesResolver(this.entityClient))
@@ -794,7 +794,8 @@ private void configureMutationResolvers(final RuntimeWiring.Builder builder) {
             .dataFetcher("updateName", new UpdateNameResolver(entityService))
             .dataFetcher("addRelatedTerms", new AddRelatedTermsResolver(this.entityService))
             .dataFetcher("removeRelatedTerms", new RemoveRelatedTermsResolver(this.entityService))
-            .dataFetcher("createNativeUserInviteToken", new CreateNativeUserInviteTokenResolver(this.nativeUserService))
+            .dataFetcher("createNativeUserInviteToken",
+                new CreateNativeUserInviteTokenResolver(this.inviteTokenService))
             .dataFetcher("createNativeUserResetToken", new CreateNativeUserResetTokenResolver(this.nativeUserService))
             .dataFetcher("batchUpdateSoftDeleted", new BatchUpdateSoftDeletedResolver(this.entityService))
             .dataFetcher("updateUserSetting", new UpdateUserSettingResolver(this.entityService))

diff --git a/...ql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/role/AcceptRoleResolver.java b/...ql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/role/AcceptRoleResolver.java
@@ -8,7 +8,6 @@
 import com.linkedin.datahub.graphql.generated.AcceptRoleInput;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
-import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
@@ -38,14 +37,14 @@ public CompletableFuture<Boolean> get(DataFetchingEnvironment environment) throw
           throw new RuntimeException(String.format("Invite token %s is invalid", inviteTokenStr));
         }
 
-        Optional<Urn> roleUrnOptional = _inviteTokenService.getRoleUrnFromInviteToken(inviteTokenUrn, authentication);
-        if (roleUrnOptional.isPresent()) {
-          _roleService.assignRoleToActor(authentication.getActor().toUrnStr(), roleUrnOptional.get(), authentication);
+        Urn roleUrn = _inviteTokenService.getInviteTokenRole(inviteTokenUrn, authentication);
+        if (roleUrn != null) {
+          _roleService.assignRoleToActor(authentication.getActor().toUrnStr(), roleUrn, authentication);
         }
 
         return true;
       } catch (Exception e) {
-        throw new RuntimeException(String.format("Failed to perform update against input %s", input), e);
+        throw new RuntimeException(String.format("Failed to accept role using invite token %s", inviteTokenStr), e);
       }
     });
   }

diff --git a/.../src/main/java/com/linkedin/datahub/graphql/resolvers/role/CreateInviteTokenResolver.java b/.../src/main/java/com/linkedin/datahub/graphql/resolvers/role/CreateInviteTokenResolver.java
@@ -2,14 +2,12 @@
 
 import com.datahub.authentication.Authentication;
 import com.datahub.authentication.invite.InviteTokenService;
-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.CreateInviteTokenInput;
 import com.linkedin.datahub.graphql.generated.InviteToken;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
-import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
@@ -37,11 +35,9 @@ public CompletableFuture<InviteToken> get(final DataFetchingEnvironment environm
 
     return CompletableFuture.supplyAsync(() -> {
       try {
-        Optional<Urn> optionalRoleUrn =
-            roleUrnStr == null ? Optional.empty() : Optional.of(Urn.createFromString(roleUrnStr));
-        return new InviteToken(_inviteTokenService.getInviteToken(optionalRoleUrn, true, authentication));
+        return new InviteToken(_inviteTokenService.getInviteToken(roleUrnStr, true, authentication));
       } catch (Exception e) {
-        throw new RuntimeException(String.format("Failed to perform update against input %s", input), e);
+        throw new RuntimeException(String.format("Failed to create invite token for role %s", roleUrnStr), e);
       }
     });
   }

diff --git a/...ore/src/main/java/com/linkedin/datahub/graphql/resolvers/role/GetInviteTokenResolver.java b/...ore/src/main/java/com/linkedin/datahub/graphql/resolvers/role/GetInviteTokenResolver.java
@@ -2,14 +2,12 @@
 
 import com.datahub.authentication.Authentication;
 import com.datahub.authentication.invite.InviteTokenService;
-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.GetInviteTokenInput;
 import com.linkedin.datahub.graphql.generated.InviteToken;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
-import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
@@ -37,11 +35,9 @@ public CompletableFuture<InviteToken> get(final DataFetchingEnvironment environm
 
     return CompletableFuture.supplyAsync(() -> {
       try {
-        Optional<Urn> optionalRoleUrn =
-            roleUrnStr == null ? Optional.empty() : Optional.of(Urn.createFromString(roleUrnStr));
-        return new InviteToken(_inviteTokenService.getInviteToken(optionalRoleUrn, false, authentication));
+        return new InviteToken(_inviteTokenService.getInviteToken(roleUrnStr, false, authentication));
       } catch (Exception e) {
-        throw new RuntimeException(String.format("Failed to perform update against input %s", input), e);
+        throw new RuntimeException(String.format("Failed to get invite token for role %s", roleUrnStr), e);
       }
     });
   }

diff --git a/...java/com/linkedin/datahub/graphql/resolvers/user/CreateNativeUserInviteTokenResolver.java b/...java/com/linkedin/datahub/graphql/resolvers/user/CreateNativeUserInviteTokenResolver.java
@@ -1,40 +1,42 @@
 package com.linkedin.datahub.graphql.resolvers.user;
 
-import com.datahub.authentication.user.NativeUserService;
+import com.datahub.authentication.Authentication;
+import com.datahub.authentication.invite.InviteTokenService;
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.exception.AuthorizationException;
 import com.linkedin.datahub.graphql.generated.InviteToken;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
 import java.util.concurrent.CompletableFuture;
+import lombok.RequiredArgsConstructor;
 
 import static com.linkedin.datahub.graphql.authorization.AuthorizationUtils.*;
 
 /**
  * Resolver responsible for creating an invite token that Admins can share with prospective users to create native
  * user accounts.
  */
+@RequiredArgsConstructor
+@Deprecated
 public class CreateNativeUserInviteTokenResolver implements DataFetcher<CompletableFuture<InviteToken>> {
-  private final NativeUserService _nativeUserService;
-
-  public CreateNativeUserInviteTokenResolver(final NativeUserService nativeUserService) {
-    _nativeUserService = nativeUserService;
-  }
+  private final InviteTokenService _inviteTokenService;
 
   @Override
   public CompletableFuture<InviteToken> get(final DataFetchingEnvironment environment) throws Exception {
     final QueryContext context = environment.getContext();
+    if (!canManagePolicies(context)) {
+      throw new AuthorizationException(
+          "Unauthorized to create invite tokens. Please contact your DataHub administrator if this needs corrective action.");
+    }
+
+    final Authentication authentication = context.getAuthentication();
+
     return CompletableFuture.supplyAsync(() -> {
-      if (!canManageUserCredentials(context)) {
-        throw new AuthorizationException(
-            "Unauthorized to perform this action. Please contact your DataHub administrator.");
-      }
 
       try {
-        String inviteToken = _nativeUserService.generateNativeUserInviteToken(context.getAuthentication());
-        return new InviteToken(inviteToken);
+        return new InviteToken(_inviteTokenService.getInviteToken(null, true, authentication));
       } catch (Exception e) {
-        throw new RuntimeException("Failed to generate new invite token");
+        throw new RuntimeException("Failed to create new invite token");
       }
     });
   }

diff --git a/...in/java/com/linkedin/datahub/graphql/resolvers/user/GetNativeUserInviteTokenResolver.java b/...in/java/com/linkedin/datahub/graphql/resolvers/user/GetNativeUserInviteTokenResolver.java
@@ -1,41 +1,41 @@
 package com.linkedin.datahub.graphql.resolvers.user;
 
-import com.datahub.authentication.user.NativeUserService;
+import com.datahub.authentication.Authentication;
+import com.datahub.authentication.invite.InviteTokenService;
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.exception.AuthorizationException;
 import com.linkedin.datahub.graphql.generated.InviteToken;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
 import java.util.concurrent.CompletableFuture;
+import lombok.RequiredArgsConstructor;
 
 import static com.linkedin.datahub.graphql.authorization.AuthorizationUtils.*;
 
 /**
  * Resolver responsible for getting an existing invite token that Admins can share with prospective users to create
  * native user accounts. If the invite token does not already exist, this resolver will create a new one.
  */
+@RequiredArgsConstructor
+@Deprecated
 public class GetNativeUserInviteTokenResolver implements DataFetcher<CompletableFuture<InviteToken>> {
-  private final NativeUserService _nativeUserService;
-
-  public GetNativeUserInviteTokenResolver(final NativeUserService nativeUserService) {
-    _nativeUserService = nativeUserService;
-  }
+  private final InviteTokenService _inviteTokenService;
 
   @Override
   public CompletableFuture<InviteToken> get(final DataFetchingEnvironment environment) throws Exception {
     final QueryContext context = environment.getContext();
+    if (!canManagePolicies(context)) {
+      throw new AuthorizationException(
+          "Unauthorized to get invite tokens. Please contact your DataHub administrator if this needs corrective action.");
+    }
 
-    return CompletableFuture.supplyAsync(() -> {
-      if (!canManageUserCredentials(context)) {
-        throw new AuthorizationException(
-            "Unauthorized to perform this action. Please contact your DataHub administrator.");
-      }
+    final Authentication authentication = context.getAuthentication();
 
+    return CompletableFuture.supplyAsync(() -> {
       try {
-        String inviteToken = _nativeUserService.getNativeUserInviteToken(context.getAuthentication());
-        return new InviteToken(inviteToken);
+        return new InviteToken(_inviteTokenService.getInviteToken(null, false, authentication));
       } catch (Exception e) {
-        throw new RuntimeException("Failed to generate new invite token");
+        throw new RuntimeException("Failed to get invite token");
       }
     });
   }

diff --git a/datahub-graphql-core/src/main/resources/entity.graphql b/datahub-graphql-core/src/main/resources/entity.graphql
@@ -167,7 +167,7 @@ type Query {
     """
     Gets the current invite token. If the optional regenerate param is set to true, generate a new invite token.
     """
-    getNativeUserInviteToken: InviteToken
+    getNativeUserInviteToken: InviteToken @deprecated(reason: "Use getInviteToken instead")
 
     """
     Gets an entity based on its urn
@@ -497,7 +497,7 @@ type Mutation {
     """
     Generates an invite token that can be shared with prospective users to create their accounts.
     """
-    createNativeUserInviteToken: InviteToken
+    createNativeUserInviteToken: InviteToken @deprecated(reason: "Use createInviteToken instead")
 
     """
     Generates a token that can be shared with existing native users to reset their credentials.

diff --git a/...ore/src/test/java/com/linkedin/datahub/graphql/resolvers/role/AcceptRoleResolverTest.java b/...ore/src/test/java/com/linkedin/datahub/graphql/resolvers/role/AcceptRoleResolverTest.java
@@ -8,7 +8,6 @@
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.generated.AcceptRoleInput;
 import graphql.schema.DataFetchingEnvironment;
-import java.util.Optional;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
 
@@ -71,8 +70,7 @@ public void testNoRoleUrn() throws Exception {
     when(mockContext.getAuthentication()).thenReturn(_authentication);
     when(_inviteTokenService.getInviteTokenUrn(eq(INVITE_TOKEN_STRING))).thenReturn(inviteTokenUrn);
     when(_inviteTokenService.isInviteTokenValid(eq(inviteTokenUrn), eq(_authentication))).thenReturn(true);
-    when(_inviteTokenService.getRoleUrnFromInviteToken(eq(inviteTokenUrn), eq(_authentication))).thenReturn(
-        Optional.empty());
+    when(_inviteTokenService.getInviteTokenRole(eq(inviteTokenUrn), eq(_authentication))).thenReturn(null);
 
     AcceptRoleInput input = new AcceptRoleInput();
     input.setInviteToken(INVITE_TOKEN_STRING);
@@ -89,8 +87,7 @@ public void testAssignRolePasses() throws Exception {
     when(mockContext.getAuthentication()).thenReturn(_authentication);
     when(_inviteTokenService.getInviteTokenUrn(eq(INVITE_TOKEN_STRING))).thenReturn(inviteTokenUrn);
     when(_inviteTokenService.isInviteTokenValid(eq(inviteTokenUrn), eq(_authentication))).thenReturn(true);
-    when(_inviteTokenService.getRoleUrnFromInviteToken(eq(inviteTokenUrn), eq(_authentication))).thenReturn(
-        Optional.of(roleUrn));
+    when(_inviteTokenService.getInviteTokenRole(eq(inviteTokenUrn), eq(_authentication))).thenReturn(roleUrn);
     Actor actor = mock(Actor.class);
     when(_authentication.getActor()).thenReturn(actor);
     when(actor.toUrnStr()).thenReturn(ACTOR_URN_STRING);

diff --git a/.../com/linkedin/datahub/graphql/resolvers/user/CreateNativeUserInviteTokenResolverTest.java b/.../com/linkedin/datahub/graphql/resolvers/user/CreateNativeUserInviteTokenResolverTest.java
@@ -1,7 +1,7 @@
 package com.linkedin.datahub.graphql.resolvers.user;
 
 import com.datahub.authentication.Authentication;
-import com.datahub.authentication.user.NativeUserService;
+import com.datahub.authentication.invite.InviteTokenService;
 import com.linkedin.datahub.graphql.QueryContext;
 import graphql.schema.DataFetchingEnvironment;
 import org.testng.annotations.BeforeMethod;
@@ -16,22 +16,22 @@ public class CreateNativeUserInviteTokenResolverTest {
 
   private static final String INVITE_TOKEN = "inviteToken";
 
-  private NativeUserService _nativeUserService;
+  private InviteTokenService _inviteTokenService;
   private CreateNativeUserInviteTokenResolver _resolver;
   private DataFetchingEnvironment _dataFetchingEnvironment;
   private Authentication _authentication;
 
   @BeforeMethod
   public void setupTest() {
-    _nativeUserService = mock(NativeUserService.class);
+    _inviteTokenService = mock(InviteTokenService.class);
     _dataFetchingEnvironment = mock(DataFetchingEnvironment.class);
     _authentication = mock(Authentication.class);
 
-    _resolver = new CreateNativeUserInviteTokenResolver(_nativeUserService);
+    _resolver = new CreateNativeUserInviteTokenResolver(_inviteTokenService);
   }
 
   @Test
-  public void testFailsCannotManageUserCredentials() {
+  public void testFailsDenyContext() {
     QueryContext mockContext = getMockDenyContext();
     when(_dataFetchingEnvironment.getContext()).thenReturn(mockContext);
 
@@ -43,7 +43,7 @@ public void testPasses() throws Exception {
     QueryContext mockContext = getMockAllowContext();
     when(_dataFetchingEnvironment.getContext()).thenReturn(mockContext);
     when(mockContext.getAuthentication()).thenReturn(_authentication);
-    when(_nativeUserService.generateNativeUserInviteToken(any())).thenReturn(INVITE_TOKEN);
+    when(_inviteTokenService.getInviteToken(any(), eq(true), eq(_authentication))).thenReturn(INVITE_TOKEN);
 
     assertEquals(INVITE_TOKEN, _resolver.get(_dataFetchingEnvironment).join().getInviteToken());
   }

diff --git a/...ava/com/linkedin/datahub/graphql/resolvers/user/GetNativeUserInviteTokenResolverTest.java b/...ava/com/linkedin/datahub/graphql/resolvers/user/GetNativeUserInviteTokenResolverTest.java
@@ -1,7 +1,7 @@
 package com.linkedin.datahub.graphql.resolvers.user;
 
 import com.datahub.authentication.Authentication;
-import com.datahub.authentication.user.NativeUserService;
+import com.datahub.authentication.invite.InviteTokenService;
 import com.linkedin.datahub.graphql.QueryContext;
 import graphql.schema.DataFetchingEnvironment;
 import org.testng.annotations.BeforeMethod;
@@ -16,22 +16,22 @@ public class GetNativeUserInviteTokenResolverTest {
 
   private static final String INVITE_TOKEN = "inviteToken";
 
-  private NativeUserService _nativeUserService;
+  private InviteTokenService _inviteTokenService;
   private GetNativeUserInviteTokenResolver _resolver;
   private DataFetchingEnvironment _dataFetchingEnvironment;
   private Authentication _authentication;
 
   @BeforeMethod
   public void setupTest() {
-    _nativeUserService = mock(NativeUserService.class);
+    _inviteTokenService = mock(InviteTokenService.class);
     _dataFetchingEnvironment = mock(DataFetchingEnvironment.class);
     _authentication = mock(Authentication.class);
 
-    _resolver = new GetNativeUserInviteTokenResolver(_nativeUserService);
+    _resolver = new GetNativeUserInviteTokenResolver(_inviteTokenService);
   }
 
   @Test
-  public void testFailsCannotManageUserCredentials() {
+  public void testFailsDenyContext() {
     QueryContext mockContext = getMockDenyContext();
     when(_dataFetchingEnvironment.getContext()).thenReturn(mockContext);
 
@@ -43,7 +43,7 @@ public void testPasses() throws Exception {
     QueryContext mockContext = getMockAllowContext();
     when(_dataFetchingEnvironment.getContext()).thenReturn(mockContext);
     when(mockContext.getAuthentication()).thenReturn(_authentication);
-    when(_nativeUserService.getNativeUserInviteToken(any())).thenReturn(INVITE_TOKEN);
+    when(_inviteTokenService.getInviteToken(any(), eq(false), eq(_authentication))).thenReturn(INVITE_TOKEN);
 
     assertEquals(INVITE_TOKEN, _resolver.get(_dataFetchingEnvironment).join().getInviteToken());
   }