keycloak · davoustp · May 2, 2022
@@ -17,16 +17,23 @@
 
 package org.keycloak.models.jpa;
 
+import org.hibernate.Hibernate;
+import org.hibernate.SynchronizeableQuery;
+import org.keycloak.connections.jpa.util.JpaUtils;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.RealmModel;
 import org.keycloak.models.RoleContainerModel;
 import org.keycloak.models.RoleModel;
+import org.keycloak.models.jpa.entities.CompositeRoleEntity;
+import org.keycloak.models.jpa.entities.CompositeRoleEntityKey;
 import org.keycloak.models.jpa.entities.RoleAttributeEntity;
 import org.keycloak.models.jpa.entities.RoleEntity;
 import org.keycloak.models.utils.KeycloakModelUtils;
 
 import javax.persistence.EntityManager;
 import javax.persistence.Query;
+import javax.persistence.TypedQuery;
+
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
@@ -41,6 +48,7 @@
  * @version $Revision: 1 $
  */
 public class RoleAdapter implements RoleModel, JpaModel<RoleEntity> {
+
     protected RoleEntity role;
     protected EntityManager em;
     protected RealmModel realm;
@@ -89,30 +97,95 @@ public void setName(String name) {
 
     @Override
     public boolean isComposite() {
-        return getCompositesStream().count() > 0;
+        // Use the composite role collection if already loaded, or use a named query to avoid triggering its lazy loading
+        if (Hibernate.isPropertyInitialized(getEntity(), "compositeRoles")) {
+            return !getEntity().getCompositeRoles().isEmpty();
+        }
+        else {
+            TypedQuery<String> query = em.createNamedQuery("getChildrenRoleIds", String.class);
+            query.setParameter("roleId", getId());
+            query.setMaxResults(1);
+            return !query.getResultList().isEmpty();
+        }
     }
 
     @Override
     public void addCompositeRole(RoleModel role) {
+        // Avoid lazy loading the composite role collection if not already done
+        // Not using Persistence.getPersistenceUtil().isLoaded(Object, String) as this is not working
+        // properly (Hibernate 5.6) - still returning false even after the collection was lazy loaded
+        if (Hibernate.isPropertyInitialized(getEntity(), "compositeRoles")) {
+            addCompositeRoleUsingLoadedCompositeCollection(role);
+        }
+        else {
+            addCompositeRoleWithoutLoadingCompositeCollection(role);
+        }
+    }
+
+    private void addCompositeRoleUsingLoadedCompositeCollection(RoleModel role) {
         RoleEntity entity = toRoleEntity(role);
+        // Why performing this loop at all? The semantic of Set.add(T) ensures that the operation
+        // is performed only if there is not entry already...
         for (RoleEntity composite : getEntity().getCompositeRoles()) {
             if (composite.equals(entity)) return;
         }
         getEntity().getCompositeRoles().add(entity);
     }
 
+    private void addCompositeRoleWithoutLoadingCompositeCollection(RoleModel role) {
+        // Ensure that the entry does not exist already - will hit the database,
+        // but it's required to maintain the semantic of method #addCompositeRole(RoleModel)
+        CompositeRoleEntityKey compositeKey = new CompositeRoleEntityKey(this.getId(), role.getId());
+        if (em.find(CompositeRoleEntity.class, compositeKey) == null) {
+            // Using a native query to perform insertion (not possible with JPQL), allowing to instruct
+            // the auto-flushing mechanism which query spaces (tables) should be flushed
+            // (and invalidated in 2nd level cache) before executing it.
+            String compositeRoleTable = JpaUtils.getTableNameForNativeQuery("COMPOSITE_ROLE", em);
+            Query q = em.createNativeQuery("insert into " + compositeRoleTable + " (COMPOSITE, CHILD_ROLE) values (:composite, :child)")
+                    .setParameter("composite", this.getId())
+                    .setParameter("child", role.getId())
+                    ;
+            // Table for the RoleEntity class is added to the query space so that
+            // any pending operation in this table is flushed as well.
+            SynchronizeableQuery<?> sq = q.unwrap(SynchronizeableQuery.class);
+            sq.addSynchronizedEntityClass(CompositeRoleEntity.class, RoleEntity.class);
+            q.executeUpdate();
+        }
+    }
+
     @Override
     public void removeCompositeRole(RoleModel role) {
+        // Avoid lazy loading the composite role collection if not already done
+        // Not using Persistence.getPersistenceUtil().isLoaded(Object, String) as this is not working
+        // properly (Hibernate 5.6) - still returning false even after the collection was lazy loaded
+        if (Hibernate.isPropertyInitialized(getEntity(), "compositeRoles")) {
+            removeCompositeRoleUsingLoadedCompositeCollection(role);
+        }
+        else {
+            removeCompositeRoleWithoutLoadingCompositeCollection(role);
+        }
+    }
+
+    private void removeCompositeRoleUsingLoadedCompositeCollection(RoleModel role) {
         RoleEntity entity = toRoleEntity(role);
         getEntity().getCompositeRoles().remove(entity);
     }
 
+    private void removeCompositeRoleWithoutLoadingCompositeCollection(RoleModel role) {
+        // Using a named query here to avoid the two-steps find() + delete() operation when using
+        // EntityManager.
+        em.createNamedQuery("removeCompositeAndChildRoleEntry")
+            .setParameter("compositeId", this.getId())
+            .setParameter("childId", role.getId())
+            .executeUpdate();
+    }
+
     @Override
     public Stream<RoleModel> getCompositesStream() {
         Stream<RoleModel> composites = getEntity().getCompositeRoles().stream().map(c -> new RoleAdapter(session, realm, em, c));
         return composites.filter(Objects::nonNull);
     }
-
+    
     @Override
     public Stream<RoleModel> getCompositesStream(String search, Integer first, Integer max) {
         return session.roles().getRolesStream(realm,
@@ -214,4 +287,5 @@ private RoleEntity toRoleEntity(RoleModel model) {
         }
         return em.getReference(RoleEntity.class, model.getId());
     }
-}
+
+}
@@ -0,0 +1,90 @@
+package org.keycloak.models.jpa.entities;
+
+import java.io.Serializable;
+import java.util.Objects;
+
+import javax.persistence.Column;
+import javax.persistence.Entity;
+import javax.persistence.Id;
+import javax.persistence.IdClass;
+import javax.persistence.NamedQueries;
+import javax.persistence.NamedQuery;
+import javax.persistence.QueryHint;
+import javax.persistence.Table;
+import javax.persistence.UniqueConstraint;
+
+import org.hibernate.jpa.QueryHints;
+
+@Entity
+@IdClass(CompositeRoleEntityKey.class)
+@Table(name="COMPOSITE_ROLE", uniqueConstraints = {
+      @UniqueConstraint(columnNames = { "COMPOSITE", "CHILD_ROLE" })
+})
+@NamedQueries({
+    @NamedQuery(name="getChildrenRoleIds", hints = @QueryHint(name = QueryHints.HINT_READONLY, value = "true"),
+            query="select compositerole.childRoleId from CompositeRoleEntity compositerole where compositerole.compositeId = :roleId"),
+    @NamedQuery(name="getCompositeRolesByCompositeIds", hints = @QueryHint(name = QueryHints.HINT_READONLY, value = "true"),
+            query="select compositerole from CompositeRoleEntity compositerole where compositerole.compositeId in :roleIds"),
+    @NamedQuery(name="getChildRoleIdsForCompositeIds", hints = @QueryHint(name = QueryHints.HINT_READONLY, value = "true"),
+            query="select compositerole.childRoleId from CompositeRoleEntity compositerole where compositerole.compositeId in :roleIds"),
+    @NamedQuery(name="removeCompositeAndChildRoleEntry",
+            query="delete from CompositeRoleEntity compositerole where compositerole.compositeId = :compositeId and compositerole.childRoleId = :childId"),
+})
+public class CompositeRoleEntity implements Serializable {
+
+    private static final long serialVersionUID = 7375648337789837909L;
+
+    @Id
+    @Column(name="COMPOSITE", length = 36, nullable = false)
+    private String compositeId;
+
+    @Id
+    @Column(name="CHILD_ROLE", length = 36, nullable = false)
+    private String childRoleId;
+
+    public CompositeRoleEntity() {
+        super();
+    }
+
+    public CompositeRoleEntity(CompositeRoleEntityKey key) {
+        super();
+        setCompositeId(key.getCompositeId());
+        setChildRoleId(key.getChildRoleId());
+    }
+
+    public String getCompositeId() {
+        return compositeId;
+    }
+
+    public void setCompositeId(String compositeId) {
+        this.compositeId = compositeId;
+    }
+
+    public String getChildRoleId() {
+        return childRoleId;
+    }
+
+    public void setChildRoleId(String childRoleId) {
+        this.childRoleId = childRoleId;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null) return false;
+        if (!(o instanceof CompositeRoleEntity)) return false;
+
+        CompositeRoleEntity that = (CompositeRoleEntity) o;
+
+        if (!compositeId.equals(that.compositeId)) return false;
+        if (!childRoleId.equals(that.childRoleId)) return false;
+
+        return true;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(compositeId, childRoleId);
+    }
+
+}
@@ -0,0 +1,64 @@
+package org.keycloak.models.jpa.entities;
+
+import java.io.Serializable;
+import java.util.Objects;
+
+/**
+ * The composite primary key representation for {@link CompositeRoleEntity}.
+ * Required to perform lookup by primary key through JPA entity manager.
+ */
+public class CompositeRoleEntityKey implements Serializable {
+
+    private static final long serialVersionUID = -6078479162595894390L;
+
+    private String compositeId;
+    private String childRoleId;
+
+    public CompositeRoleEntityKey() {
+        super();
+    }
+
+    public CompositeRoleEntityKey(String compositeId, String childRoleId) {
+        super();
+        this.compositeId = compositeId;
+        this.childRoleId = childRoleId;
+    }
+
+    public String getCompositeId() {
+        return compositeId;
+    }
+
+    public String getChildRoleId() {
+        return childRoleId;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null) return false;
+        if (!(o instanceof CompositeRoleEntityKey)) return false;
+
+        CompositeRoleEntityKey that = (CompositeRoleEntityKey) o;
+
+        if (!compositeId.equals(that.compositeId)) return false;
+        if (!childRoleId.equals(that.childRoleId)) return false;
+
+        return true;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(compositeId, childRoleId);
+    }
+
+    @Override
+    public String toString() {
+        StringBuilder builder = new StringBuilder();
+        builder.append("CompositeRolePK [")
+                .append("compositeId=").append(compositeId)
+                .append(", childRoleId=").append(childRoleId)
+                .append("]");
+        return builder.toString();
+    }
+
+}
@@ -32,6 +32,7 @@
         <class>org.keycloak.models.jpa.entities.UserFederationMapperEntity</class>
         <class>org.keycloak.models.jpa.entities.RoleEntity</class>
         <class>org.keycloak.models.jpa.entities.RoleAttributeEntity</class>
+        <class>org.keycloak.models.jpa.entities.CompositeRoleEntity</class>
         <class>org.keycloak.models.jpa.entities.FederatedIdentityEntity</class>
         <class>org.keycloak.models.jpa.entities.MigrationModelEntity</class>
         <class>org.keycloak.models.jpa.entities.UserEntity</class>