[KEYCLOAK-4902] - Using streams when fetching resources

model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/authorization/StoreFactoryCacheSession.java

@@ -628,6 +628,13 @@ public class StoreFactoryCacheSession implements CachedStoreFactoryProvider {
                     (revision, resources) -> new ResourceListQuery(revision, cacheKey, resources.stream().map(resource -> resource.getId()).collect(Collectors.toSet()), resourceServerId), resourceServerId);
         }
 
+        @Override
+        public void findByOwner(String ownerId, String resourceServerId, Consumer<Resource> consumer) {
+            String cacheKey = getResourceByOwnerCacheKey(ownerId, resourceServerId);
+            cacheQuery(cacheKey, ResourceListQuery.class, () -> getResourceStoreDelegate().findByOwner(ownerId, resourceServerId),
+                    (revision, resources) -> new ResourceListQuery(revision, cacheKey, resources.stream().map(resource -> resource.getId()).collect(Collectors.toSet()), resourceServerId), resourceServerId, consumer);
+        }
+
         @Override
         public List<Resource> findByUri(String uri, String resourceServerId) {
             if (uri == null) return null;
@@ -667,7 +674,19 @@ public class StoreFactoryCacheSession implements CachedStoreFactoryProvider {
                      (revision, resources) -> new ResourceListQuery(revision, cacheKey, resources.stream().map(resource -> resource.getId()).collect(Collectors.toSet()), resourceServerId), resourceServerId);
         }
 
+        @Override
+        public void findByType(String type, String resourceServerId, Consumer<Resource> consumer) {
+            if (type == null) return;
+            String cacheKey = getResourceByTypeCacheKey(type, resourceServerId);
+            cacheQuery(cacheKey, ResourceListQuery.class, () -> getResourceStoreDelegate().findByType(type, resourceServerId),
+                    (revision, resources) -> new ResourceListQuery(revision, cacheKey, resources.stream().map(resource -> resource.getId()).collect(Collectors.toSet()), resourceServerId), resourceServerId, consumer);
+        }
+
         private <R, Q extends ResourceQuery> List<R> cacheQuery(String cacheKey, Class<Q> queryType, Supplier<List<R>> resultSupplier, BiFunction<Long, List<R>, Q> querySupplier, String resourceServerId) {
+            return cacheQuery(cacheKey, queryType, resultSupplier, querySupplier, resourceServerId, null);
+        }
+
+        private <R, Q extends ResourceQuery> List<R> cacheQuery(String cacheKey, Class<Q> queryType, Supplier<List<R>> resultSupplier, BiFunction<Long, List<R>, Q> querySupplier, String resourceServerId, Consumer<R> consumer) {
             Q query = cache.get(cacheKey, queryType);
             if (query != null) {
                 logger.tracev("cache hit for key: {0}", cacheKey);
@@ -679,11 +698,31 @@ public class StoreFactoryCacheSession implements CachedStoreFactoryProvider {
                 if (invalidations.contains(cacheKey)) return model;
                 query = querySupplier.apply(loaded, model);
                 cache.addRevisioned(query, startupRevision);
+                if (consumer != null) {
+                    for (R resource : model) {
+                        consumer.accept(resource);
+                    }
+                }
                 return model;
             } else if (query.isInvalid(invalidations)) {
-                return resultSupplier.get();
+                List<R> result = resultSupplier.get();
+
+                if (consumer != null) {
+                    for (R resource : result) {
+                        consumer.accept(resource);
+                    }
+                }
+
+                return result;
             } else {
-                return query.getResources().stream().map(resourceId -> (R) findById(resourceId, resourceServerId)).collect(Collectors.toList());
+                Set<String> resources = query.getResources();
+
+                if (consumer != null) {
+                    resources.stream().map(resourceId -> (R) findById(resourceId, resourceServerId)).forEach(consumer);
+                    return Collections.emptyList();
+                }
+
+                return resources.stream().map(resourceId -> (R) findById(resourceId, resourceServerId)).collect(Collectors.toList());
             }
         }
     }

model/jpa/src/main/java/org/keycloak/authorization/jpa/store/JPAResourceStore.java

@@ -38,6 +38,8 @@ import java.util.ArrayList;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
+import java.util.function.Consumer;
 
 /**
  * @author <a href="mailto:psilva@redhat.com">Pedro Igor</a>
@@ -99,6 +101,15 @@ public class JPAResourceStore implements ResourceStore {
 
     @Override
     public List<Resource> findByOwner(String ownerId, String resourceServerId) {
+        List<Resource> list = new LinkedList<>();
+
+        findByOwner(ownerId, resourceServerId, list::add);
+
+        return list;
+    }
+
+    @Override
+    public void findByOwner(String ownerId, String resourceServerId, Consumer<Resource> consumer) {
         String queryName = "findResourceIdByOwner";
 
         if (resourceServerId == null) {
@@ -114,19 +125,12 @@ public class JPAResourceStore implements ResourceStore {
             query.setParameter("serverId", resourceServerId);
         }
 
-        List<String> result = query.getResultList();
-        List<Resource> list = new LinkedList<>();
         ResourceStore resourceStore = provider.getStoreFactory().getResourceStore();
 
-        for (String id : result) {
-            Resource resource = resourceStore.findById(id, resourceServerId);
-
-            if (resource != null) {
-                list.add(resource);
-            }
-        }
-
-        return list;
+        query.getResultList().stream()
+                .map(id -> resourceStore.findById(id, resourceServerId))
+                .filter(Objects::nonNull)
+                .forEach(consumer);
     }
 
     @Override
@@ -279,6 +283,15 @@ public class JPAResourceStore implements ResourceStore {
 
     @Override
     public List<Resource> findByType(String type, String resourceServerId) {
+        List<Resource> list = new LinkedList<>();
+
+        findByType(type, resourceServerId, list::add);
+
+        return list;
+    }
+
+    @Override
+    public void findByType(String type, String resourceServerId, Consumer<Resource> consumer) {
         TypedQuery<String> query = entityManager.createNamedQuery("findResourceIdByType", String.class);
 
         query.setFlushMode(FlushModeType.COMMIT);
@@ -286,18 +299,11 @@ public class JPAResourceStore implements ResourceStore {
         query.setParameter("ownerId", resourceServerId);
         query.setParameter("serverId", resourceServerId);
 
-        List<String> result = query.getResultList();
-        List<Resource> list = new LinkedList<>();
         ResourceStore resourceStore = provider.getStoreFactory().getResourceStore();
 
-        for (String id : result) {
-            Resource resource = resourceStore.findById(id, resourceServerId);
-
-            if (resource != null) {
-                list.add(resource);
-            }
-        }
-
-        return list;
+        query.getResultList().stream()
+                .map(id -> resourceStore.findById(id, resourceServerId))
+                .filter(Objects::nonNull)
+                .forEach(consumer);
     }
 }