Bot roles missing after edition (#9351)

openmetadata-service/src/main/java/org/openmetadata/service/resources/bots/BotResource.java

@@ -13,6 +13,8 @@
 
 package org.openmetadata.service.resources.bots;
 
+import static org.openmetadata.service.util.UserUtil.getRoleForBot;
+
 import io.swagger.annotations.Api;
 import io.swagger.v3.oas.annotations.ExternalDocumentation;
 import io.swagger.v3.oas.annotations.Operation;
@@ -23,7 +25,6 @@ import io.swagger.v3.oas.annotations.media.Schema;
 import io.swagger.v3.oas.annotations.parameters.RequestBody;
 import io.swagger.v3.oas.annotations.responses.ApiResponse;
 import java.io.IOException;
-import java.util.Arrays;
 import java.util.List;
 import java.util.UUID;
 import javax.json.JsonPatch;
@@ -95,7 +96,7 @@ public class BotResource extends EntityResource<Bot, BotRepository> {
 
       // Add role corresponding to the bot to the user
       // we need to set a mutable list here
-      user.setRoles(Arrays.asList(RoleResource.getRole(getRoleForBot(bot.getName()))));
+      user.setRoles(getRoleForBot(bot.getName()));
       user = UserUtil.addOrUpdateBotUser(user, config);
 
       bot.withId(UUID.randomUUID())
@@ -121,19 +122,6 @@ public class BotResource extends EntityResource<Bot, BotRepository> {
     }
   }
 
-  private static String getRoleForBot(String botName) {
-    switch (botName) {
-      case Entity.INGESTION_BOT_NAME:
-        return Entity.INGESTION_BOT_ROLE;
-      case Entity.QUALITY_BOT_NAME:
-        return Entity.QUALITY_BOT_ROLE;
-      case Entity.PROFILER_BOT_NAME:
-        return Entity.PROFILER_BOT_ROLE;
-      default:
-        throw new IllegalArgumentException("No role found for the bot " + botName);
-    }
-  }
-
   @Override
   public Bot addHref(UriInfo uriInfo, Bot entity) {
     Entity.withHref(uriInfo, entity.getBotUser());
@@ -422,6 +410,10 @@ public class BotResource extends EntityResource<Bot, BotRepository> {
       throw new IllegalArgumentException(
           String.format("Bot user [%s] is already used by [%s] bot", botUser.getName(), bot.getName()));
     }
+    // TODO: review this flow on https://github.com/open-metadata/OpenMetadata/issues/8321
+    if (originalBot != null) {
+      bot.setProvider(originalBot.getProvider());
+    }
     return bot;
   }
 

openmetadata-service/src/main/java/org/openmetadata/service/resources/teams/UserResource.java

@@ -90,6 +90,7 @@ import org.openmetadata.schema.type.EntityHistory;
 import org.openmetadata.schema.type.EntityReference;
 import org.openmetadata.schema.type.Include;
 import org.openmetadata.schema.type.MetadataOperation;
+import org.openmetadata.schema.type.ProviderType;
 import org.openmetadata.schema.type.Relationship;
 import org.openmetadata.service.Entity;
 import org.openmetadata.service.OpenMetadataApplicationConfig;
@@ -119,6 +120,7 @@ import org.openmetadata.service.util.JsonUtils;
 import org.openmetadata.service.util.PasswordUtil;
 import org.openmetadata.service.util.RestUtil;
 import org.openmetadata.service.util.ResultList;
+import org.openmetadata.service.util.UserUtil;
 
 @Slf4j
 @Path("/v1/users")
@@ -1040,6 +1042,12 @@ public class UserResource extends EntityResource<User, UserRepository> {
       throw new IllegalArgumentException(
           String.format("Bot user [%s] is already used by [%s] bot.", user.getName(), bot.getName()));
     }
+    // TODO: review this flow on https://github.com/open-metadata/OpenMetadata/issues/8321
+    if (original != null) {
+      user.setRoles(original.getRoles());
+    } else if (bot != null && ProviderType.SYSTEM.equals(bot.getProvider())) {
+      user.setRoles(UserUtil.getRoleForBot(botName));
+    }
     // TODO remove this
     addAuthMechanismToBot(user, create, uriInfo);
     RestUtil.PutResponse<User> response = dao.createOrUpdate(uriInfo, user);
@@ -1116,7 +1124,7 @@ public class UserResource extends EntityResource<User, UserRepository> {
   private User retrieveBotUser(User user, UriInfo uriInfo) {
     User original;
     try {
-      original = dao.getByName(uriInfo, user.getFullyQualifiedName(), new Fields(List.of("authenticationMechanism")));
+      original = dao.getByName(uriInfo, user.getFullyQualifiedName(), dao.getFieldsWithUserAuth("*"));
     } catch (EntityNotFoundException | IOException exc) {
       LOG.debug(String.format("User not found when adding auth mechanism for: [%s]", user.getName()));
       original = null;

openmetadata-service/src/main/java/org/openmetadata/service/util/UserUtil.java

@@ -13,6 +13,7 @@ import static org.openmetadata.service.resources.teams.UserResource.USER_PROTECT
 
 import at.favre.lib.crypto.bcrypt.BCrypt;
 import java.io.IOException;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Objects;
 import java.util.Set;
@@ -27,11 +28,13 @@ import org.openmetadata.schema.auth.SSOAuthMechanism;
 import org.openmetadata.schema.entity.teams.AuthenticationMechanism;
 import org.openmetadata.schema.entity.teams.User;
 import org.openmetadata.schema.security.client.OpenMetadataJWTClientConfig;
+import org.openmetadata.schema.type.EntityReference;
 import org.openmetadata.service.Entity;
 import org.openmetadata.service.OpenMetadataApplicationConfig;
 import org.openmetadata.service.exception.EntityNotFoundException;
 import org.openmetadata.service.jdbi3.EntityRepository;
 import org.openmetadata.service.jdbi3.UserRepository;
+import org.openmetadata.service.resources.teams.RoleResource;
 import org.openmetadata.service.security.jwt.JWTTokenGenerator;
 
 @Slf4j
@@ -219,4 +222,22 @@ public final class UserUtil {
       return null;
     }
   }
+
+  public static List<EntityReference> getRoleForBot(String botName) {
+    String botRole;
+    switch (botName) {
+      case Entity.INGESTION_BOT_NAME:
+        botRole = Entity.INGESTION_BOT_ROLE;
+        break;
+      case Entity.QUALITY_BOT_NAME:
+        botRole = Entity.QUALITY_BOT_ROLE;
+        break;
+      case Entity.PROFILER_BOT_NAME:
+        botRole = Entity.PROFILER_BOT_ROLE;
+        break;
+      default:
+        throw new IllegalArgumentException("No role found for the bot " + botName);
+    }
+    return Arrays.asList(RoleResource.getRole(botRole));
+  }
 }