GEN-927 - Add bot default roles (#18256)

* GEN-927 - Add bot default roles * GEN-927 - Add bot default roles
open-metadata · Oct 21, 2024 · 93b4c66 · 93b4c66
1 parent a3224f2
commit 93b4c66
Show file tree

Hide file tree

Showing 6 changed files with 104 additions and 51 deletions.
diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/RoleRepository.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/RoleRepository.java
@@ -35,6 +35,7 @@
 @Slf4j
 public class RoleRepository extends EntityRepository<Role> {
   public static final String DOMAIN_ONLY_ACCESS_ROLE = "DomainOnlyAccessRole";
+  public static final String DEFAULT_BOT_ROLE = "DefaultBotRole";
 
   public RoleRepository() {
     super(

diff --git a/...metadata-service/src/main/java/org/openmetadata/service/resources/teams/UserResource.java b/...metadata-service/src/main/java/org/openmetadata/service/resources/teams/UserResource.java
@@ -18,12 +18,15 @@
 import static javax.ws.rs.core.Response.Status.FORBIDDEN;
 import static javax.ws.rs.core.Response.Status.OK;
 import static org.openmetadata.common.utils.CommonUtil.listOf;
+import static org.openmetadata.common.utils.CommonUtil.nullOrEmpty;
 import static org.openmetadata.schema.api.teams.CreateUser.CreatePasswordType.ADMIN_CREATE;
 import static org.openmetadata.schema.auth.ChangePasswordRequest.RequestType.SELF;
 import static org.openmetadata.schema.entity.teams.AuthenticationMechanism.AuthType.BASIC;
 import static org.openmetadata.schema.entity.teams.AuthenticationMechanism.AuthType.JWT;
 import static org.openmetadata.schema.type.Include.ALL;
 import static org.openmetadata.service.exception.CatalogExceptionMessage.EMAIL_SENDING_ISSUE;
+import static org.openmetadata.service.jdbi3.RoleRepository.DEFAULT_BOT_ROLE;
+import static org.openmetadata.service.jdbi3.RoleRepository.DOMAIN_ONLY_ACCESS_ROLE;
 import static org.openmetadata.service.jdbi3.UserRepository.AUTH_MECHANISM_FIELD;
 import static org.openmetadata.service.secrets.ExternalSecretsManager.NULL_SECRET_STRING;
 import static org.openmetadata.service.security.jwt.JWTTokenGenerator.getExpiryDate;
@@ -51,6 +54,7 @@
 import java.io.IOException;
 import java.time.LocalDateTime;
 import java.time.ZoneId;
+import java.util.ArrayList;
 import java.util.Base64;
 import java.util.Date;
 import java.util.List;
@@ -104,7 +108,6 @@
 import org.openmetadata.schema.auth.RegistrationRequest;
 import org.openmetadata.schema.auth.RevokePersonalTokenRequest;
 import org.openmetadata.schema.auth.RevokeTokenRequest;
-import org.openmetadata.schema.auth.SSOAuthMechanism;
 import org.openmetadata.schema.auth.ServiceTokenType;
 import org.openmetadata.schema.auth.TokenRefreshRequest;
 import org.openmetadata.schema.auth.TokenType;
@@ -125,6 +128,7 @@
 import org.openmetadata.service.exception.EntityNotFoundException;
 import org.openmetadata.service.jdbi3.CollectionDAO;
 import org.openmetadata.service.jdbi3.ListFilter;
+import org.openmetadata.service.jdbi3.RoleRepository;
 import org.openmetadata.service.jdbi3.TokenRepository;
 import org.openmetadata.service.jdbi3.UserRepository;
 import org.openmetadata.service.jdbi3.UserRepository.UserCsv;
@@ -173,6 +177,7 @@ public class UserResource extends EntityResource<User, UserRepository> {
   public static final String USER_PROTECTED_FIELDS = "authenticationMechanism";
   private final JWTTokenGenerator jwtTokenGenerator;
   private final TokenRepository tokenRepository;
+  private final RoleRepository roleRepository;
   private AuthenticationConfiguration authenticationConfiguration;
   private AuthorizerConfiguration authorizerConfiguration;
   private final AuthenticatorHandler authHandler;
@@ -197,6 +202,7 @@ public UserResource(
     jwtTokenGenerator = JWTTokenGenerator.getInstance();
     allowedFields.remove(USER_PROTECTED_FIELDS);
     tokenRepository = Entity.getTokenRepository();
+    roleRepository = Entity.getRoleRepository();
     UserTokenCache.initialize();
     authHandler = authenticatorHandler;
   }
@@ -567,6 +573,7 @@ public Response createUser(
     User user = getUser(securityContext.getUserPrincipal().getName(), create);
     if (Boolean.TRUE.equals(user.getIsBot())) {
       addAuthMechanismToBot(user, create, uriInfo);
+      addRolesToBot(user, uriInfo);
     }
 
     //
@@ -696,8 +703,8 @@ public Response createOrUpdateUser(
           new OperationContext(entityType, EntityUtil.createOrUpdateOperation(resourceContext));
       authorizer.authorize(securityContext, createOperationContext, resourceContext);
     }
-    if (Boolean.TRUE.equals(create.getIsBot())) { // TODO expect bot to be created separately
-      return createOrUpdateBot(user, create, uriInfo, securityContext);
+    if (Boolean.TRUE.equals(create.getIsBot())) {
+      return createOrUpdateBotUser(user, create, uriInfo, securityContext);
     }
     PutResponse<User> response = repository.createOrUpdate(uriInfo, user);
     addHref(uriInfo, response.getEntity());
@@ -1454,7 +1461,7 @@ public void validateEmailAlreadyExists(String email) {
     }
   }
 
-  private Response createOrUpdateBot(
+  private Response createOrUpdateBotUser(
       User user, CreateUser create, UriInfo uriInfo, SecurityContext securityContext) {
     User original = retrieveBotUser(user, uriInfo);
     String botName = create.getBotName();
@@ -1488,8 +1495,9 @@ && userHasRelationshipWithAnyBot(original, bot)) {
               original.getAuthenticationMechanism());
       user.setRoles(original.getRoles());
     }
-    // TODO remove this
+    // TODO remove this -> Still valid TODO?
     addAuthMechanismToBot(user, create, uriInfo);
+    addRolesToBot(user, uriInfo);
     PutResponse<User> response = repository.createOrUpdate(uriInfo, user);
     decryptOrNullify(securityContext, response.getEntity());
     return response.toResponse();
@@ -1531,41 +1539,54 @@ private List<CollectionDAO.EntityRelationshipRecord> retrieveBotRelationshipsFor
     return repository.findToRecords(bot.getId(), Entity.BOT, Relationship.CONTAINS, Entity.USER);
   }
 
-  // TODO remove this
+  // TODO remove this -> still valid TODO?
   private void addAuthMechanismToBot(User user, @Valid CreateUser create, UriInfo uriInfo) {
     if (!Boolean.TRUE.equals(user.getIsBot())) {
       throw new IllegalArgumentException(
           "Authentication mechanism change is only supported for bot users");
     }
-    if (isValidAuthenticationMechanism(create)) {
-      AuthenticationMechanism authMechanism = create.getAuthenticationMechanism();
-      AuthenticationMechanism.AuthType authType = authMechanism.getAuthType();
-      switch (authType) {
-        case JWT -> {
-          User original = retrieveBotUser(user, uriInfo);
-          if (original == null
-              || !hasAJWTAuthMechanism(user, original.getAuthenticationMechanism())) {
-            JWTAuthMechanism jwtAuthMechanism =
-                JsonUtils.convertValue(authMechanism.getConfig(), JWTAuthMechanism.class);
-            authMechanism.setConfig(
-                jwtTokenGenerator.generateJWTToken(user, jwtAuthMechanism.getJWTTokenExpiry()));
-          } else {
-            authMechanism = original.getAuthenticationMechanism();
-          }
-        }
-        case SSO -> {
-          SSOAuthMechanism ssoAuthMechanism =
-              JsonUtils.convertValue(authMechanism.getConfig(), SSOAuthMechanism.class);
-          authMechanism.setConfig(ssoAuthMechanism);
-        }
-        default -> throw new IllegalArgumentException(
-            String.format("Not supported authentication mechanism type: [%s]", authType.value()));
-      }
-      user.setAuthenticationMechanism(authMechanism);
+    AuthenticationMechanism authMechanism = create.getAuthenticationMechanism();
+    User original = retrieveBotUser(user, uriInfo);
+    if (original == null || !hasAJWTAuthMechanism(user, original.getAuthenticationMechanism())) {
+      JWTAuthMechanism jwtAuthMechanism =
+          JsonUtils.convertValue(authMechanism.getConfig(), JWTAuthMechanism.class);
+      authMechanism.setConfig(
+          jwtTokenGenerator.generateJWTToken(user, jwtAuthMechanism.getJWTTokenExpiry()));
     } else {
-      throw new IllegalArgumentException(
-          String.format("Authentication mechanism is empty bot user: [%s]", user.getName()));
+      authMechanism = original.getAuthenticationMechanism();
+    }
+    user.setAuthenticationMechanism(authMechanism);
+  }
+
+  private void addRolesToBot(User user, UriInfo uriInfo) {
+    if (!Boolean.TRUE.equals(user.getIsBot())) {
+      throw new IllegalArgumentException("Bot roles are only supported for bot users");
+    }
+    User original = retrieveBotUser(user, uriInfo);
+    ArrayList<EntityReference> defaultBotRoles = getDefaultBotRoles(user);
+    // Keep the incoming roles of the created user
+    if (!nullOrEmpty(user.getRoles())) {
+      defaultBotRoles.addAll(user.getRoles());
+    }
+    // If user existed, merge roles
+    if (original != null && !nullOrEmpty(original.getRoles())) {
+      defaultBotRoles.addAll(original.getRoles());
+    }
+    user.setRoles(defaultBotRoles);
+  }
+
+  private ArrayList<EntityReference> getDefaultBotRoles(User user) {
+    ArrayList<EntityReference> defaultBotRoles = new ArrayList<>();
+    EntityReference defaultBotRole =
+        roleRepository.getReferenceByName(DEFAULT_BOT_ROLE, Include.NON_DELETED);
+    defaultBotRoles.add(defaultBotRole);
+
+    if (!nullOrEmpty(user.getDomains())) {
+      EntityReference domainOnlyAccessRole =
+          roleRepository.getReferenceByName(DOMAIN_ONLY_ACCESS_ROLE, Include.NON_DELETED);
+      defaultBotRoles.add(domainOnlyAccessRole);
     }
+    return defaultBotRoles;
   }
 
   @Nullable

diff --git a/openmetadata-service/src/main/resources/json/data/role/DefaultBotRole.json b/openmetadata-service/src/main/resources/json/data/role/DefaultBotRole.json
@@ -0,0 +1,17 @@
+{
+  "name": "DefaultBotRole",
+  "displayName": "Default Bot Role",
+  "description": "Role Corresponding to a Bot by default.",
+  "allowDelete": false,
+  "provider": "system",
+  "policies" : [
+    {
+      "type" : "policy",
+      "name" : "DefaultBotPolicy"
+    },
+    {
+      "type" : "policy",
+      "name" : "DataConsumerPolicy"
+    }
+  ]
+}
diff --git a/...metadata-service/src/test/java/org/openmetadata/service/resources/EntityResourceTest.java b/...metadata-service/src/test/java/org/openmetadata/service/resources/EntityResourceTest.java
@@ -272,6 +272,8 @@ public abstract class EntityResourceTest<T extends EntityInterface, K extends Cr
   public static EntityReference USER2_REF;
   public static User USER_TEAM21;
   public static User BOT_USER;
+  public static EntityReference DEFAULT_BOT_ROLE_REF;
+  public static EntityReference DOMAIN_ONLY_ACCESS_ROLE_REF;
 
   public static Team ORG_TEAM;
   public static Team TEAM1;

diff --git a/...a-service/src/test/java/org/openmetadata/service/resources/system/SystemResourceTest.java b/...a-service/src/test/java/org/openmetadata/service/resources/system/SystemResourceTest.java
@@ -40,12 +40,12 @@
 import org.openmetadata.schema.api.teams.CreateTeam;
 import org.openmetadata.schema.api.teams.CreateUser;
 import org.openmetadata.schema.api.tests.CreateTestSuite;
-import org.openmetadata.schema.auth.SSOAuthMechanism;
+import org.openmetadata.schema.auth.JWTAuthMechanism;
+import org.openmetadata.schema.auth.JWTTokenExpiry;
 import org.openmetadata.schema.email.SmtpSettings;
 import org.openmetadata.schema.entity.data.Table;
 import org.openmetadata.schema.entity.teams.AuthenticationMechanism;
 import org.openmetadata.schema.profiler.MetricType;
-import org.openmetadata.schema.security.client.GoogleSSOClientConfig;
 import org.openmetadata.schema.settings.Settings;
 import org.openmetadata.schema.settings.SettingsType;
 import org.openmetadata.schema.system.ValidationResponse;
@@ -321,13 +321,10 @@ void botUserCountCheck(TestInfo test) throws HttpResponseException {
             .withIsBot(true)
             .withAuthenticationMechanism(
                 new AuthenticationMechanism()
-                    .withAuthType(AuthenticationMechanism.AuthType.SSO)
+                    .withAuthType(AuthenticationMechanism.AuthType.JWT)
                     .withConfig(
-                        new SSOAuthMechanism()
-                            .withSsoServiceType(SSOAuthMechanism.SsoServiceType.GOOGLE)
-                            .withAuthConfig(
-                                new GoogleSSOClientConfig()
-                                    .withSecretKey("/fake/path/secret.json"))));
+                        new JWTAuthMechanism().withJWTTokenExpiry(JWTTokenExpiry.Unlimited)));
+
     userResourceTest.createEntity(createUser, ADMIN_AUTH_HEADERS);
 
     int afterUserCount = getEntitiesCount().getUserCount();

diff --git a/...data-service/src/test/java/org/openmetadata/service/resources/teams/UserResourceTest.java b/...data-service/src/test/java/org/openmetadata/service/resources/teams/UserResourceTest.java
@@ -44,6 +44,8 @@
 import static org.openmetadata.service.exception.CatalogExceptionMessage.notAdmin;
 import static org.openmetadata.service.exception.CatalogExceptionMessage.operationNotAllowed;
 import static org.openmetadata.service.exception.CatalogExceptionMessage.permissionNotAllowed;
+import static org.openmetadata.service.jdbi3.RoleRepository.DEFAULT_BOT_ROLE;
+import static org.openmetadata.service.jdbi3.RoleRepository.DOMAIN_ONLY_ACCESS_ROLE;
 import static org.openmetadata.service.resources.teams.UserResource.USER_PROTECTED_FIELDS;
 import static org.openmetadata.service.security.SecurityUtil.authHeaders;
 import static org.openmetadata.service.util.EntityUtil.fieldAdded;
@@ -109,18 +111,17 @@
 import org.openmetadata.schema.auth.RegistrationRequest;
 import org.openmetadata.schema.auth.RevokePersonalTokenRequest;
 import org.openmetadata.schema.auth.RevokeTokenRequest;
-import org.openmetadata.schema.auth.SSOAuthMechanism;
 import org.openmetadata.schema.entity.data.Table;
 import org.openmetadata.schema.entity.teams.AuthenticationMechanism;
 import org.openmetadata.schema.entity.teams.AuthenticationMechanism.AuthType;
 import org.openmetadata.schema.entity.teams.Role;
 import org.openmetadata.schema.entity.teams.Team;
 import org.openmetadata.schema.entity.teams.User;
-import org.openmetadata.schema.security.client.GoogleSSOClientConfig;
 import org.openmetadata.schema.type.ApiStatus;
 import org.openmetadata.schema.type.ChangeDescription;
 import org.openmetadata.schema.type.EntityReference;
 import org.openmetadata.schema.type.ImageList;
+import org.openmetadata.schema.type.Include;
 import org.openmetadata.schema.type.MetadataOperation;
 import org.openmetadata.schema.type.Profile;
 import org.openmetadata.schema.type.Webhook;
@@ -129,6 +130,7 @@
 import org.openmetadata.service.Entity;
 import org.openmetadata.service.auth.JwtResponse;
 import org.openmetadata.service.exception.CatalogExceptionMessage;
+import org.openmetadata.service.jdbi3.RoleRepository;
 import org.openmetadata.service.jdbi3.TeamRepository.TeamCsv;
 import org.openmetadata.service.jdbi3.UserRepository.UserCsv;
 import org.openmetadata.service.resources.EntityResourceTest;
@@ -150,11 +152,13 @@ public class UserResourceTest extends EntityResourceTest<User, CreateUser> {
   private static final Profile PROFILE =
       new Profile().withImages(new ImageList().withImage(URI.create("https://image.com")));
   private static final TeamResourceTest TEAM_TEST = new TeamResourceTest();
+  private final RoleRepository roleRepository;
 
   public UserResourceTest() {
     super(USER, User.class, UserList.class, "users", UserResource.FIELDS);
     supportedNameCharacters = "_-.";
     supportsSearchIndex = true;
+    roleRepository = Entity.getRoleRepository();
   }
 
   public void setupUsers(TestInfo test) throws HttpResponseException {
@@ -193,6 +197,11 @@ public void setupUsers(TestInfo test) throws HttpResponseException {
     Set<String> userFields = Entity.getEntityFields(User.class);
     userFields.remove("authenticationMechanism");
     BOT_USER = getEntityByName(INGESTION_BOT, String.join(",", userFields), ADMIN_AUTH_HEADERS);
+
+    // Get the bot roles
+    DEFAULT_BOT_ROLE_REF = roleRepository.getReferenceByName(DEFAULT_BOT_ROLE, Include.NON_DELETED);
+    DOMAIN_ONLY_ACCESS_ROLE_REF =
+        roleRepository.getReferenceByName(DOMAIN_ONLY_ACCESS_ROLE, Include.NON_DELETED);
   }
 
   @Test
@@ -886,28 +895,26 @@ protected void validateCommonEntityFields(User entity, CreateEntity create, Stri
   void put_generateToken_bot_user_200_ok() throws HttpResponseException {
     AuthenticationMechanism authMechanism =
         new AuthenticationMechanism()
-            .withAuthType(AuthType.SSO)
-            .withConfig(
-                new SSOAuthMechanism()
-                    .withSsoServiceType(SSOAuthMechanism.SsoServiceType.GOOGLE)
-                    .withAuthConfig(
-                        new GoogleSSOClientConfig().withSecretKey("/path/to/secret.json")));
+            .withAuthType(AuthType.JWT)
+            .withConfig(new JWTAuthMechanism().withJWTTokenExpiry(JWTTokenExpiry.Unlimited));
     CreateUser create =
         createBotUserRequest("ingestion-bot-jwt")
             .withEmail("ingestion-bot-jwt@email.com")
             .withRoles(List.of(ROLE1_REF.getId()))
             .withAuthenticationMechanism(authMechanism);
     User user = createEntity(create, USER_WITH_CREATE_HEADERS);
     user = getEntity(user.getId(), "*", ADMIN_AUTH_HEADERS);
-    assertEquals(1, user.getRoles().size());
+    // Has the given role and the default bot role
+    assertEquals(2, user.getRoles().size());
     TestUtils.put(
         getResource(String.format("users/generateToken/%s", user.getId())),
         new GenerateTokenRequest().withJWTTokenExpiry(JWTTokenExpiry.Seven),
         OK,
         ADMIN_AUTH_HEADERS);
     user = getEntity(user.getId(), "*", ADMIN_AUTH_HEADERS);
     assertNull(user.getAuthenticationMechanism());
-    assertEquals(1, user.getRoles().size());
+    // Has the given role and the default bot role
+    assertEquals(2, user.getRoles().size());
     JWTAuthMechanism jwtAuthMechanism =
         TestUtils.get(
             getResource(String.format("users/token/%s", user.getId())),
@@ -1441,6 +1448,14 @@ public void validateCreatedEntity(
     for (UUID roleId : listOrEmpty(createRequest.getRoles())) {
       expectedRoles.add(new EntityReference().withId(roleId).withType(Entity.ROLE));
     }
+
+    // bots are created with default roles
+    if (createRequest.getIsBot()) {
+      expectedRoles.add(DEFAULT_BOT_ROLE_REF);
+      if (!nullOrEmpty(createRequest.getDomains())) {
+        expectedRoles.add(DOMAIN_ONLY_ACCESS_ROLE_REF);
+      }
+    }
     assertRoles(user, expectedRoles);
 
     List<EntityReference> expectedTeams = new ArrayList<>();