fix: Optimize memberships fetching and refine search handling (#32)

Fixes #29
anarsultanov · Feb 22, 2024 · a2ef026 · a2ef026
1 parent 41bffce
commit a2ef026
Show file tree

Hide file tree

Showing 8 changed files with 79 additions and 36 deletions.
diff --git a/...multitenancy/authentication/authenticators/IdpTenantMembershipsCreatingAuthenticator.java b/...multitenancy/authentication/authenticators/IdpTenantMembershipsCreatingAuthenticator.java
@@ -3,7 +3,6 @@
 import dev.sultanov.keycloak.multitenancy.authentication.IdentityProviderTenantsConfig;
 import dev.sultanov.keycloak.multitenancy.model.TenantProvider;
 import dev.sultanov.keycloak.multitenancy.util.Constants;
-import jakarta.ws.rs.core.Response;
 import jakarta.ws.rs.core.Response.Status;
 import java.util.Optional;
 import java.util.Set;
@@ -59,7 +58,7 @@ private void doAuthenticate(AuthenticationFlowContext context, BrokeredIdentityC
                 if (tenantById.isEmpty()) {
                     log.warn("Tenant with ID %s, configured in IDP with alias %s, does not exist. Skipping membership creation."
                             .formatted(tenantId, brokerContext.getIdpConfig().getAlias()));
-                } else if (tenantById.get().getMembership(user).isPresent()) {
+                } else if (tenantById.get().getMembershipByUser(user).isPresent()) {
                     log.debug("User is already a member of tenant with ID %s. Skipping membership creation.".formatted(tenantId));
                 } else {
                     tenantById.get().grantMembership(user, Set.of(Constants.TENANT_USER_ROLE));

diff --git a/src/main/java/dev/sultanov/keycloak/multitenancy/model/TenantModel.java b/src/main/java/dev/sultanov/keycloak/multitenancy/model/TenantModel.java
@@ -20,18 +20,16 @@ public interface TenantModel {
 
     TenantMembershipModel grantMembership(UserModel user, Set<String> roles);
 
-    Stream<TenantMembershipModel> getMembershipsStream();
+    Stream<TenantMembershipModel> getMembershipsStream(Integer firstResult, Integer maxResults);
 
-    default Optional<TenantMembershipModel> getMembershipById(String membershipId) {
-        return getMembershipsStream().filter(membership -> membership.getId().equals(membershipId)).findFirst();
-    };
+    Stream<TenantMembershipModel> getMembershipsStream(String email, Integer firstResult, Integer maxResults);
 
-    default boolean hasMembership(UserModel user) {
-        return getMembershipsStream().anyMatch(membership -> membership.getUser().getId().equals(user.getId()));
-    }
+    Optional<TenantMembershipModel> getMembershipById(String membershipId);
 
-    default Optional<TenantMembershipModel> getMembership(UserModel user) {
-        return getMembershipsStream().filter(membership -> membership.getUser().getId().equals(user.getId())).findFirst();
+    Optional<TenantMembershipModel> getMembershipByUser(UserModel user);
+
+    default boolean hasMembership(UserModel user) {
+        return getMembershipByUser(user).isPresent();
     }
 
     boolean revokeMembership(String membershipId);

diff --git a/src/main/java/dev/sultanov/keycloak/multitenancy/model/entity/TenantMembershipEntity.java b/src/main/java/dev/sultanov/keycloak/multitenancy/model/entity/TenantMembershipEntity.java
@@ -8,6 +8,7 @@
 import jakarta.persistence.Id;
 import jakarta.persistence.JoinColumn;
 import jakarta.persistence.ManyToOne;
+import jakarta.persistence.NamedQueries;
 import jakarta.persistence.NamedQuery;
 import jakarta.persistence.OneToOne;
 import jakarta.persistence.Table;
@@ -19,8 +20,12 @@
 
 @Table(name = "TENANT_MEMBERSHIP", uniqueConstraints = {@UniqueConstraint(columnNames = {"TENANT_ID", "USER_ID"})})
 @Entity
-@NamedQuery(name = "getMembershipsByRealmAndUserId",
-        query = "SELECT m FROM TenantMembershipEntity m WHERE m.tenant.realmId = :realmId AND m.user.id = :userId")
+@NamedQueries({
+        @NamedQuery(name = "getMembershipsByRealmIdAndUserId", query = "SELECT m FROM TenantMembershipEntity m WHERE m.tenant.realmId = :realmId AND m.user.id = :userId"),
+        @NamedQuery(name = "getMembershipsByTenantId", query = "SELECT m FROM TenantMembershipEntity m WHERE m.tenant.id = :tenantId"),
+        @NamedQuery(name = "getMembershipsByTenantIdAndUserId", query = "SELECT m FROM TenantMembershipEntity m WHERE m.tenant.id = :tenantId AND m.user.id = :userId"),
+        @NamedQuery(name = "getMembershipsByTenantIdAndUserEmail", query = "SELECT m FROM TenantMembershipEntity m WHERE m.tenant.id = :tenantId AND m.user.email = :email")
+})
 public class TenantMembershipEntity {
 
     @Id

diff --git a/src/main/java/dev/sultanov/keycloak/multitenancy/model/jpa/JpaTenantProvider.java b/src/main/java/dev/sultanov/keycloak/multitenancy/model/jpa/JpaTenantProvider.java
@@ -82,7 +82,7 @@ public Stream<TenantInvitationModel> getTenantInvitationsStream(RealmModel realm
 
     @Override
     public Stream<TenantMembershipModel> getTenantMembershipsStream(RealmModel realm, UserModel user) {
-        TypedQuery<TenantMembershipEntity> query = em.createNamedQuery("getMembershipsByRealmAndUserId", TenantMembershipEntity.class);
+        TypedQuery<TenantMembershipEntity> query = em.createNamedQuery("getMembershipsByRealmIdAndUserId", TenantMembershipEntity.class);
         query.setParameter("realmId", realm.getId());
         query.setParameter("userId", user.getId());
         return query.getResultStream().map(m -> new TenantMembershipAdapter(session, realm, em, m));

diff --git a/src/main/java/dev/sultanov/keycloak/multitenancy/model/jpa/TenantAdapter.java b/src/main/java/dev/sultanov/keycloak/multitenancy/model/jpa/TenantAdapter.java
@@ -7,13 +7,16 @@
 import dev.sultanov.keycloak.multitenancy.model.entity.TenantInvitationEntity;
 import dev.sultanov.keycloak.multitenancy.model.entity.TenantMembershipEntity;
 import jakarta.persistence.EntityManager;
+import jakarta.persistence.TypedQuery;
 import java.util.HashSet;
+import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Stream;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.RealmModel;
 import org.keycloak.models.UserModel;
 import org.keycloak.models.jpa.JpaModel;
+import org.keycloak.models.jpa.PaginationUtils;
 import org.keycloak.models.jpa.entities.UserEntity;
 import org.keycloak.models.utils.KeycloakModelUtils;
 
@@ -47,32 +50,64 @@ public RealmModel getRealm() {
     }
 
     @Override
-    public TenantMembershipAdapter grantMembership(UserModel user, Set<String> roles) {
+    public TenantMembershipModel grantMembership(UserModel user, Set<String> roles) {
         TenantMembershipEntity entity = new TenantMembershipEntity();
         entity.setId(KeycloakModelUtils.generateId());
         entity.setUser(em.getReference(UserEntity.class, user.getId()));
         entity.setTenant(tenant);
         entity.setRoles(new HashSet<>(roles));
         em.persist(entity);
+        em.flush();
         tenant.getMemberships().add(entity);
         return new TenantMembershipAdapter(session, realm, em, entity);
     }
 
     @Override
-    public Stream<TenantMembershipModel> getMembershipsStream() {
-        return tenant.getMemberships().stream()
-                .map(membership -> new TenantMembershipAdapter(session, realm, em, membership));
+    public Stream<TenantMembershipModel> getMembershipsStream(Integer first, Integer max) {
+        TypedQuery<TenantMembershipEntity> query = em.createNamedQuery("getMembershipsByTenantId", TenantMembershipEntity.class);
+        query.setParameter("tenantId", tenant.getId());
+        return PaginationUtils.paginateQuery(query, first, max).getResultStream()
+                .map((membership) -> new TenantMembershipAdapter(session, realm, em, membership));
+    }
+
+    @Override
+    public Stream<TenantMembershipModel> getMembershipsStream(String email, Integer first, Integer max) {
+        TypedQuery<TenantMembershipEntity> query = em.createNamedQuery("getMembershipsByTenantIdAndUserEmail", TenantMembershipEntity.class);
+        query.setParameter("tenantId", tenant.getId());
+        query.setParameter("email", email);
+        return PaginationUtils.paginateQuery(query, first, max).getResultStream()
+                .map((membership) -> new TenantMembershipAdapter(session, realm, em, membership));
+    }
+
+    @Override
+    public Optional<TenantMembershipModel> getMembershipById(String membershipId) {
+        TenantMembershipEntity membership = em.find(TenantMembershipEntity.class, membershipId);
+        if (membership != null && realm.getId().equals(membership.getTenant().getRealmId())) {
+            return Optional.of(new TenantMembershipAdapter(session, realm, em, membership));
+        } else {
+            return Optional.empty();
+        }
+    }
+
+    @Override
+    public Optional<TenantMembershipModel> getMembershipByUser(UserModel user) {
+        TypedQuery<TenantMembershipEntity> query = em.createNamedQuery("getMembershipsByTenantIdAndUserId", TenantMembershipEntity.class);
+        query.setParameter("tenantId", tenant.getId());
+        query.setParameter("userId", user.getId());
+        return query.getResultStream().map(m -> (TenantMembershipModel) new TenantMembershipAdapter(session, realm, em, m)).findFirst();
     }
 
     @Override
     public boolean revokeMembership(String membershipId) {
-        var optionalMembership = getMembershipById(membershipId);
-        if (optionalMembership.isPresent()) {
-            var membershipEmail = optionalMembership.get().getUser().getEmail();
-            tenant.getMemberships().removeIf(entity -> entity.getId().equals(membershipId));
+        var membershipEntity = em.find(TenantMembershipEntity.class, membershipId);
+        if (membershipEntity != null) {
+            var membershipEmail = membershipEntity.getUser().getEmail();
+
             if (membershipEmail != null) {
                 revokeInvitations(membershipEmail);
             }
+            em.remove(membershipEntity);
+            em.flush();
             return true;
         }
         return false;
@@ -87,6 +122,7 @@ public TenantInvitationModel addInvitation(String email, UserModel inviter, Set<
         entity.setInvitedBy(inviter.getId());
         entity.setRoles(new HashSet<>(roles));
         em.persist(entity);
+        em.flush();
         tenant.getInvitations().add(entity);
         return new TenantInvitationAdapter(session, realm, em, entity);
     }

diff --git a/src/main/java/dev/sultanov/keycloak/multitenancy/resource/TenantAdminAuth.java b/src/main/java/dev/sultanov/keycloak/multitenancy/resource/TenantAdminAuth.java
@@ -1,7 +1,7 @@
 package dev.sultanov.keycloak.multitenancy.resource;
 
-import dev.sultanov.keycloak.multitenancy.util.Constants;
 import dev.sultanov.keycloak.multitenancy.model.TenantModel;
+import dev.sultanov.keycloak.multitenancy.util.Constants;
 import org.keycloak.models.ClientModel;
 import org.keycloak.models.RealmModel;
 import org.keycloak.models.UserModel;
@@ -15,10 +15,10 @@ public TenantAdminAuth(RealmModel realm, AccessToken token, UserModel user, Clie
     }
 
     boolean isTenantAdmin(TenantModel tenantModel) {
-        return tenantModel.getMembership(getUser()).filter(membership -> membership.getRoles().contains(Constants.TENANT_ADMIN_ROLE)).isPresent();
+        return tenantModel.getMembershipByUser(getUser()).filter(membership -> membership.getRoles().contains(Constants.TENANT_ADMIN_ROLE)).isPresent();
     }
 
     boolean isTenantMember(TenantModel tenantModel) {
-        return tenantModel.getMembership(getUser()).isPresent();
+        return tenantModel.getMembershipByUser(getUser()).isPresent();
     }
 }
diff --git a/src/main/java/dev/sultanov/keycloak/multitenancy/resource/TenantMembershipsResource.java b/src/main/java/dev/sultanov/keycloak/multitenancy/resource/TenantMembershipsResource.java
@@ -13,7 +13,8 @@
 import jakarta.ws.rs.QueryParam;
 import jakarta.ws.rs.core.MediaType;
 import jakarta.ws.rs.core.Response;
-import java.util.Optional;
+import java.net.URLDecoder;
+import java.nio.charset.Charset;
 import java.util.stream.Stream;
 import org.eclipse.microprofile.openapi.annotations.Operation;
 import org.eclipse.microprofile.openapi.annotations.enums.SchemaType;
@@ -24,7 +25,7 @@
 import org.eclipse.microprofile.openapi.annotations.responses.APIResponse;
 import org.keycloak.events.admin.OperationType;
 import org.keycloak.models.Constants;
-import org.keycloak.models.KeycloakSession;
+import org.keycloak.utils.StringUtil;
 
 public class TenantMembershipsResource extends AbstractAdminResource<TenantAdminAuth> {
 
@@ -40,17 +41,21 @@ public TenantMembershipsResource(AbstractAdminResource<TenantAdminAuth> parent,
     @Operation(operationId = "listMemberships", summary = "List tenant memberships")
     @APIResponse(responseCode = "200", description = "OK", content = @Content(schema = @Schema(type = SchemaType.ARRAY, implementation = TenantMembershipRepresentation.class)))
     public Stream<TenantMembershipRepresentation> listMemberships(
-            @Parameter(description = "Member email") @QueryParam("search") String searchQuery,
+            @Parameter(description = "Member email") @QueryParam("search") String search,
             @Parameter(description = "Pagination offset") @QueryParam("first") Integer firstResult,
             @Parameter(description = "Maximum results size (defaults to 100)") @QueryParam("max") Integer maxResults) {
-        Optional<String> search = Optional.ofNullable(searchQuery);
+
         firstResult = firstResult != null ? firstResult : 0;
         maxResults = maxResults != null ? maxResults : Constants.DEFAULT_MAX_RESULTS;
-        return tenant.getMembershipsStream()
-                .filter(m -> search.isEmpty() || m.getUser().getEmail().contains(search.get()))
-                .skip(firstResult)
-                .limit(maxResults)
-                .map(ModelMapper::toRepresentation);
+
+        if (StringUtil.isNotBlank(search)) {
+            search = URLDecoder.decode(search, Charset.defaultCharset()).trim().toLowerCase();
+            return tenant.getMembershipsStream(search, firstResult, maxResults)
+                    .map(ModelMapper::toRepresentation);
+        } else {
+            return tenant.getMembershipsStream(firstResult, maxResults)
+                    .map(ModelMapper::toRepresentation);
+        }
     }
 
     @PATCH

diff --git a/src/test/java/dev/sultanov/keycloak/multitenancy/ApiIntegrationTest.java b/src/test/java/dev/sultanov/keycloak/multitenancy/ApiIntegrationTest.java
@@ -44,7 +44,7 @@ void admin_shouldBeAbleToRevokeMembership_whenUserAcceptsInvitation() {
         assertThat(nextPage).isInstanceOf(ReviewInvitationsPage.class);
         ((ReviewInvitationsPage) nextPage).accept();
 
-        var userMembership = tenantResource.memberships().listMemberships("", null, null).stream()
+        var userMembership = tenantResource.memberships().listMemberships(user.getUserData().getEmail(), null, null).stream()
                 .filter(membership -> membership.getUser().getEmail().equalsIgnoreCase(user.getUserData().getEmail()))
                 .findFirst()
                 .orElseThrow();
@@ -54,7 +54,7 @@ void admin_shouldBeAbleToRevokeMembership_whenUserAcceptsInvitation() {
 
             // then
             assertThat(response.getStatus()).isEqualTo(HttpStatus.SC_NO_CONTENT);
-            assertThat(tenantResource.memberships().listMemberships("", null, null))
+            assertThat(tenantResource.memberships().listMemberships(null, null, null))
                     .extracting(TenantMembershipRepresentation::getUser)
                     .extracting(UserRepresentation::getEmail)
                     .extracting(String::toLowerCase)