From 33331788a457c4f7674a24f586c4a2f7ba783a25 Mon Sep 17 00:00:00 2001 From: Martin Kanis Date: Tue, 4 Jun 2024 14:49:35 +0200 Subject: [PATCH] Introduce count method to avoid fetching all organization upon checking for existence Closes #29697 Signed-off-by: Martin Kanis --- .../models/jpa/entities/OrganizationEntity.java | 3 ++- .../organization/jpa/JpaOrganizationProvider.java | 12 +++++++++--- .../organization/OrganizationProvider.java | 10 ++++++++-- .../keycloak/organization/utils/Organizations.java | 5 ++--- .../organization/admin/OrganizationTest.java | 14 ++++++++++++++ 5 files changed, 35 insertions(+), 9 deletions(-) diff --git a/model/jpa/src/main/java/org/keycloak/models/jpa/entities/OrganizationEntity.java b/model/jpa/src/main/java/org/keycloak/models/jpa/entities/OrganizationEntity.java index 3ce9cd0288..e4a1638e06 100644 --- a/model/jpa/src/main/java/org/keycloak/models/jpa/entities/OrganizationEntity.java +++ b/model/jpa/src/main/java/org/keycloak/models/jpa/entities/OrganizationEntity.java @@ -42,7 +42,8 @@ import jakarta.persistence.Table; @NamedQuery(name="getByNameOrDomain", query="select distinct o from OrganizationEntity o inner join OrganizationDomainEntity d ON o.id = d.organization.id" + " where o.realmId = :realmId AND (o.name = :search OR d.name = :search) order by o.name ASC"), @NamedQuery(name="getByNameOrDomainContained", query="select distinct o from OrganizationEntity o inner join OrganizationDomainEntity d ON o.id = d.organization.id" + - " where o.realmId = :realmId AND (lower(o.name) like concat('%',:search,'%') OR d.name like concat('%',:search,'%')) order by o.name ASC") + " where o.realmId = :realmId AND (lower(o.name) like concat('%',:search,'%') OR d.name like concat('%',:search,'%')) order by o.name ASC"), + @NamedQuery(name="getCount", query="select count(o) from OrganizationEntity o where o.realmId = :realmId") }) public class OrganizationEntity { diff --git a/model/jpa/src/main/java/org/keycloak/organization/jpa/JpaOrganizationProvider.java b/model/jpa/src/main/java/org/keycloak/organization/jpa/JpaOrganizationProvider.java index 070ba2c1db..8aa96e48da 100644 --- a/model/jpa/src/main/java/org/keycloak/organization/jpa/JpaOrganizationProvider.java +++ b/model/jpa/src/main/java/org/keycloak/organization/jpa/JpaOrganizationProvider.java @@ -26,8 +26,6 @@ import static org.keycloak.utils.StreamsUtil.closing; import java.util.ArrayList; import java.util.List; import java.util.Map; -import java.util.Set; -import java.util.stream.Collectors; import java.util.stream.Stream; import jakarta.persistence.EntityManager; @@ -47,7 +45,6 @@ import org.keycloak.models.KeycloakSession; import org.keycloak.models.ModelDuplicateException; import org.keycloak.models.ModelException; import org.keycloak.models.ModelValidationException; -import org.keycloak.models.OrganizationDomainModel; import org.keycloak.models.OrganizationModel; import org.keycloak.models.RealmModel; import org.keycloak.models.UserModel; @@ -394,6 +391,15 @@ public class JpaOrganizationProvider implements OrganizationProvider { return true; } + @Override + public long count() { + TypedQuery query; + query = em.createNamedQuery("getCount", Long.class); + query.setParameter("realmId", realm.getId()); + + return query.getSingleResult(); + } + @Override public boolean isEnabled() { return realm.isOrganizationsEnabled(); diff --git a/server-spi-private/src/main/java/org/keycloak/organization/OrganizationProvider.java b/server-spi-private/src/main/java/org/keycloak/organization/OrganizationProvider.java index 94e27025b5..7fd84d889e 100644 --- a/server-spi-private/src/main/java/org/keycloak/organization/OrganizationProvider.java +++ b/server-spi-private/src/main/java/org/keycloak/organization/OrganizationProvider.java @@ -138,7 +138,7 @@ public interface OrganizationProvider extends Provider { /** * Associate the given {@link IdentityProviderModel} with the given {@link OrganizationModel}. - * + * * @param organization the organization * @param identityProvider the identityProvider * @return {@code true} if the identityProvider was associated with the organization. Otherwise, returns {@code false} @@ -153,7 +153,7 @@ public interface OrganizationProvider extends Provider { /** * Removes the link between the given {@link OrganizationModel} and identity provider associated with it if such a link exists. - * + * * @param organization the organization * @param identityProvider the identity provider * @return {@code true} if the link was removed, {@code false} otherwise @@ -197,4 +197,10 @@ public interface OrganizationProvider extends Provider { * @return {@code true} if the given {@code member} is a member and was successfully removed from the organization. Otherwise, returns {@code false} */ boolean removeMember(OrganizationModel organization, UserModel member); + + /** + * Returns number of organizations in the realm. + * @return long Number of organizations + */ + long count(); } diff --git a/services/src/main/java/org/keycloak/organization/utils/Organizations.java b/services/src/main/java/org/keycloak/organization/utils/Organizations.java index efe7b9bdf0..655f352922 100644 --- a/services/src/main/java/org/keycloak/organization/utils/Organizations.java +++ b/services/src/main/java/org/keycloak/organization/utils/Organizations.java @@ -130,9 +130,8 @@ public class Organizations { }; } - public static boolean isEnabledAndOrganizationsPresent(OrganizationProvider provider) { - // todo replace getAllStream().findAny().isPresent() with count query - return provider != null && provider.isEnabled() && provider.getAllStream().findAny().isPresent(); + public static boolean isEnabledAndOrganizationsPresent(OrganizationProvider orgProvider) { + return orgProvider != null && orgProvider.isEnabled() && orgProvider.count() != 0; } public static void checkEnabled(OrganizationProvider provider) { diff --git a/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/organization/admin/OrganizationTest.java b/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/organization/admin/OrganizationTest.java index e4fc98b317..6768d72db9 100755 --- a/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/organization/admin/OrganizationTest.java +++ b/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/organization/admin/OrganizationTest.java @@ -47,11 +47,13 @@ import org.keycloak.admin.client.resource.OrganizationResource; import org.keycloak.admin.client.resource.RealmResource; import org.keycloak.common.Profile.Feature; import org.keycloak.models.utils.KeycloakModelUtils; +import org.keycloak.organization.OrganizationProvider; import org.keycloak.representations.idm.IdentityProviderRepresentation; import org.keycloak.representations.idm.OrganizationDomainRepresentation; import org.keycloak.representations.idm.OrganizationRepresentation; import org.keycloak.representations.idm.RealmRepresentation; import org.keycloak.testsuite.arquillian.annotation.EnableFeature; +import org.keycloak.testsuite.runonserver.RunOnServer; import org.keycloak.testsuite.updaters.RealmAttributeUpdater; import org.keycloak.testsuite.util.RealmBuilder; @@ -447,4 +449,16 @@ public class OrganizationTest extends AbstractOrganizationTest { realmRes.remove(); } } + + @Test + public void testCount() { + for (int i = 0; i < 10; i++) { + createOrganization("kc.org." + i); + } + + getTestingClient().server(TEST_REALM_NAME).run((RunOnServer) session -> { + OrganizationProvider orgProvider = session.getProvider(OrganizationProvider.class); + assertEquals(10, orgProvider.count()); + }); + } }