From ee20ab886b553b56ccabe4e296cb86a028de5e5b Mon Sep 17 00:00:00 2001 From: Pedro Ruivo Date: Thu, 8 Jan 2026 08:56:20 +0000 Subject: [PATCH] Admin UI: slow response time listing second user page Fixes #44860 Signed-off-by: Pedro Ruivo <1492066+pruivo@users.noreply.github.com> Signed-off-by: Alexander Schwartz Co-authored-by: Pedro Ruivo <1492066+pruivo@users.noreply.github.com> Co-authored-by: Alexander Schwartz --- .../keycloak/models/jpa/JpaUserProvider.java | 93 ++++---- .../keycloak/storage/UserStorageManager.java | 225 ++++++++++-------- .../model/user/UserPaginationTest.java | 5 +- 3 files changed, 170 insertions(+), 153 deletions(-) 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 c1bfdb5fc6b..1abfb907d2b 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 @@ -33,7 +33,6 @@ import jakarta.persistence.LockModeType; import jakarta.persistence.TypedQuery; import jakarta.persistence.criteria.CriteriaBuilder; import jakarta.persistence.criteria.CriteriaQuery; -import jakarta.persistence.criteria.Expression; import jakarta.persistence.criteria.From; import jakarta.persistence.criteria.Join; import jakarta.persistence.criteria.JoinType; @@ -98,7 +97,7 @@ public class JpaUserProvider implements UserProvider, UserCredentialStore, JpaUs private static final char ESCAPE_BACKSLASH = '\\'; private final KeycloakSession session; - protected EntityManager em; + protected final EntityManager em; private final JpaUserCredentialStore credentialStore; public JpaUserProvider(KeycloakSession session, EntityManager em) { @@ -156,7 +155,6 @@ public class JpaUserProvider implements UserProvider, UserCredentialStore, JpaUs } private void removeUser(UserEntity user) { - String id = user.getId(); em.createNamedQuery("deleteUserRoleMappingsByUser").setParameter("user", user).executeUpdate(); em.createNamedQuery("deleteUserGroupMembershipsByUser").setParameter("user", user).executeUpdate(); em.createNamedQuery("deleteUserConsentClientScopesByUser").setParameter("user", user).executeUpdate(); @@ -183,6 +181,9 @@ public class JpaUserProvider implements UserProvider, UserCredentialStore, JpaUs @Override public void updateFederatedIdentity(RealmModel realm, UserModel federatedUser, FederatedIdentityModel federatedIdentityModel) { FederatedIdentityEntity federatedIdentity = findFederatedIdentity(federatedUser, federatedIdentityModel.getIdentityProvider(), LockModeType.PESSIMISTIC_WRITE); + if (federatedIdentity == null) { + return; + } federatedIdentity.setUserName(federatedIdentityModel.getUserName()); federatedIdentity.setToken(federatedIdentityModel.getToken()); @@ -378,7 +379,7 @@ public class JpaUserProvider implements UserProvider, UserCredentialStore, JpaUs @Override public void grantToAllUsers(RealmModel realm, RoleModel role) { - if (realm.equals(role.isClientRole() ? ((ClientModel)role.getContainer()).getRealm() : (RealmModel)role.getContainer())) { + if (realm.equals(role.isClientRole() ? ((ClientModel)role.getContainer()).getRealm() : role.getContainer())) { em.createNamedQuery("grantRoleToAllUsers") .setParameter("realmId", realm.getId()) .setParameter("roleId", role.getId()) @@ -596,7 +597,7 @@ public class JpaUserProvider implements UserProvider, UserCredentialStore, JpaUs } predicates.addAll(AdminPermissionsSchema.SCHEMA.applyAuthorizationFilters(session, AdminPermissionsSchema.USERS, this, realm, cb, query, root)); - query.select(cb.count(root)).where(predicates.toArray(Predicate[]::new)); + query.select(cb.count(root)).where(predicates); return em.createQuery(query).getSingleResult().intValue(); } @@ -621,21 +622,13 @@ public class JpaUserProvider implements UserProvider, UserCredentialStore, JpaUs CriteriaQuery queryBuilder = builder.createQuery(Long.class); Root root = queryBuilder.from(UserEntity.class); - queryBuilder.select(builder.countDistinct(root)); - List predicates = new ArrayList<>(); predicates.add(builder.equal(root.get("realmId"), realm.getId())); - - for (String stringToSearch : search.trim().split("\\s+")) { - predicates.add(builder.or(getSearchOptionPredicateArray(stringToSearch, builder, root))); - } - + addSearchPredicates(search, builder, root, predicates); predicates.addAll(AdminPermissionsSchema.SCHEMA.applyAuthorizationFilters(session, AdminPermissionsSchema.USERS, this, realm, builder, queryBuilder, root)); - queryBuilder.where(predicates.toArray(Predicate[]::new)); - - return em.createQuery(queryBuilder).getSingleResult().intValue(); + return em.createQuery(countQuery(queryBuilder, builder, root, predicates)).getSingleResult().intValue(); } @Override @@ -650,21 +643,13 @@ public class JpaUserProvider implements UserProvider, UserCredentialStore, JpaUs Root groupMembership = queryBuilder.from(UserGroupMembershipEntity.class); Join userJoin = groupMembership.join("user"); - queryBuilder.select(builder.countDistinct(userJoin)); - List predicates = new ArrayList<>(); predicates.add(builder.equal(userJoin.get("realmId"), realm.getId())); - - for (String stringToSearch : search.trim().split("\\s+")) { - predicates.add(builder.or(getSearchOptionPredicateArray(stringToSearch, builder, userJoin))); - } - + addSearchPredicates(search, builder, userJoin, predicates); predicates.add(groupMembership.get("groupId").in(groupIds)); - queryBuilder.where(predicates.toArray(Predicate[]::new)); - - return em.createQuery(queryBuilder).getSingleResult().intValue(); + return em.createQuery(countQuery(queryBuilder, builder, userJoin, predicates)).getSingleResult().intValue(); } @Override @@ -672,15 +657,12 @@ public class JpaUserProvider implements UserProvider, UserCredentialStore, JpaUs CriteriaBuilder qb = em.getCriteriaBuilder(); CriteriaQuery userQuery = qb.createQuery(Long.class); Root from = userQuery.from(UserEntity.class); - Expression count = qb.countDistinct(from); - userQuery = userQuery.select(count); List restrictions = predicates(params, from, Map.of()); restrictions.add(qb.equal(from.get("realmId"), realm.getId())); restrictions.addAll(AdminPermissionsSchema.SCHEMA.applyAuthorizationFilters(session, AdminPermissionsSchema.USERS, this, realm, qb, userQuery, from)); - userQuery = userQuery.where(restrictions.toArray(Predicate[]::new)); - TypedQuery query = em.createQuery(userQuery); + TypedQuery query = em.createQuery(countQuery(userQuery, qb, from, restrictions)); Long result = query.getSingleResult(); return result.intValue(); @@ -706,7 +688,7 @@ public class JpaUserProvider implements UserProvider, UserCredentialStore, JpaUs restrictions.addAll(AdminPermissionsSchema.SCHEMA.applyAuthorizationFilters(session, AdminPermissionsSchema.USERS, this, realm, cb, countQuery, root)); - countQuery.where(restrictions.toArray(Predicate[]::new)); + countQuery.where(restrictions); TypedQuery query = em.createQuery(countQuery); Long result = query.getSingleResult(); @@ -728,7 +710,7 @@ public class JpaUserProvider implements UserProvider, UserCredentialStore, JpaUs predicates.addAll(AdminPermissionsSchema.SCHEMA.applyAuthorizationFilters(session, AdminPermissionsSchema.USERS, this, realm, builder, queryBuilder, userPath)); - queryBuilder.where(predicates.toArray(Predicate[]::new)).orderBy(builder.asc(userPath.get(UserModel.USERNAME))); + queryBuilder.where(predicates).orderBy(builder.asc(userPath.get(UserModel.USERNAME))); return closing(paginateQuery(em.createQuery(queryBuilder), firstResult, maxResults).getResultStream().map(user -> new UserAdapter(session, realm, em, user))); } @@ -772,7 +754,7 @@ public class JpaUserProvider implements UserProvider, UserCredentialStore, JpaUs predicates.addAll(AdminPermissionsSchema.SCHEMA.applyAuthorizationFilters(session, AdminPermissionsSchema.USERS, this, realm, builder, queryBuilder, userPath)); - queryBuilder.where(predicates.toArray(Predicate[]::new)).orderBy(builder.asc(userPath.get(UserModel.USERNAME))); + queryBuilder.where(predicates).orderBy(builder.asc(userPath.get(UserModel.USERNAME))); return closing(paginateQuery(em.createQuery(queryBuilder), first, max).getResultStream().map(user -> new UserAdapter(session, realm, em, user))); } @@ -790,7 +772,7 @@ public class JpaUserProvider implements UserProvider, UserCredentialStore, JpaUs predicates.addAll(AdminPermissionsSchema.SCHEMA.applyAuthorizationFilters(session, AdminPermissionsSchema.USERS, this, realm, cb, cq, user)); cq.select(user) - .where(predicates.toArray(Predicate[]::new)) + .where(predicates) .orderBy(cb.asc(user.get("username"))); TypedQuery query = em.createQuery(cq); @@ -824,7 +806,7 @@ public class JpaUserProvider implements UserProvider, UserCredentialStore, JpaUs predicates.addAll(AdminPermissionsSchema.SCHEMA.applyAuthorizationFilters(session, AdminPermissionsSchema.USERS, this, realm, builder, queryBuilder, root)); - queryBuilder.distinct(true).where(predicates.toArray(Predicate[]::new)).orderBy(builder.asc(root.get(UserModel.USERNAME))); + queryBuilder.distinct(true).where(predicates).orderBy(builder.asc(root.get(UserModel.USERNAME))); TypedQuery query = em.createQuery(queryBuilder); @@ -992,7 +974,7 @@ public class JpaUserProvider implements UserProvider, UserCredentialStore, JpaUs } } - private Predicate[] getSearchOptionPredicateArray(String value, CriteriaBuilder builder, From from) { + private static Predicate getSearchOptionPredicate(String value, CriteriaBuilder builder, From from) { value = value.toLowerCase(); List orPredicates = new ArrayList<>(); @@ -1016,7 +998,7 @@ public class JpaUserProvider implements UserProvider, UserCredentialStore, JpaUs orPredicates.add(builder.like(builder.lower(from.get(LAST_NAME)), value, ESCAPE_BACKSLASH)); } - return orPredicates.toArray(Predicate[]::new); + return builder.or(orPredicates); } private UserEntity userInEntityManagerContext(String id) { @@ -1042,9 +1024,7 @@ public class JpaUserProvider implements UserProvider, UserCredentialStore, JpaUs switch (key) { case UserModel.SEARCH: - for (String stringToSearch : value.trim().split("\\s+")) { - predicates.add(builder.or(getSearchOptionPredicateArray(stringToSearch, builder, root))); - } + addSearchPredicates(value, builder, root, predicates); break; case FIRST_NAME: case LAST_NAME: @@ -1082,8 +1062,15 @@ public class JpaUserProvider implements UserProvider, UserCredentialStore, JpaUs break; case UserModel.EXACT: break; - // All unknown attributes will be assumed as custom attributes + case UserModel.INCLUDE_SERVICE_ACCOUNT: { + if (!attributes.containsKey(UserModel.INCLUDE_SERVICE_ACCOUNT) + || !Boolean.parseBoolean(attributes.get(UserModel.INCLUDE_SERVICE_ACCOUNT))) { + predicates.add(root.get("serviceAccountClientLink").isNull()); + } + break; + } default: + // All unknown attributes will be assumed as custom attributes Join attributesJoin = root.join("attributes", JoinType.LEFT); if (value.length() > 255) { customLongValueSearchAttributes.put(key, value); @@ -1102,18 +1089,11 @@ public class JpaUserProvider implements UserProvider, UserCredentialStore, JpaUs } } break; - case UserModel.INCLUDE_SERVICE_ACCOUNT: { - if (!attributes.containsKey(UserModel.INCLUDE_SERVICE_ACCOUNT) - || !Boolean.parseBoolean(attributes.get(UserModel.INCLUDE_SERVICE_ACCOUNT))) { - predicates.add(root.get("serviceAccountClientLink").isNull()); - } - break; - } } } if (!attributePredicates.isEmpty()) { - predicates.add(builder.and(attributePredicates.toArray(Predicate[]::new))); + predicates.add(builder.and(attributePredicates)); } return predicates; @@ -1128,4 +1108,21 @@ public class JpaUserProvider implements UserProvider, UserCredentialStore, JpaUs public EntityManager getEntityManager() { return em; } + + private static void addSearchPredicates(String search, CriteriaBuilder builder, From from, List predicates) { + if (search == null) { + return; + } + for (String stringToSearch : search.trim().split("\\s+")) { + predicates.add(getSearchOptionPredicate(stringToSearch, builder, from)); + } + } + + private static CriteriaQuery countQuery(CriteriaQuery query, CriteriaBuilder builder, From from, List predicates) { + // When joining multiple tables, issuing a "distinct" is required to get the correct result. + // At the same time it is more expensive than a regular count as it would need to sort/keep all keys in memory at the database for the duration of the query. + // Therefore, we use a standard count where possible. + return query.select(from.getJoins().isEmpty() ? builder.count(from) : builder.countDistinct(from)) + .where(predicates); + } } diff --git a/model/storage-private/src/main/java/org/keycloak/storage/UserStorageManager.java b/model/storage-private/src/main/java/org/keycloak/storage/UserStorageManager.java index d8bd36e44d3..da2ca0819ad 100755 --- a/model/storage-private/src/main/java/org/keycloak/storage/UserStorageManager.java +++ b/model/storage-private/src/main/java/org/keycloak/storage/UserStorageManager.java @@ -23,11 +23,9 @@ import java.util.Map; import java.util.Objects; import java.util.Optional; import java.util.Set; -import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Function; import java.util.function.Predicate; -import java.util.stream.Collectors; import java.util.stream.Stream; import org.keycloak.common.Profile; @@ -111,10 +109,6 @@ public class UserStorageManager extends AbstractStorageManager Stream getCredentialProviders(KeycloakSession session, Class type) { return session.getKeycloakSessionFactory().getProviderFactoriesStream(CredentialProvider.class) .filter(f -> Types.supports(type, f, CredentialProviderFactory.class)) - .map(f -> (T) session.getProvider(CredentialProvider.class, f.getId())); + .map(f -> session.getProvider(CredentialProvider.class, f.getId())) + .map(type::cast); } @Override @@ -232,7 +227,7 @@ public class UserStorageManager extends AbstractStorageManager credentialAuthentication.supportsCredentialAuthenticationFor(input.getType())) - .collect(Collectors.toList())) { + .toList()) { CredentialValidationOutput validationOutput = session.getProvider(TracingProvider.class).trace(credentialAuthentication.getClass(), "authenticate", span -> { CredentialValidationOutput output = credentialAuthentication.authenticate(realm, input); @@ -302,109 +297,133 @@ public class UserStorageManager extends AbstractStorageManager query(Object provider, Integer firstResult, Integer maxResults); } @FunctionalInterface - interface CountQuery { + protected interface CountQuery { int query(Object provider, Integer firstResult, Integer maxResult); } protected Stream query(PaginatedQuery pagedQuery, RealmModel realm, Integer firstResult, Integer maxResults) { - return query(pagedQuery, ((provider, first, max) -> (int) pagedQuery.query(provider, first, max).count()), realm, firstResult, maxResults); + return query(pagedQuery, ((provider, first, max) -> (int) pagedQuery.query(provider, first, max).count()), realm, firstResult, maxResults, true); } - protected Stream query(PaginatedQuery pagedQuery, CountQuery countQuery, RealmModel realm, Integer firstResult, Integer maxResults) { - if (maxResults != null && maxResults == 0) return Stream.empty(); - - Stream providersStream = Stream.concat(Stream.of((Object) localStorage()), getEnabledStorageProviders(realm, UserQueryMethodsProvider.class)); - - UserFederatedStorageProvider federatedStorageProvider = getFederatedStorage(); - if (federatedStorageProvider != null) { - providersStream = Stream.concat(providersStream, Stream.of(federatedStorageProvider)); + protected Stream query(PaginatedQuery pagedQuery, CountQuery countQuery, RealmModel realm, Integer firstResult, Integer maxResults, boolean requiresFederatedStorage) { + if (maxResults != null && maxResults == 0) { + return Stream.empty(); } - final AtomicInteger currentFirst; - final AtomicBoolean needsAdditionalFirstResultFiltering = new AtomicBoolean(false); + var storageProviders = getEnabledStorageProviders(realm, UserQueryMethodsProvider.class).toList(); - if (firstResult == null || firstResult <= 0) { // We don't want to skip any users so we don't need to do firstResult filtering - currentFirst = new AtomicInteger(0); - } else { - // This is an optimization using count query to skip querying users if we can use count method to determine how many users can be provided by each provider - AtomicBoolean droppingProviders = new AtomicBoolean(true); - currentFirst = new AtomicInteger(firstResult); - - providersStream = providersStream - .filter(provider -> { // This is basically dropWhile - if (!droppingProviders.get()) return true; // We have already gathered enough users to pass firstResult number in previous providers, we can take all following providers - - if (!(provider instanceof UserCountMethodsProvider)) { - logger.tracef("We encountered a provider (%s) that does not implement count queries therefore we can't say how many users it can provide.", provider.getClass().getSimpleName()); - // for this reason we need to start querying this provider and all following providers - droppingProviders.set(false); - needsAdditionalFirstResultFiltering.set(true); - return true; // don't filter out this provider because we are unable to say how many users it can provide - } - - long expectedNumberOfUsersForProvider = countQueryWithGracefulDegradation(provider, countQuery, 0, currentFirst.get() + 1); // check how many users we can obtain from this provider - logger.tracef("This provider (%s) is able to return %d users.", provider.getClass().getSimpleName(), expectedNumberOfUsersForProvider); - - if (expectedNumberOfUsersForProvider == currentFirst.get()) { // This provider provides exactly the amount of users we need for passing firstResult, we can set currentFirst to 0 and drop this provider - currentFirst.set(0); - droppingProviders.set(false); - return false; - } - - if (expectedNumberOfUsersForProvider > currentFirst.get()) { // If we can obtain enough enough users from this provider to fulfill our need we can stop dropping providers - droppingProviders.set(false); - return true; // don't filter out this provider because we are going to return some users from it - } - - logger.tracef("This provider (%s) cannot provide enough users to pass firstResult so we are going to filter it out and change " - + "firstResult for next provider: %d - %d = %d", provider.getClass().getSimpleName(), - currentFirst.get(), expectedNumberOfUsersForProvider, currentFirst.get() - expectedNumberOfUsersForProvider); - currentFirst.set((int) (currentFirst.get() - expectedNumberOfUsersForProvider)); - return false; - }) - // collecting stream of providers to ensure the filtering (above) is evaluated before we move forward to actual querying - .collect(Collectors.toList()).stream(); + if (firstResult == null || firstResult <= 0) { + // we don't have a first result set, so we start from the beginning and go through all providers. + var providers = Stream.concat(Stream.of(localStorage()), concatExternalWithFederated(storageProviders, 0, requiresFederatedStorage)); + return queryProviders(providers, pagedQuery, 0, maxResults, false); } - if (needsAdditionalFirstResultFiltering.get() && currentFirst.get() > 0) { - logger.tracef("In the providerStream there is a provider that does not support count queries and we need to skip some users."); - // we need to make sure, we skip firstResult users from this or the following providers - if (maxResults == null || maxResults < 0) { - return paginatedStream(providersStream - .flatMap(provider -> queryWithGracefulDegradation(provider, pagedQuery, null, null)), currentFirst.get(), null); - } else { - final AtomicInteger currentMax = new AtomicInteger(currentFirst.get() + maxResults); - - return paginatedStream(providersStream - .flatMap(provider -> queryWithGracefulDegradation(provider, pagedQuery, null, currentMax.get())) - .peek(userModel -> { - currentMax.updateAndGet(i -> i > 0 ? i - 1 : i); - }), currentFirst.get(), maxResults); + if (storageProviders.isEmpty()) { + if (requiresFederatedStorage) { + // we need to count the database. + return queryLocalAndFederatedStorage(pagedQuery, countQuery, firstResult, maxResults); } + // fast path, stream from the database only without counting anything. + return queryLocalStorage(pagedQuery, firstResult, maxResults); } - // Actual user querying + var localStorage = localStorage(); + var count = countQueryWithGracefulDegradation(localStorage, countQuery, 0, firstResult + 1); + + if (count > firstResult) { + // we need some users from the database + var providers = Stream.concat(Stream.of(localStorage()), concatExternalWithFederated(storageProviders, 0, requiresFederatedStorage)); + return queryProviders(providers, pagedQuery, firstResult, maxResults, false); + } + + if (count == firstResult) { + // we have the exact users in the database + // we don't need to count anything else, and we start from the first external storage provider. + var providers = concatExternalWithFederated(storageProviders, 0, requiresFederatedStorage); + return queryProviders(providers, pagedQuery, 0, maxResults, false); + } + + // we need to count the external providers + firstResult -= count; + int lastProviderToCount = requiresFederatedStorage ? + storageProviders.size(): + storageProviders.size() - 1; + int startIdx = 0; + + for (; startIdx < lastProviderToCount; ++startIdx) { + var provider = storageProviders.get(startIdx); + if (!(provider instanceof UserCountMethodsProvider)) { + assert firstResult > 0; + Stream providers = concatExternalWithFederated(storageProviders, startIdx, requiresFederatedStorage); + return queryProviders(providers, pagedQuery, firstResult, maxResults, true); + } + count = countQueryWithGracefulDegradation(storageProviders.get(startIdx), countQuery, 0, firstResult + 1); + if (count > firstResult) { + // we start on this provider as we need some users from it. + break; + } + if (count == firstResult) { + // exact users required to offset, start querying the next provider. + startIdx++; + firstResult = 0; + break; + } + firstResult -= count; + } + + var providers = concatExternalWithFederated(storageProviders, startIdx, requiresFederatedStorage); + return queryProviders(providers, pagedQuery, firstResult, maxResults, false); + } + + private Stream queryLocalAndFederatedStorage(PaginatedQuery pagedQuery, CountQuery countQuery, int firstResult, Integer maxResults) { + assert firstResult > 0; + var localStorage = localStorage(); + // check how many users we can obtain from the local provider + var count = countQueryWithGracefulDegradation(localStorage, countQuery, 0, firstResult + 1); + + if (count <= firstResult) { + // local provider does not have enough user to skip the first users, querying the federated provider only. + return queryProviders(Stream.of(getFederatedStorage()), pagedQuery, firstResult - count, maxResults, false); + } + + return queryProviders(Stream.of(localStorage, getFederatedStorage()), pagedQuery, firstResult, maxResults, false); + } + + private Stream queryLocalStorage(PaginatedQuery pagedQuery, int firstResult, Integer maxResults) { + assert firstResult > 0; + return queryWithGracefulDegradation(localStorage(), pagedQuery, firstResult, maxResults); + } + + private static Stream queryProviders(Stream providersStream, PaginatedQuery pagedQuery, int offset, Integer maxResults, boolean useStreamSkip) { + var firstResults = useStreamSkip ? new AtomicInteger(0) : new AtomicInteger(offset); if (maxResults == null || maxResults < 0) { - // No maxResult set, we want all users - return providersStream - .flatMap(provider -> queryWithGracefulDegradation(provider, pagedQuery, currentFirst.getAndSet(0), null)); - } else { - final AtomicInteger currentMax = new AtomicInteger(maxResults); - - // Query users with currentMax variable counting how many users we return - return providersStream - .filter(provider -> currentMax.get() != 0) // If we reach currentMax == 0, we can skip querying all following providers - .flatMap(provider -> queryWithGracefulDegradation(provider, pagedQuery, currentFirst.getAndSet(0), currentMax.get())) - .peek(userModel -> { - currentMax.updateAndGet(i -> i > 0 ? i - 1 : i); - }); + var users = providersStream + .flatMap(provider -> queryWithGracefulDegradation(provider, pagedQuery, firstResults.getAndSet(0), null)); + return useStreamSkip ? users.skip(offset) : users; } + var currentMax = useStreamSkip ? + new AtomicInteger(offset + maxResults) : + new AtomicInteger(maxResults); + + // Query users with currentMax variable counting how many users we return + var users = providersStream + .filter(provider -> currentMax.get() != 0) // If we reach currentMax == 0, we can skip querying all following providers + .flatMap(provider -> queryWithGracefulDegradation(provider, pagedQuery, firstResults.getAndSet(0), currentMax.get())) + .peek(userModel -> currentMax.updateAndGet(i -> i > 0 ? i - 1 : i)); + return useStreamSkip ? users.skip(offset).limit(maxResults) : users.limit(maxResults); + } + + private Stream concatExternalWithFederated(List externalStorage, int startIdx, boolean requiresFederatedStorage) { + Stream providers = externalStorage.subList(startIdx, externalStorage.size()).stream(); + return requiresFederatedStorage ? + Stream.concat(providers, Stream.of(getFederatedStorage())) : + providers; } /** @@ -412,7 +431,7 @@ public class UserStorageManager extends AbstractStorageManager queryWithGracefulDegradation(Object provider, PaginatedQuery pagedQuery, + private static Stream queryWithGracefulDegradation(Object provider, PaginatedQuery pagedQuery, Integer firstResult, Integer maxResults) { try { return pagedQuery.query(provider, firstResult, maxResults); @@ -565,7 +584,7 @@ public class UserStorageManager extends AbstractStorageManager getFederatedIdentitiesStream(RealmModel realm, UserModel user) { if (user == null) throw new IllegalStateException("Federated user no longer valid"); - Stream stream = StorageId.isLocalStorage(user) ? + Stream stream = StorageId.isLocalStorage(user.getId()) ? localStorage().getFederatedIdentitiesStream(realm, user) : Stream.empty(); if (getFederatedStorage() != null) stream = Stream.concat(stream, getFederatedStorage().getFederatedIdentitiesStream(user.getId(), realm)); @@ -958,7 +977,7 @@ public class UserStorageManager extends AbstractStorageManager factory = ComponentUtil.getComponentFactory(session, model); if (!(factory instanceof UserStorageProviderFactory)) return; // enlistAfterCompletion(..) as we need to ensure that the realm is available in the system @@ -1051,12 +1070,12 @@ public class UserStorageManager extends AbstractStorageManager