Skip to content

Commit 905759f

Browse files
anarsultanovanarsultanov
committed
Optimize membership queries and fix search query
Fixes #29
1 parent 8f23b7f commit 905759f

File tree

8 files changed

+79
-36
lines changed

8 files changed

+79
-36
lines changed

‎src/main/java/dev/sultanov/keycloak/multitenancy/authentication/authenticators/IdpTenantMembershipsCreatingAuthenticator.java

+1-2
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33
import dev.sultanov.keycloak.multitenancy.authentication.IdentityProviderTenantsConfig;
44
import dev.sultanov.keycloak.multitenancy.model.TenantProvider;
55
import dev.sultanov.keycloak.multitenancy.util.Constants;
6-
import jakarta.ws.rs.core.Response;
76
import jakarta.ws.rs.core.Response.Status;
87
import java.util.Optional;
98
import java.util.Set;
@@ -59,7 +58,7 @@ private void doAuthenticate(AuthenticationFlowContext context, BrokeredIdentityC
5958
if (tenantById.isEmpty()) {
6059
log.warn("Tenant with ID %s, configured in IDP with alias %s, does not exist. Skipping membership creation."
6160
.formatted(tenantId, brokerContext.getIdpConfig().getAlias()));
62-
} else if (tenantById.get().getMembership(user).isPresent()) {
61+
} else if (tenantById.get().getMembershipByUser(user).isPresent()) {
6362
log.debug("User is already a member of tenant with ID %s. Skipping membership creation.".formatted(tenantId));
6463
} else {
6564
tenantById.get().grantMembership(user, Set.of(Constants.TENANT_USER_ROLE));

‎src/main/java/dev/sultanov/keycloak/multitenancy/model/TenantModel.java

+7-9
Original file line numberDiff line numberDiff line change
@@ -20,18 +20,16 @@ public interface TenantModel {
2020

2121
TenantMembershipModel grantMembership(UserModel user, Set<String> roles);
2222

23-
Stream<TenantMembershipModel> getMembershipsStream();
23+
Stream<TenantMembershipModel> getMembershipsStream(Integer firstResult, Integer maxResults);
2424

25-
default Optional<TenantMembershipModel> getMembershipById(String membershipId) {
26-
return getMembershipsStream().filter(membership -> membership.getId().equals(membershipId)).findFirst();
27-
};
25+
Stream<TenantMembershipModel> getMembershipsStream(String email, Integer firstResult, Integer maxResults);
2826

29-
default boolean hasMembership(UserModel user) {
30-
return getMembershipsStream().anyMatch(membership -> membership.getUser().getId().equals(user.getId()));
31-
}
27+
Optional<TenantMembershipModel> getMembershipById(String membershipId);
3228

33-
default Optional<TenantMembershipModel> getMembership(UserModel user) {
34-
return getMembershipsStream().filter(membership -> membership.getUser().getId().equals(user.getId())).findFirst();
29+
Optional<TenantMembershipModel> getMembershipByUser(UserModel user);
30+
31+
default boolean hasMembership(UserModel user) {
32+
return getMembershipByUser(user).isPresent();
3533
}
3634

3735
boolean revokeMembership(String membershipId);

‎src/main/java/dev/sultanov/keycloak/multitenancy/model/entity/TenantMembershipEntity.java

+7-2
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import jakarta.persistence.Id;
99
import jakarta.persistence.JoinColumn;
1010
import jakarta.persistence.ManyToOne;
11+
import jakarta.persistence.NamedQueries;
1112
import jakarta.persistence.NamedQuery;
1213
import jakarta.persistence.OneToOne;
1314
import jakarta.persistence.Table;
@@ -19,8 +20,12 @@
1920

2021
@Table(name = "TENANT_MEMBERSHIP", uniqueConstraints = {@UniqueConstraint(columnNames = {"TENANT_ID", "USER_ID"})})
2122
@Entity
22-
@NamedQuery(name = "getMembershipsByRealmAndUserId",
23-
query = "SELECT m FROM TenantMembershipEntity m WHERE m.tenant.realmId = :realmId AND m.user.id = :userId")
23+
@NamedQueries({
24+
@NamedQuery(name = "getMembershipsByRealmIdAndUserId", query = "SELECT m FROM TenantMembershipEntity m WHERE m.tenant.realmId = :realmId AND m.user.id = :userId"),
25+
@NamedQuery(name = "getMembershipsByTenantId", query = "SELECT m FROM TenantMembershipEntity m WHERE m.tenant.id = :tenantId"),
26+
@NamedQuery(name = "getMembershipsByTenantIdAndUserId", query = "SELECT m FROM TenantMembershipEntity m WHERE m.tenant.id = :tenantId AND m.user.id = :userId"),
27+
@NamedQuery(name = "getMembershipsByTenantIdAndUserEmail", query = "SELECT m FROM TenantMembershipEntity m WHERE m.tenant.id = :tenantId AND m.user.email = :email")
28+
})
2429
public class TenantMembershipEntity {
2530

2631
@Id

‎src/main/java/dev/sultanov/keycloak/multitenancy/model/jpa/JpaTenantProvider.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,7 @@ public Stream<TenantInvitationModel> getTenantInvitationsStream(RealmModel realm
8282

8383
@Override
8484
public Stream<TenantMembershipModel> getTenantMembershipsStream(RealmModel realm, UserModel user) {
85-
TypedQuery<TenantMembershipEntity> query = em.createNamedQuery("getMembershipsByRealmAndUserId", TenantMembershipEntity.class);
85+
TypedQuery<TenantMembershipEntity> query = em.createNamedQuery("getMembershipsByRealmIdAndUserId", TenantMembershipEntity.class);
8686
query.setParameter("realmId", realm.getId());
8787
query.setParameter("userId", user.getId());
8888
return query.getResultStream().map(m -> new TenantMembershipAdapter(session, realm, em, m));

‎src/main/java/dev/sultanov/keycloak/multitenancy/model/jpa/TenantAdapter.java

+44-8
Original file line numberDiff line numberDiff line change
@@ -7,13 +7,16 @@
77
import dev.sultanov.keycloak.multitenancy.model.entity.TenantInvitationEntity;
88
import dev.sultanov.keycloak.multitenancy.model.entity.TenantMembershipEntity;
99
import jakarta.persistence.EntityManager;
10+
import jakarta.persistence.TypedQuery;
1011
import java.util.HashSet;
12+
import java.util.Optional;
1113
import java.util.Set;
1214
import java.util.stream.Stream;
1315
import org.keycloak.models.KeycloakSession;
1416
import org.keycloak.models.RealmModel;
1517
import org.keycloak.models.UserModel;
1618
import org.keycloak.models.jpa.JpaModel;
19+
import org.keycloak.models.jpa.PaginationUtils;
1720
import org.keycloak.models.jpa.entities.UserEntity;
1821
import org.keycloak.models.utils.KeycloakModelUtils;
1922

@@ -47,32 +50,64 @@ public RealmModel getRealm() {
4750
}
4851

4952
@Override
50-
public TenantMembershipAdapter grantMembership(UserModel user, Set<String> roles) {
53+
public TenantMembershipModel grantMembership(UserModel user, Set<String> roles) {
5154
TenantMembershipEntity entity = new TenantMembershipEntity();
5255
entity.setId(KeycloakModelUtils.generateId());
5356
entity.setUser(em.getReference(UserEntity.class, user.getId()));
5457
entity.setTenant(tenant);
5558
entity.setRoles(new HashSet<>(roles));
5659
em.persist(entity);
60+
em.flush();
5761
tenant.getMemberships().add(entity);
5862
return new TenantMembershipAdapter(session, realm, em, entity);
5963
}
6064

6165
@Override
62-
public Stream<TenantMembershipModel> getMembershipsStream() {
63-
return tenant.getMemberships().stream()
64-
.map(membership -> new TenantMembershipAdapter(session, realm, em, membership));
66+
public Stream<TenantMembershipModel> getMembershipsStream(Integer first, Integer max) {
67+
TypedQuery<TenantMembershipEntity> query = em.createNamedQuery("getMembershipsByTenantId", TenantMembershipEntity.class);
68+
query.setParameter("tenantId", tenant.getId());
69+
return PaginationUtils.paginateQuery(query, first, max).getResultStream()
70+
.map((membership) -> new TenantMembershipAdapter(session, realm, em, membership));
71+
}
72+
73+
@Override
74+
public Stream<TenantMembershipModel> getMembershipsStream(String email, Integer first, Integer max) {
75+
TypedQuery<TenantMembershipEntity> query = em.createNamedQuery("getMembershipsByTenantIdAndUserEmail", TenantMembershipEntity.class);
76+
query.setParameter("tenantId", tenant.getId());
77+
query.setParameter("email", email);
78+
return PaginationUtils.paginateQuery(query, first, max).getResultStream()
79+
.map((membership) -> new TenantMembershipAdapter(session, realm, em, membership));
80+
}
81+
82+
@Override
83+
public Optional<TenantMembershipModel> getMembershipById(String membershipId) {
84+
TenantMembershipEntity membership = em.find(TenantMembershipEntity.class, membershipId);
85+
if (membership != null && realm.getId().equals(membership.getTenant().getRealmId())) {
86+
return Optional.of(new TenantMembershipAdapter(session, realm, em, membership));
87+
} else {
88+
return Optional.empty();
89+
}
90+
}
91+
92+
@Override
93+
public Optional<TenantMembershipModel> getMembershipByUser(UserModel user) {
94+
TypedQuery<TenantMembershipEntity> query = em.createNamedQuery("getMembershipsByTenantIdAndUserId", TenantMembershipEntity.class);
95+
query.setParameter("tenantId", tenant.getId());
96+
query.setParameter("userId", user.getId());
97+
return query.getResultStream().map(m -> (TenantMembershipModel) new TenantMembershipAdapter(session, realm, em, m)).findFirst();
6598
}
6699

67100
@Override
68101
public boolean revokeMembership(String membershipId) {
69-
var optionalMembership = getMembershipById(membershipId);
70-
if (optionalMembership.isPresent()) {
71-
var membershipEmail = optionalMembership.get().getUser().getEmail();
72-
tenant.getMemberships().removeIf(entity -> entity.getId().equals(membershipId));
102+
var membershipEntity = em.find(TenantMembershipEntity.class, membershipId);
103+
if (membershipEntity != null) {
104+
var membershipEmail = membershipEntity.getUser().getEmail();
105+
73106
if (membershipEmail != null) {
74107
revokeInvitations(membershipEmail);
75108
}
109+
em.remove(membershipEntity);
110+
em.flush();
76111
return true;
77112
}
78113
return false;
@@ -87,6 +122,7 @@ public TenantInvitationModel addInvitation(String email, UserModel inviter, Set<
87122
entity.setInvitedBy(inviter.getId());
88123
entity.setRoles(new HashSet<>(roles));
89124
em.persist(entity);
125+
em.flush();
90126
tenant.getInvitations().add(entity);
91127
return new TenantInvitationAdapter(session, realm, em, entity);
92128
}

‎src/main/java/dev/sultanov/keycloak/multitenancy/resource/TenantAdminAuth.java

+3-3
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
package dev.sultanov.keycloak.multitenancy.resource;
22

3-
import dev.sultanov.keycloak.multitenancy.util.Constants;
43
import dev.sultanov.keycloak.multitenancy.model.TenantModel;
4+
import dev.sultanov.keycloak.multitenancy.util.Constants;
55
import org.keycloak.models.ClientModel;
66
import org.keycloak.models.RealmModel;
77
import org.keycloak.models.UserModel;
@@ -15,10 +15,10 @@ public TenantAdminAuth(RealmModel realm, AccessToken token, UserModel user, Clie
1515
}
1616

1717
boolean isTenantAdmin(TenantModel tenantModel) {
18-
return tenantModel.getMembership(getUser()).filter(membership -> membership.getRoles().contains(Constants.TENANT_ADMIN_ROLE)).isPresent();
18+
return tenantModel.getMembershipByUser(getUser()).filter(membership -> membership.getRoles().contains(Constants.TENANT_ADMIN_ROLE)).isPresent();
1919
}
2020

2121
boolean isTenantMember(TenantModel tenantModel) {
22-
return tenantModel.getMembership(getUser()).isPresent();
22+
return tenantModel.getMembershipByUser(getUser()).isPresent();
2323
}
2424
}

‎src/main/java/dev/sultanov/keycloak/multitenancy/resource/TenantMembershipsResource.java

+14-9
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,8 @@
1313
import jakarta.ws.rs.QueryParam;
1414
import jakarta.ws.rs.core.MediaType;
1515
import jakarta.ws.rs.core.Response;
16-
import java.util.Optional;
16+
import java.net.URLDecoder;
17+
import java.nio.charset.Charset;
1718
import java.util.stream.Stream;
1819
import org.eclipse.microprofile.openapi.annotations.Operation;
1920
import org.eclipse.microprofile.openapi.annotations.enums.SchemaType;
@@ -24,7 +25,7 @@
2425
import org.eclipse.microprofile.openapi.annotations.responses.APIResponse;
2526
import org.keycloak.events.admin.OperationType;
2627
import org.keycloak.models.Constants;
27-
import org.keycloak.models.KeycloakSession;
28+
import org.keycloak.utils.StringUtil;
2829

2930
public class TenantMembershipsResource extends AbstractAdminResource<TenantAdminAuth> {
3031

@@ -40,17 +41,21 @@ public TenantMembershipsResource(AbstractAdminResource<TenantAdminAuth> parent,
4041
@Operation(operationId = "listMemberships", summary = "List tenant memberships")
4142
@APIResponse(responseCode = "200", description = "OK", content = @Content(schema = @Schema(type = SchemaType.ARRAY, implementation = TenantMembershipRepresentation.class)))
4243
public Stream<TenantMembershipRepresentation> listMemberships(
43-
@Parameter(description = "Member email") @QueryParam("search") String searchQuery,
44+
@Parameter(description = "Member email") @QueryParam("search") String search,
4445
@Parameter(description = "Pagination offset") @QueryParam("first") Integer firstResult,
4546
@Parameter(description = "Maximum results size (defaults to 100)") @QueryParam("max") Integer maxResults) {
46-
Optional<String> search = Optional.ofNullable(searchQuery);
47+
4748
firstResult = firstResult != null ? firstResult : 0;
4849
maxResults = maxResults != null ? maxResults : Constants.DEFAULT_MAX_RESULTS;
49-
return tenant.getMembershipsStream()
50-
.filter(m -> search.isEmpty() || m.getUser().getEmail().contains(search.get()))
51-
.skip(firstResult)
52-
.limit(maxResults)
53-
.map(ModelMapper::toRepresentation);
50+
51+
if (StringUtil.isNotBlank(search)) {
52+
search = URLDecoder.decode(search, Charset.defaultCharset()).trim().toLowerCase();
53+
return tenant.getMembershipsStream(search, firstResult, maxResults)
54+
.map(ModelMapper::toRepresentation);
55+
} else {
56+
return tenant.getMembershipsStream(firstResult, maxResults)
57+
.map(ModelMapper::toRepresentation);
58+
}
5459
}
5560

5661
@PATCH

‎src/test/java/dev/sultanov/keycloak/multitenancy/ApiIntegrationTest.java

+2-2
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ void admin_shouldBeAbleToRevokeMembership_whenUserAcceptsInvitation() {
4444
assertThat(nextPage).isInstanceOf(ReviewInvitationsPage.class);
4545
((ReviewInvitationsPage) nextPage).accept();
4646

47-
var userMembership = tenantResource.memberships().listMemberships("", null, null).stream()
47+
var userMembership = tenantResource.memberships().listMemberships(user.getUserData().getEmail(), null, null).stream()
4848
.filter(membership -> membership.getUser().getEmail().equalsIgnoreCase(user.getUserData().getEmail()))
4949
.findFirst()
5050
.orElseThrow();
@@ -54,7 +54,7 @@ void admin_shouldBeAbleToRevokeMembership_whenUserAcceptsInvitation() {
5454

5555
// then
5656
assertThat(response.getStatus()).isEqualTo(HttpStatus.SC_NO_CONTENT);
57-
assertThat(tenantResource.memberships().listMemberships("", null, null))
57+
assertThat(tenantResource.memberships().listMemberships(null, null, null))
5858
.extracting(TenantMembershipRepresentation::getUser)
5959
.extracting(UserRepresentation::getEmail)
6060
.extracting(String::toLowerCase)

0 commit comments

Comments
 (0)