simplify storage changes

Signed-off-by: Michal Hajas <mhajas@redhat.com>
keycloak · Oct 24, 2023 · ba1d7c2 · ba1d7c2
1 parent fb97b64
commit ba1d7c2
Show file tree

Hide file tree

Showing 19 changed files with 137 additions and 194 deletions.
diff --git a/docs/documentation/upgrading/topics/keycloak/changes-23_0_0.adoc b/docs/documentation/upgrading/topics/keycloak/changes-23_0_0.adoc
@@ -112,10 +112,18 @@ hence `RegistrationProfile` is not needed anymore. There is usually no further a
 
 * `RealmModel#getTopLevelGroupsStream()` and overloaded methods are now deprecated
 
-= `GroupProvider` and `GroupLookupProvider` changes
+= `GroupProvider` changes
 
-Methods have been added to allow for searching and paging through groups and subgroups.
-If you implement these interfaces you will need to implement the new method `GroupLookupProvider#searchForSubgroupsByParentIdNameStream()`
+A new method has been added to allow for searching and paging through top level groups.
+If you implement this interface you will need to implement the following method:
+[source,java]
+----
+Stream<GroupModel> getTopLevelGroupsStream(RealmModel realm,
+                                           String search,
+                                           Boolean exact,
+                                           Integer firstResult,
+                                           Integer maxResults)
+----
 
 = `GroupRepresentation` changes
 

diff --git a/.../main/java/org/keycloak/storage/ldap/mappers/membership/group/GroupLDAPStorageMapper.java b/.../main/java/org/keycloak/storage/ldap/mappers/membership/group/GroupLDAPStorageMapper.java
@@ -815,7 +815,7 @@ protected Stream<GroupModel> getKcSubGroups(RealmModel realm, GroupModel parentG
             parentGroup = getKcGroupsPathGroup(realm);
         }
         return parentGroup == null ? session.groups().getTopLevelGroupsStream(realm) :
-            session.groups().searchForSubgroupsByParentIdStream(realm, parentGroup.getId(), null, null);
+            parentGroup.getSubGroupsStream();
     }
 
     /**

diff --git a/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/GroupAdapter.java b/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/GroupAdapter.java
@@ -237,13 +237,29 @@ public Stream<GroupModel> getSubGroupsStream() {
         return subGroups.stream().sorted(GroupModel.COMPARE_BY_NAME);
     }
 
+    @Override
+    public Stream<GroupModel> getSubGroupsStream(String search, Integer firstResult, Integer maxResults) {
+        if (isUpdated()) return updated.getSubGroupsStream(search, firstResult, maxResults);
+        return modelSupplier.get().getSubGroupsStream(search, firstResult, maxResults);
+    }
+
+    @Override
+    public Stream<GroupModel> getSubGroupsStream(Integer firstResult, Integer maxResults) {
+        if (isUpdated()) return updated.getSubGroupsStream(firstResult, maxResults);
+        return modelSupplier.get().getSubGroupsStream(firstResult, maxResults);
+    }
+
     @Override
     public Stream<GroupModel> getSubGroupsStream(String search, Boolean exact, Integer firstResult, Integer maxResults) {
         if (isUpdated()) return updated.getSubGroupsStream(search, exact, firstResult, maxResults);
-        return cached.getSubGroups(modelSupplier, search, exact, firstResult, maxResults);
+        return modelSupplier.get().getSubGroupsStream(search, exact, firstResult, maxResults);
     }
 
-
+    @Override
+    public Long getSubGroupsCount() {
+        if (isUpdated()) return updated.getSubGroupsCount();
+        return modelSupplier.get().getSubGroupsCount();
+    }
 
     @Override
     public void setParent(GroupModel group) {

diff --git a/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/RealmCacheSession.java b/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/RealmCacheSession.java
@@ -989,12 +989,6 @@ public Long getGroupsCount(RealmModel realm, Stream<String> ids, String search)
     public Long getGroupsCount(RealmModel realm, Boolean onlyTopGroups) {
         return getGroupDelegate().getGroupsCount(realm, onlyTopGroups);
     }
-
-    @Override
-    public Long getSubGroupsCount(RealmModel realm, String parentId) {
-        return getGroupDelegate().getSubGroupsCount(realm, parentId);
-    }
-
     @Override
     public long getClientsCount(RealmModel realm) {
         return getClientDelegate().getClientsCount(realm);
@@ -1058,11 +1052,6 @@ public Stream<GroupModel> searchForGroupByNameStream(RealmModel realm, String se
        return getGroupDelegate().searchForGroupByNameStream(realm, search, exact, firstResult, maxResults);
     }
 
-    @Override
-    public Stream<GroupModel> searchForSubgroupsByParentIdNameStream(RealmModel realm, String id, String search, Boolean exact, Integer firstResult, Integer maxResults) {
-        return getGroupDelegate().searchForSubgroupsByParentIdNameStream(realm, id, search, exact, firstResult, maxResults);
-    }
-
     @Override
     public Stream<GroupModel> searchGroupsByAttributes(RealmModel realm, Map<String, String> attributes, Integer firstResult, Integer maxResults) {
         return getGroupDelegate().searchGroupsByAttributes(realm, attributes, firstResult, maxResults);

diff --git a/...l/infinispan/src/main/java/org/keycloak/models/cache/infinispan/entities/CachedGroup.java b/...l/infinispan/src/main/java/org/keycloak/models/cache/infinispan/entities/CachedGroup.java
@@ -17,7 +17,6 @@
 
 package org.keycloak.models.cache.infinispan.entities;
 
-import java.util.stream.Stream;
 import org.keycloak.common.util.MultivaluedHashMap;
 import org.keycloak.models.GroupModel;
 import org.keycloak.models.RealmModel;
@@ -80,8 +79,4 @@ public String getParentId() {
     public Set<String> getSubGroups(Supplier<GroupModel> group) {
         return subGroups.get(group);
     }
-
-    public Stream<GroupModel> getSubGroups(Supplier<GroupModel> group, String search, Boolean exact, Integer firstResult, Integer maxResults) {
-        return group.get().getSubGroupsStream(search, exact, firstResult, maxResults);
-    }
 }
diff --git a/model/jpa/src/main/java/org/keycloak/models/jpa/GroupAdapter.java b/model/jpa/src/main/java/org/keycloak/models/jpa/GroupAdapter.java
@@ -39,6 +39,7 @@
 import java.util.stream.Stream;
 import jakarta.persistence.LockModeType;
 
+import static org.keycloak.models.jpa.PaginationUtils.paginateQuery;
 import static org.keycloak.utils.StreamsUtil.closing;
 
 /**
@@ -47,13 +48,11 @@
  */
 public class GroupAdapter implements GroupModel , JpaModel<GroupEntity> {
 
-    protected KeycloakSession session;
     protected GroupEntity group;
     protected EntityManager em;
     protected RealmModel realm;
 
-    public GroupAdapter(KeycloakSession session, RealmModel realm, EntityManager em, GroupEntity group) {
-        this.session = session;
+    public GroupAdapter(RealmModel realm, EntityManager em, GroupEntity group) {
         this.em = em;
         this.group = group;
         this.realm = realm;
@@ -124,15 +123,35 @@ public void removeChild(GroupModel subGroup) {
 
     @Override
     public Stream<GroupModel> getSubGroupsStream() {
-        return session.groups().searchForSubgroupsByParentIdStream(realm, group.getId(), -1, -1);
+        return getSubGroupsStream("", false, -1, -1);
     }
 
     @Override
     public Stream<GroupModel> getSubGroupsStream(String search, Boolean exact, Integer firstResult, Integer maxResults) {
-        if(search == null || search.isEmpty()) {
-            return session.groups().searchForSubgroupsByParentIdStream(realm, group.getId(), firstResult, maxResults);
+        TypedQuery<String> query;
+        if (Boolean.TRUE.equals(exact)) {
+            query = em.createNamedQuery("getGroupIdsByParentAndName", String.class);
+        } else {
+            query = em.createNamedQuery("getGroupIdsByParentAndNameContaining", String.class);
         }
-        return session.groups().searchForSubgroupsByParentIdNameStream(realm, group.getId(), search, exact, firstResult, maxResults);
+        query.setParameter("realm", realm.getId())
+                .setParameter("parent", group.getId())
+                .setParameter("search", search);
+
+        return closing(paginateQuery(query, firstResult, maxResults).getResultStream()
+                .map(realm::getGroupById)
+                // In concurrent tests, the group might be deleted in another thread, therefore, skip those null values.
+                .filter(Objects::nonNull)
+                .sorted(GroupModel.COMPARE_BY_NAME)
+        );
+    }
+
+    @Override
+    public Long getSubGroupsCount() {
+        return em.createNamedQuery("getGroupCountByParent", Long.class)
+                .setParameter("realm", realm.getId())
+                .setParameter("parent", group.getId())
+                .getSingleResult();
     }
 
     @Override

diff --git a/model/jpa/src/main/java/org/keycloak/models/jpa/JpaRealmProvider.java b/model/jpa/src/main/java/org/keycloak/models/jpa/JpaRealmProvider.java
@@ -431,7 +431,7 @@ public GroupModel getGroupById(RealmModel realm, String id) {
         GroupEntity groupEntity = em.find(GroupEntity.class, id);
         if (groupEntity == null) return null;
         if (!groupEntity.getRealm().equals(realm.getId())) return null;
-        GroupAdapter adapter =  new GroupAdapter(session, realm, em, groupEntity);
+        GroupAdapter adapter =  new GroupAdapter(realm, em, groupEntity);
         return adapter;
     }
 
@@ -577,14 +577,6 @@ public Long getGroupsCount(RealmModel realm, Boolean onlyTopGroups) {
         }
     }
 
-    @Override
-    public Long getSubGroupsCount(RealmModel realm, String parentId) {
-        return em.createNamedQuery("getGroupCountByParent", Long.class)
-            .setParameter("realm", realm.getId())
-            .setParameter("parent", parentId)
-            .getSingleResult();
-    }
-
     @Override
     public long getClientsCount(RealmModel realm) {
         final Long res = em.createNamedQuery("getRealmClientsCount", Long.class)
@@ -606,7 +598,7 @@ public Stream<GroupModel> getGroupsByRoleStream(RealmModel realm, RoleModel role
         Stream<GroupEntity> results = paginateQuery(query, firstResult, maxResults).getResultStream();
 
         return closing(results
-        		.map(g -> (GroupModel) new GroupAdapter(session, realm, em, g))
+        		.map(g -> (GroupModel) new GroupAdapter(realm, em, g))
                 .sorted(GroupModel.COMPARE_BY_NAME));
     }
 
@@ -659,7 +651,7 @@ public KeycloakSession getKeycloakSession() {
 
         realm.removeDefaultGroup(group);
 
-        session.groups().searchForSubgroupsByParentIdStream(realm, group.getId(), -1, -1).forEach(realm::removeGroup);
+        group.getSubGroupsStream().forEach(realm::removeGroup);
 
         GroupEntity groupEntity = em.find(GroupEntity.class, group.getId(), LockModeType.PESSIMISTIC_WRITE);
         if ((groupEntity == null) || (!groupEntity.getRealm().equals(realm.getId()))) {
@@ -689,7 +681,7 @@ public GroupModel createGroup(RealmModel realm, String id, String name, GroupMod
         em.persist(groupEntity);
         em.flush();
 
-        return new GroupAdapter(session, realm, em, groupEntity);
+        return new GroupAdapter(realm, em, groupEntity);
     }
 
     @Override
@@ -1029,26 +1021,6 @@ public Stream<GroupModel> searchForGroupByNameStream(RealmModel realm, String se
         return closing(groups.map(id -> session.groups().getGroupById(realm, id)).sorted(GroupModel.COMPARE_BY_NAME).distinct());
     }
 
-    @Override
-    public Stream<GroupModel> searchForSubgroupsByParentIdNameStream(RealmModel realm, String id, String search, Boolean exact, Integer firstResult, Integer maxResults) {
-        TypedQuery<String> query;
-        if (Boolean.TRUE.equals(exact)) {
-            query = em.createNamedQuery("getGroupIdsByParentAndName", String.class);
-        } else {
-            query = em.createNamedQuery("getGroupIdsByParentAndNameContaining", String.class);
-        }
-        query.setParameter("realm", realm.getId())
-            .setParameter("parent", id)
-            .setParameter("search", search);
-
-        return closing(paginateQuery(query, firstResult, maxResults).getResultStream()
-            .map(realm::getGroupById)
-            // In concurrent tests, the group might be deleted in another thread, therefore, skip those null values.
-            .filter(Objects::nonNull)
-            .sorted(GroupModel.COMPARE_BY_NAME)
-        );
-    }
-
     @Override
     public Stream<GroupModel> searchGroupsByAttributes(RealmModel realm, Map<String, String> attributes, Integer firstResult, Integer maxResults) {
         Map<String, String> filteredAttributes = groupSearchableAttributes == null || groupSearchableAttributes.isEmpty()
@@ -1083,7 +1055,7 @@ public Stream<GroupModel> searchGroupsByAttributes(RealmModel realm, Map<String,
 
         TypedQuery<GroupEntity> query = em.createQuery(queryBuilder);
         return closing(paginateQuery(query, firstResult, maxResults).getResultStream())
-                .map(g -> new GroupAdapter(session, realm, em, g));
+                .map(g -> new GroupAdapter(realm, em, g));
     }
 
     @Override

diff --git a/model/legacy-private/src/main/java/org/keycloak/storage/GroupStorageManager.java b/model/legacy-private/src/main/java/org/keycloak/storage/GroupStorageManager.java
@@ -103,11 +103,6 @@ public Long getGroupsCount(RealmModel realm, Boolean onlyTopGroups) {
         return localStorage().getGroupsCount(realm, onlyTopGroups);
     }
 
-    @Override
-    public Long getSubGroupsCount(RealmModel realm, String parentId) {
-        return localStorage().getSubGroupsCount(realm, parentId);
-    }
-
     @Override
     public Long getGroupsCountByNameContaining(RealmModel realm, String search) {
         return localStorage().getGroupsCountByNameContaining(realm, search);
@@ -123,11 +118,6 @@ public Stream<GroupModel> getTopLevelGroupsStream(RealmModel realm, String searc
         return localStorage().getTopLevelGroupsStream(realm, search, exact, firstResult, maxResults);
     }
 
-    @Override
-    public Stream<GroupModel> searchForSubgroupsByParentIdNameStream(RealmModel realm, String id, String search, Boolean exact, Integer firstResult, Integer maxResults) {
-        return localStorage().searchForSubgroupsByParentIdNameStream(realm, id, search, exact, firstResult, maxResults);
-    }
-
     @Override
     public GroupModel createGroup(RealmModel realm, String id, String name, GroupModel toParent) {
         return localStorage().createGroup(realm, id, name, toParent);

diff --git a/model/map/src/main/java/org/keycloak/models/map/group/MapGroupAdapter.java b/model/map/src/main/java/org/keycloak/models/map/group/MapGroupAdapter.java
@@ -32,7 +32,7 @@
 import java.util.stream.Stream;
 
 
-public class MapGroupAdapter extends AbstractGroupModel<MapGroupEntity> {
+public abstract class MapGroupAdapter extends AbstractGroupModel<MapGroupEntity> {
     public MapGroupAdapter(KeycloakSession session, RealmModel realm, MapGroupEntity entity) {
         super(session, realm, entity);
     }
@@ -100,19 +100,6 @@ public String getParentId() {
         return entity.getParentId();
     }
 
-    @Override
-    public Stream<GroupModel> getSubGroupsStream() {
-        return session.groups().searchForSubgroupsByParentIdStream(realm, getId(), -1, -1);
-    }
-
-    @Override
-    public Stream<GroupModel> getSubGroupsStream(String search, Boolean exact, Integer firstResult, Integer maxResults) {
-        if(search == null || search.isEmpty()) {
-            return session.groups().searchForSubgroupsByParentIdStream(realm, getId(), firstResult, maxResults);
-        }
-        return session.groups().searchForSubgroupsByParentIdNameStream(realm, getId(), search, exact, firstResult, maxResults);
-    }
-
     @Override
     public void setParent(GroupModel group) {
         if (group == null) {
@@ -172,7 +159,6 @@ public void grantRole(RoleModel role) {
     @Override
     public Stream<RoleModel> getRoleMappingsStream() {
         Set<String> grantedRoles = entity.getGrantedRoles();
-        // TODO: doing lookups on an adapter for a model object, should this be moved to a lookup provider to fit the general pattern?
         return grantedRoles == null ? Stream.empty() : grantedRoles.stream()
             .map(roleId -> session.roles().getRoleById(realm, roleId));
     }

diff --git a/model/map/src/main/java/org/keycloak/models/map/group/MapGroupProvider.java b/model/map/src/main/java/org/keycloak/models/map/group/MapGroupProvider.java
@@ -70,7 +70,12 @@ private MapStorage<MapGroupEntity, GroupModel> storeWithRealm(RealmModel realm)
 
     private Function<MapGroupEntity, GroupModel> entityToAdapterFunc(RealmModel realm) {
         // Clone entity before returning back, to avoid giving away a reference to the live object to the caller
-        return origEntity -> new MapGroupAdapter(session, realm, origEntity);
+        return origEntity -> new MapGroupAdapter(session, realm, origEntity) {
+            @Override
+            public Stream<GroupModel> getSubGroupsStream() {
+                return getGroupsByParentId(realm, this.getId());
+            }
+        };
     }
 
     @Override
@@ -160,13 +165,6 @@ public Long getGroupsCount(RealmModel realm, Boolean onlyTopGroups) {
         return storeWithRealm(realm).getCount(withCriteria(mcb));
     }
 
-    @Override
-    public Long getSubGroupsCount(RealmModel realm, String parentId) {
-        DefaultModelCriteria<GroupModel> mcb = criteria();
-        mcb = mcb.compare(SearchableFields.REALM_ID, Operator.EQ, realm.getId()).compare(SearchableFields.PARENT_ID, Operator.EQ, parentId);
-        return storeWithRealm(realm).getCount(withCriteria(mcb));
-    }
-
     @Override
     public Long getGroupsCountByNameContaining(RealmModel realm, String search) {
         LOG.tracef("getGroupsCountByNameContaining(%s, %s, %s)%s", realm, session, search, getShortStackTrace());
@@ -225,22 +223,6 @@ public Stream<GroupModel> searchForGroupByNameStream(RealmModel realm, String se
             .map(entityToAdapterFunc(realm));
     }
 
-    @Override
-    public Stream<GroupModel> searchForSubgroupsByParentIdNameStream(RealmModel realm, String id, String search, Boolean exact, Integer firstResult, Integer maxResults) {
-        DefaultModelCriteria<GroupModel> mcb = criteria();
-        mcb = mcb.compare(SearchableFields.REALM_ID, Operator.EQ, realm.getId())
-            .compare(SearchableFields.PARENT_ID, Operator.EQ, id);
-
-        if(Boolean.TRUE.equals(exact)) {
-            mcb.compare(SearchableFields.NAME, Operator.EQ, search);
-        } else {
-            mcb.compare(SearchableFields.NAME, Operator.ILIKE, "%" + search + "%");
-        }
-
-        return storeWithRealm(realm).read(withCriteria(mcb).pagination(firstResult, maxResults, SearchableFields.NAME))
-            .map(entityToAdapterFunc(realm));
-    }
-
     @Override
     public Stream<GroupModel> searchGroupsByAttributes(RealmModel realm, Map<String, String> attributes, Integer firstResult, Integer maxResults) {
         DefaultModelCriteria<GroupModel> mcb = criteria();

diff --git a/model/map/src/main/java/org/keycloak/models/map/group/MapGroupProviderFactory.java b/model/map/src/main/java/org/keycloak/models/map/group/MapGroupProviderFactory.java
@@ -64,9 +64,9 @@ public void invalidate(KeycloakSession session, InvalidableObjectType type, Obje
 
             // TODO: Should the batch size be a config option?
             // batch and remove subgroups to avoid grinding server to a halt at scale
-            long batches = (long) Math.ceil(create(session).getSubGroupsCount(realm, group.getId()) / 1000.0);
+            long batches = (long) Math.ceil(group.getSubGroupsCount() / 1000.0);
             for(int i = 0; i < batches; i++) {
-                create(session).searchForSubgroupsByParentIdStream(realm, group.getId(), i * 1000, 1000)
+                group.getSubGroupsStream(i * 1000, 1000)
                     .forEach(subGroup -> create(session).removeGroup(realm, subGroup));
             }
         } else if (type == GROUP_AFTER_REMOVE) {