keycloak-aplcache

KEYCLOAK-6589: Optimize jpql in User search API This commit

2/13/2018 8:06:06 AM

Details

diff --git a/model/jpa/src/main/java/org/keycloak/models/jpa/entities/UserEntity.java b/model/jpa/src/main/java/org/keycloak/models/jpa/entities/UserEntity.java
index d0662cb..82062cc 100755
--- a/model/jpa/src/main/java/org/keycloak/models/jpa/entities/UserEntity.java
+++ b/model/jpa/src/main/java/org/keycloak/models/jpa/entities/UserEntity.java
@@ -17,6 +17,8 @@
 
 package org.keycloak.models.jpa.entities;
 
+import org.hibernate.annotations.Fetch;
+import org.hibernate.annotations.FetchMode;
 import org.keycloak.models.utils.KeycloakModelUtils;
 
 import javax.persistence.Access;
@@ -90,12 +92,15 @@ public class UserEntity {
     protected String realmId;
 
     @OneToMany(cascade = CascadeType.REMOVE, orphanRemoval = true, mappedBy="user")
+    @Fetch(FetchMode.SUBSELECT)
     protected Collection<UserAttributeEntity> attributes = new ArrayList<UserAttributeEntity>();
 
     @OneToMany(cascade = CascadeType.REMOVE, orphanRemoval = true, mappedBy="user")
+    @Fetch(FetchMode.SUBSELECT)
     protected Collection<UserRequiredActionEntity> requiredActions = new ArrayList<UserRequiredActionEntity>();
 
     @OneToMany(cascade = CascadeType.REMOVE, orphanRemoval = true, mappedBy="user")
+    @Fetch(FetchMode.SUBSELECT)
     protected Collection<CredentialEntity> credentials = new ArrayList<CredentialEntity>();
 
     @Column(name="FEDERATION_LINK")
diff --git a/model/jpa/src/main/java/org/keycloak/models/jpa/JpaUserProvider.java b/model/jpa/src/main/java/org/keycloak/models/jpa/JpaUserProvider.java
index 8eb8102..f20e87a 100755
--- a/model/jpa/src/main/java/org/keycloak/models/jpa/JpaUserProvider.java
+++ b/model/jpa/src/main/java/org/keycloak/models/jpa/JpaUserProvider.java
@@ -59,6 +59,7 @@ import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 /**
  * @author <a href="mailto:bill@burkecentral.com">Bill Burke</a>
@@ -945,6 +946,7 @@ public class JpaUserProvider implements UserProvider, UserCredentialStore {
         UserEntity userRef = em.getReference(UserEntity.class, user.getId());
         entity.setUser(userRef);
         em.persist(entity);
+
         MultivaluedHashMap<String, String> config = cred.getConfig();
         if (config != null && !config.isEmpty()) {
 
@@ -962,6 +964,11 @@ public class JpaUserProvider implements UserProvider, UserCredentialStore {
             }
 
         }
+
+        UserEntity userEntity = userInEntityManagerContext(user.getId());
+        if (userEntity != null) {
+            userEntity.getCredentials().add(entity);
+        }
         return toModel(entity);
     }
 
@@ -970,6 +977,10 @@ public class JpaUserProvider implements UserProvider, UserCredentialStore {
         CredentialEntity entity = em.find(CredentialEntity.class, id);
         if (entity == null) return false;
         em.remove(entity);
+        UserEntity userEntity = userInEntityManagerContext(user.getId());
+        if (userEntity != null) {
+            userEntity.getCredentials().remove(entity);
+        }
         return true;
     }
 
@@ -1017,11 +1028,19 @@ public class JpaUserProvider implements UserProvider, UserCredentialStore {
 
     @Override
     public List<CredentialModel> getStoredCredentialsByType(RealmModel realm, UserModel user, String type) {
-        UserEntity userEntity = em.getReference(UserEntity.class, user.getId());
-        TypedQuery<CredentialEntity> query = em.createNamedQuery("credentialByUserAndType", CredentialEntity.class)
-                .setParameter("type", type)
-                .setParameter("user", userEntity);
-        List<CredentialEntity> results = query.getResultList();
+        List<CredentialEntity> results;
+        UserEntity userEntity = userInEntityManagerContext(user.getId());
+        if (userEntity != null) {
+
+            // user already in persistence context, no need to execute a query
+            results = userEntity.getCredentials().stream().filter(it -> it.getType().equals(type)).collect(Collectors.toList());
+        } else {
+            userEntity = em.getReference(UserEntity.class, user.getId());
+            TypedQuery<CredentialEntity> query = em.createNamedQuery("credentialByUserAndType", CredentialEntity.class)
+                    .setParameter("type", type)
+                    .setParameter("user", userEntity);
+            results = query.getResultList();
+        }
         List<CredentialModel> rtn = new LinkedList<>();
         for (CredentialEntity entity : results) {
             rtn.add(toModel(entity));
@@ -1062,4 +1081,10 @@ public class JpaUserProvider implements UserProvider, UserCredentialStore {
             em.persist(user);
         }  
     }
+
+    private UserEntity userInEntityManagerContext(String id) {
+        UserEntity user = em.getReference(UserEntity.class, id);
+        boolean isLoaded = em.getEntityManagerFactory().getPersistenceUnitUtil().isLoaded(user);
+        return isLoaded ? user : null;
+    }
 }