Refactor UserAwareKnowledgeManager to get user from UserManager inste…

…ad of repository Signed-off-by: Florian Hotze <[email protected]>
llamara-ai · Feb 3, 2025 · 23f19ae · 23f19ae
1 parent 0c9ed24
commit 23f19ae
Show file tree

Hide file tree

Showing 7 changed files with 43 additions and 29 deletions.
diff --git a/...main/java/com/github/llamara/ai/internal/security/knowledge/UserKnowledgeManagerImpl.java b/...main/java/com/github/llamara/ai/internal/security/knowledge/UserKnowledgeManagerImpl.java
@@ -31,7 +31,6 @@
 import com.github.llamara.ai.internal.security.user.User;
 import com.github.llamara.ai.internal.security.user.UserManager;
 import com.github.llamara.ai.internal.security.user.UserNotFoundException;
-import com.github.llamara.ai.internal.security.user.UserRepository;
 
 import java.io.IOException;
 import java.nio.file.Path;
@@ -58,8 +57,7 @@
 public class UserKnowledgeManagerImpl implements UserKnowledgeManager {
     private final KnowledgeManager delegate;
     private final SecurityConfig config;
-    private final UserAwareKnowledgeRepository userAwareKnowledgeRepository;
-    private final UserRepository userRepository;
+    private final UserAwareKnowledgeRepository userAwareRepository;
     private final UserManager userManager;
 
     private final SecurityIdentity identity;
@@ -68,14 +66,12 @@ public class UserKnowledgeManagerImpl implements UserKnowledgeManager {
     UserKnowledgeManagerImpl(
             KnowledgeManager delegate,
             SecurityConfig config,
-            UserAwareKnowledgeRepository userAwareKnowledgeRepository,
-            UserRepository userRepository,
+            UserAwareKnowledgeRepository userAwareRepository,
             UserManager userManager,
             SecurityIdentity identity) {
         this.delegate = delegate;
         this.config = config;
-        this.userAwareKnowledgeRepository = userAwareKnowledgeRepository;
-        this.userRepository = userRepository;
+        this.userAwareRepository = userAwareRepository;
         this.userManager = userManager;
         this.identity = identity;
     }
@@ -98,15 +94,15 @@ public Collection<Knowledge> getAllKnowledge() {
                 "Authenticated, non-admin user '%s' requested knowledge, returning user knowledge"
                         + " and public knowledge.",
                 username);
-        Set<Knowledge> userKnowledge = new HashSet<>(userManager.getUser().getKnowledge());
+        Set<Knowledge> userKnowledge = new HashSet<>(userManager.getCurrentUser().getKnowledge());
         userKnowledge.addAll(publicKnowledge);
         return userKnowledge;
     }
 
     @Override
     public Knowledge getKnowledge(UUID id) throws KnowledgeNotFoundException {
         userManager.enforceRegistered();
-        Knowledge knowledge = userAwareKnowledgeRepository.findById(id);
+        Knowledge knowledge = userAwareRepository.findById(id);
         if (knowledge == null) {
             throw new KnowledgeNotFoundException(id);
         }
@@ -167,8 +163,7 @@ public UUID addSource(Path file, String fileName, String contentType)
         }
 
         String checksum = Utils.generateChecksum(file);
-        Optional<Knowledge> existingKnowledge =
-                userAwareKnowledgeRepository.existsChecksum(checksum);
+        Optional<Knowledge> existingKnowledge = userAwareRepository.existsChecksum(checksum);
         if (existingKnowledge.isPresent()) {
             return existingKnowledge.get().getId();
         }
@@ -196,10 +191,7 @@ public void setPermission(UUID id, String username, Permission permission)
             throws KnowledgeNotFoundException,
                     UserNotFoundException,
                     IllegalPermissionModificationException {
-        User user = userRepository.findByUsername(username);
-        if (user == null) {
-            throw new UserNotFoundException(username);
-        }
+        User user = userManager.getUser(username);
         setPermission(id, user, permission);
     }
 
@@ -215,10 +207,7 @@ public void removePermission(UUID id, String username)
             throws KnowledgeNotFoundException,
                     UserNotFoundException,
                     IllegalPermissionModificationException {
-        User user = userRepository.findByUsername(username);
-        if (user == null) {
-            throw new UserNotFoundException(username);
-        }
+        User user = userManager.getUser(username);
         removePermission(id, user);
     }
 

diff --git a/src/main/java/com/github/llamara/ai/internal/security/user/AnonymousUserManagerImpl.java b/src/main/java/com/github/llamara/ai/internal/security/user/AnonymousUserManagerImpl.java
@@ -58,10 +58,19 @@ public void delete() {
     }
 
     @Override
-    public User getUser() {
+    public User getCurrentUser() {
         return Users.ANY;
     }
 
+    @Override
+    public User getUser(String username) throws UserNotFoundException {
+        User user = userRepository.findByUsername(username);
+        if (user == null) {
+            throw new UserNotFoundException(username);
+        }
+        return user;
+    }
+
     @Override
     public User getUserAny() {
         return userRepository.findByUsername(Users.ANY_USERNAME);

diff --git a/src/main/java/com/github/llamara/ai/internal/security/user/AuthenticatedUserManagerImpl.java b/src/main/java/com/github/llamara/ai/internal/security/user/AuthenticatedUserManagerImpl.java
@@ -97,7 +97,7 @@ public void enforceRegistered() {
 
     @Override
     public void delete() {
-        User user = getUser();
+        User user = getCurrentUser();
         for (Session session : user.getSessions()) {
             try {
                 authenticatedUserSessionManager.deleteSession(session.getId());
@@ -138,14 +138,23 @@ public void delete() {
     }
 
     @Override
-    public User getUser() throws UserNotRegisteredException {
+    public User getCurrentUser() throws UserNotRegisteredException {
         User user = userRepository.findByUsername(identity.getPrincipal().getName());
         if (user == null) {
             throw new UserNotRegisteredException(identity.getPrincipal().getName());
         }
         return user;
     }
 
+    @Override
+    public User getUser(String username) throws UserNotFoundException {
+        User user = userRepository.findByUsername(username);
+        if (user == null) {
+            throw new UserNotFoundException(username);
+        }
+        return user;
+    }
+
     @Override
     public User getUserAny() {
         return userRepository.findByUsername(Users.ANY_USERNAME);

diff --git a/src/main/java/com/github/llamara/ai/internal/security/user/UserManager.java b/src/main/java/com/github/llamara/ai/internal/security/user/UserManager.java
@@ -59,7 +59,16 @@ public interface UserManager {
      * @return the user
      * @throws UserNotRegisteredException if the user is not registered
      */
-    User getUser() throws UserNotRegisteredException;
+    User getCurrentUser() throws UserNotRegisteredException;
+
+    /**
+     * Get the user with the given username.
+     *
+     * @param username name of the user
+     * @return the user
+     * @throws UserNotFoundException if no user with the given username was found
+     */
+    User getUser(String username) throws UserNotFoundException;
 
     /**
      * Get {@link com.github.llamara.ai.internal.security.Users#ANY} from persistence.

diff --git a/.../java/com/github/llamara/ai/internal/security/knowledge/UserKnowledgeManagerImplTest.java b/.../java/com/github/llamara/ai/internal/security/knowledge/UserKnowledgeManagerImplTest.java
@@ -116,7 +116,6 @@ void setup() {
                         knowledgeManager,
                         config,
                         userAwareKnowledgeRepository,
-                        userRepository,
                         authenticatedUserManager,
                         identity);
 
@@ -524,7 +523,6 @@ void setup()
                             knowledgeManager,
                             config,
                             userAwareKnowledgeRepository,
-                            userRepository,
                             anonymousUserManager,
                             identity);
 

diff --git a/src/test/java/com/github/llamara/ai/internal/security/user/AnonymousUserManagerTest.java b/src/test/java/com/github/llamara/ai/internal/security/user/AnonymousUserManagerTest.java
@@ -77,8 +77,8 @@ protected User getUserAnyFromPersistence() {
     }
 
     @Test
-    void getUserReturnsUserAnyFromPersistence() {
-        assertEquals(getUserAnyFromPersistence(), userManager.getUser());
+    void getCurrentUserReturnsUserAnyFromPersistence() {
+        assertEquals(getUserAnyFromPersistence(), userManager.getCurrentUser());
     }
 
     @Test

diff --git a/src/test/java/com/github/llamara/ai/internal/security/user/AuthenticatedUserManagerTest.java b/src/test/java/com/github/llamara/ai/internal/security/user/AuthenticatedUserManagerTest.java
@@ -104,8 +104,8 @@ void deleteThrowsAndDoesNothingIfNotExists()
     }
 
     @Test
-    void getUserThrowsIfNotRegistered() {
-        assertThrows(UserNotRegisteredException.class, () -> userManager.getUser());
+    void getCurrentUserThrowsIfNotRegistered() {
+        assertThrows(UserNotRegisteredException.class, () -> userManager.getCurrentUser());
     }
 
     @Transactional