Fixes #3702 Save user in authentication context to avoid multiple database lookups (#3703)

catalog-rest-service/src/main/java/org/openmetadata/catalog/exception/CatalogExceptionMessage.java

@@ -19,6 +19,9 @@ import java.util.UUID;
 public final class CatalogExceptionMessage {
   public static final String ENTITY_ALREADY_EXISTS = "Entity already exists";
   public static final String FERNET_KEY_NULL = "Fernet key is null";
+  public static final String FIELD_NOT_TOKENIZED = "Field is not tokenized";
+  public static final String FIELD_ALREADY_TOKENIZED = "Field is already tokenized";
+  public static final String INVALID_ENTITY_LINK = "Entity link must have both {arrayFieldName} and {arrayFieldValue}";
 
   private CatalogExceptionMessage() {}
 
@@ -66,18 +69,6 @@ public final class CatalogExceptionMessage {
     return String.format("Invalid service type `%s` for %s. Expected %s.", serviceType, entityType, expected);
   }
 
-  public static String invalidEntityLink() {
-    return "Entity link must have both {arrayFieldName} and {arrayFieldValue}";
-  }
-
-  public static String isNotTokenized() {
-    return "The field is not tokenized";
-  }
-
-  public static String isAlreadyTokenized() {
-    return "The field is already tokenized";
-  }
-
   public static String glossaryTermMismatch(String parentId, String glossaryId) {
     return String.format(
         "Invalid queryParameters - glossary term `parent` %s is not in the `glossary` %s", parentId, glossaryId);

catalog-rest-service/src/main/java/org/openmetadata/catalog/fernet/Fernet.java

@@ -14,8 +14,8 @@
 package org.openmetadata.catalog.fernet;
 
 import static org.openmetadata.catalog.exception.CatalogExceptionMessage.FERNET_KEY_NULL;
-import static org.openmetadata.catalog.exception.CatalogExceptionMessage.isAlreadyTokenized;
+import static org.openmetadata.catalog.exception.CatalogExceptionMessage.FIELD_ALREADY_TOKENIZED;
-import static org.openmetadata.catalog.exception.CatalogExceptionMessage.isNotTokenized;
+import static org.openmetadata.catalog.exception.CatalogExceptionMessage.FIELD_NOT_TOKENIZED;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.macasaet.fernet.Key;
@@ -86,7 +86,7 @@ public class Fernet {
 
   public String encrypt(@NonNull String secret) {
     if (secret.startsWith(FERNET_PREFIX)) {
-      throw new IllegalArgumentException(isAlreadyTokenized());
+      throw new IllegalArgumentException(FIELD_ALREADY_TOKENIZED);
     }
     if (isKeyDefined()) {
       Key key = new Key(fernetKey.split(",")[0]);
@@ -109,7 +109,7 @@ public class Fernet {
       List<Key> keys = Arrays.stream(fernetKey.split(",")).map(Key::new).collect(Collectors.toList());
       return token.validateAndDecrypt(keys, validator);
     }
-    throw new IllegalArgumentException(isNotTokenized());
+    throw new IllegalArgumentException(FIELD_NOT_TOKENIZED);
   }
 
   public static String decryptIfTokenized(String tokenized) {

catalog-rest-service/src/main/java/org/openmetadata/catalog/jdbi3/CollectionDAO.java

@@ -760,7 +760,8 @@ public interface CollectionDAO {
 
     default void deleteAllByPrefix(String fqnPrefix, int relation) {
       String prefix = String.format("%s%s%%", fqnPrefix, Entity.SEPARATOR);
-      String cond = String.format("WHERE (toFQN LIKE %s OR fromFQN LIKE %s) AND relation = %s)", prefix, relation);
+      String cond =
+          String.format("WHERE (toFQN LIKE %s OR fromFQN LIKE %s) AND relation = %s)", prefix, prefix, relation);
       deleteAllByPrefix(cond);
     }
 

catalog-rest-service/src/main/java/org/openmetadata/catalog/resources/feeds/MessageParser.java

@@ -13,7 +13,7 @@
 
 package org.openmetadata.catalog.resources.feeds;
 
-import static org.openmetadata.catalog.exception.CatalogExceptionMessage.invalidEntityLink;
+import static org.openmetadata.catalog.exception.CatalogExceptionMessage.INVALID_ENTITY_LINK;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -80,7 +80,7 @@ public final class MessageParser {
 
       if (arrayFieldValue != null) {
         if (arrayFieldName == null) {
-          throw new IllegalArgumentException(invalidEntityLink());
+          throw new IllegalArgumentException(INVALID_ENTITY_LINK);
         }
         this.linkType = LinkType.ENTITY_ARRAY_FIELD;
         this.fullyQualifiedFieldType = String.format("%s.%s.member", entityType, fieldName);

catalog-rest-service/src/main/java/org/openmetadata/catalog/security/AuthenticationContext.java

@@ -14,10 +14,14 @@
 package org.openmetadata.catalog.security;
 
 import java.security.Principal;
+import org.openmetadata.catalog.entity.teams.User;
+import org.openmetadata.catalog.util.EntityUtil.Fields;
 
 /** Holds context information of authenticated user, which will be used for authorization. */
 public final class AuthenticationContext {
   private final Principal principal;
+  private User user;
+  private Fields userFields;
 
   public AuthenticationContext(Principal principal) {
     this.principal = principal;
@@ -31,4 +35,20 @@ public final class AuthenticationContext {
   public String toString() {
     return "AuthenticationContext{" + ", principal=" + principal + '}';
   }
+
+  public User getUser() {
+    return user;
+  }
+
+  public void setUser(User user) {
+    this.user = user;
+  }
+
+  public Fields getUserFields() {
+    return userFields;
+  }
+
+  public void setUserFields(Fields userFields) {
+    this.userFields = userFields;
+  }
 }

catalog-rest-service/src/main/java/org/openmetadata/catalog/security/DefaultAuthorizer.java

@@ -38,7 +38,6 @@ import org.openmetadata.catalog.security.policyevaluator.PolicyEvaluator;
 import org.openmetadata.catalog.type.EntityReference;
 import org.openmetadata.catalog.type.Include;
 import org.openmetadata.catalog.type.MetadataOperation;
-import org.openmetadata.catalog.util.EntityUtil;
 import org.openmetadata.catalog.util.EntityUtil.Fields;
 import org.openmetadata.catalog.util.RestUtil;
 
@@ -52,7 +51,8 @@ public class DefaultAuthorizer implements Authorizer {
   private UserRepository userRepository;
 
   private PolicyEvaluator policyEvaluator;
-  private static final String FIELDS_PARAM = "roles,teams";
+  private Fields fieldsTeams;
+  private Fields fieldsRolesAndTeams;
 
   @Override
   public void init(AuthorizerConfiguration config, Jdbi dbi) throws IOException {
@@ -71,6 +71,8 @@ public class DefaultAuthorizer implements Authorizer {
     mayBeAddAdminUsers();
     mayBeAddBotUsers();
     this.policyEvaluator = PolicyEvaluator.getInstance();
+    this.fieldsTeams = userRepository.getFields("teams");
+    this.fieldsRolesAndTeams = userRepository.getFields("roles,teams");
   }
 
   private void mayBeAddAdminUsers() {
@@ -125,40 +127,29 @@ public class DefaultAuthorizer implements Authorizer {
 
   @Override
   public boolean hasPermissions(AuthenticationContext ctx, EntityReference owner) {
-    validateAuthenticationContext(ctx);
     // Since we have roles and operations. An Admin could enable updateDescription, tags, ownership permissions to
     // a role and assign that to the users who can update the entities. With this we can look at the owner as a strict
     // requirement to manage entities. So if owner is null we will not allow users to update entities. They can get a
     // role that allows them to update the entity.
-    if (owner == null) {
+    return isOwner(ctx, owner);
-      return false;
-    }
-    try {
-      User user = getUserFromAuthenticationContext(ctx);
-      return isOwnedByUser(user, owner);
-    } catch (IOException | EntityNotFoundException | ParseException ex) {
-      return false;
-    }
   }
 
   @Override
   public boolean hasPermissions(
       AuthenticationContext ctx, EntityReference entityReference, MetadataOperation operation) {
-    validateAuthenticationContext(ctx);
+    validate(ctx);
     try {
-      User user = getUserFromAuthenticationContext(ctx);
+      User user = getUserFromAuthenticationContext(ctx, fieldsRolesAndTeams);
       if (entityReference == null) {
         // In some cases there is no specific entity being acted upon. Eg: Lineage.
         return policyEvaluator.hasPermission(user, null, operation);
       }
 
-      Object entity =
+      Object entity = Entity.getEntity(entityReference, new Fields(List.of("tags", FIELD_OWNER)), Include.NON_DELETED);
-          Entity.getEntity(entityReference, new EntityUtil.Fields(List.of("tags", FIELD_OWNER)), Include.NON_DELETED);
       EntityReference owner = Entity.getEntityInterface(entity).getOwner();
 
       if (Entity.shouldHaveOwner(entityReference.getType()) && owner != null && isOwnedByUser(user, owner)) {
-        // Entity is owned by the user.
+        return true; // Entity is owned by the user.
-        return true;
       }
       return policyEvaluator.hasPermission(user, entity, operation);
     } catch (IOException | EntityNotFoundException | ParseException ex) {
@@ -168,7 +159,7 @@ public class DefaultAuthorizer implements Authorizer {
 
   @Override
   public List<MetadataOperation> listPermissions(AuthenticationContext ctx, EntityReference entityReference) {
-    validateAuthenticationContext(ctx);
+    validate(ctx);
 
     if (isAdmin(ctx) || isBot(ctx)) {
       // Admins and bots have permissions to do all operations.
@@ -176,12 +167,11 @@ public class DefaultAuthorizer implements Authorizer {
     }
 
     try {
-      User user = getUserFromAuthenticationContext(ctx);
+      User user = getUserFromAuthenticationContext(ctx, fieldsRolesAndTeams);
       if (entityReference == null) {
         return policyEvaluator.getAllowedOperations(user, null);
       }
-      Object entity =
+      Object entity = Entity.getEntity(entityReference, new Fields(List.of("tags", FIELD_OWNER)), Include.NON_DELETED);
-          Entity.getEntity(entityReference, new EntityUtil.Fields(List.of("tags", FIELD_OWNER)), Include.NON_DELETED);
       EntityReference owner = Entity.getEntityInterface(entity).getOwner();
       if (owner == null || isOwnedByUser(user, owner)) {
         // Entity does not have an owner or is owned by the user - allow all operations.
@@ -196,14 +186,12 @@ public class DefaultAuthorizer implements Authorizer {
   /** Checks if the user is same as owner or part of the team that is the owner. */
   private boolean isOwnedByUser(User user, EntityReference owner) {
     if (owner.getType().equals(Entity.USER) && owner.getName().equals(user.getName())) {
-      // Owner is same as user.
+      return true; // Owner is same as user.
-      return true;
     }
     if (owner.getType().equals(Entity.TEAM)) {
       for (EntityReference userTeam : user.getTeams()) {
         if (userTeam.getName().equals(owner.getName())) {
-          // Owner is a team, and the user is part of this team.
+          return true; // Owner is a team, and the user is part of this team.
-          return true;
         }
       }
     }
@@ -212,14 +200,10 @@ public class DefaultAuthorizer implements Authorizer {
 
   @Override
   public boolean isAdmin(AuthenticationContext ctx) {
-    validateAuthenticationContext(ctx);
+    validate(ctx);
-    String userName = SecurityUtil.getUserName(ctx);
     try {
-      User user = userRepository.getByName(null, userName, Fields.EMPTY_FIELDS);
+      User user = getUserFromAuthenticationContext(ctx, Fields.EMPTY_FIELDS);
-      if (user.getIsAdmin() == null) {
+      return Boolean.TRUE.equals(user.getIsAdmin());
-        return false;
-      }
-      return user.getIsAdmin();
     } catch (IOException | EntityNotFoundException | ParseException ex) {
       return false;
     }
@@ -227,14 +211,10 @@ public class DefaultAuthorizer implements Authorizer {
 
   @Override
   public boolean isBot(AuthenticationContext ctx) {
-    validateAuthenticationContext(ctx);
+    validate(ctx);
-    String userName = SecurityUtil.getUserName(ctx);
     try {
-      User user = userRepository.getByName(null, userName, Fields.EMPTY_FIELDS);
+      User user = getUserFromAuthenticationContext(ctx, Fields.EMPTY_FIELDS);
-      if (user.getIsBot() == null) {
+      return Boolean.TRUE.equals(user.getIsBot());
-        return false;
-      }
-      return user.getIsBot();
     } catch (IOException | EntityNotFoundException | ParseException ex) {
       return false;
     }
@@ -242,29 +222,41 @@ public class DefaultAuthorizer implements Authorizer {
 
   @Override
   public boolean isOwner(AuthenticationContext ctx, EntityReference owner) {
-    validateAuthenticationContext(ctx);
+    if (owner == null) {
-    String userName = SecurityUtil.getUserName(ctx);
+      return false;
+    }
+    validate(ctx);
     try {
-      User user = userRepository.getByName(null, userName, Fields.EMPTY_FIELDS);
+      User user = getUserFromAuthenticationContext(ctx, fieldsTeams);
-      if (owner == null) {
-        return false;
-      }
       return isOwnedByUser(user, owner);
     } catch (IOException | EntityNotFoundException | ParseException ex) {
       return false;
     }
   }
 
-  private void validateAuthenticationContext(AuthenticationContext ctx) {
+  private void validate(AuthenticationContext ctx) {
     if (ctx == null || ctx.getPrincipal() == null) {
       throw new AuthenticationException("No principal in AuthenticationContext");
     }
   }
 
-  private User getUserFromAuthenticationContext(AuthenticationContext ctx) throws IOException, ParseException {
+  private User getUserFromAuthenticationContext(AuthenticationContext ctx, Fields fields)
+      throws IOException, ParseException {
+    if (ctx.getUser() != null) {
+      // If a requested field is not present in the user, then add it
+      for (String field : fields.getList()) {
+        if (!ctx.getUserFields().contains(field)) {
+          userRepository.setFields(ctx.getUser(), userRepository.getFields(field));
+          ctx.getUserFields().add(fields);
+        }
+      }
+      return ctx.getUser();
+    }
     String userName = SecurityUtil.getUserName(ctx);
-    EntityUtil.Fields fields = userRepository.getFields(FIELDS_PARAM);
+    User user = userRepository.getByName(null, userName, fields);
-    return userRepository.getByName(null, userName, fields);
+    ctx.setUser(user);
+    ctx.setUserFields(fields);
+    return user;
   }
 
   private void addOrUpdateUser(User user) {

catalog-rest-service/src/main/java/org/openmetadata/catalog/util/EntityUtil.java

@@ -18,7 +18,6 @@ import static org.openmetadata.catalog.type.Include.ALL;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collections;
 import java.util.Comparator;
 import java.util.Date;
 import java.util.List;
@@ -327,21 +326,33 @@ public final class EntityUtil {
 
     public Fields(List<String> allowedFields, String fieldsParam) {
       if (fieldsParam == null || fieldsParam.isEmpty()) {
-        fieldList = Collections.emptyList();
+        fieldList = new ArrayList<>();
         return;
       }
       fieldList = Arrays.asList(fieldsParam.replace(" ", "").split(","));
       for (String field : fieldList) {
         if (!allowedFields.contains(field)) {
-
           throw new IllegalArgumentException(CatalogExceptionMessage.invalidField(field));
         }
       }
     }
 
+    @Override
+    public String toString() {
+      return fieldList.toString();
+    }
+
+    public void add(Fields fields) {
+      fieldList.addAll(fields.fieldList);
+    }
+
     public boolean contains(String field) {
       return fieldList.contains(field);
     }
+
+    public List<String> getList() {
+      return fieldList;
+    }
   }
 
   public static List<UUID> getIDList(List<EntityReference> refList) {