From e9e6b73bd27b71be91731e1e311543118f29d05f Mon Sep 17 00:00:00 2001
From: Alexander Schwartz <aschwart@redhat.com>
Date: Wed, 11 Jan 2023 10:46:46 +0100
Subject: [PATCH] Avoid using Hibernate APIs to cache query results as the API
 changes in Hibernate 6

Closes #16332
---
 .../jpa/JpaMapKeycloakTransaction.java        | 85 ++++++++-----------
 .../models/map/storage/QueryParameters.java   | 14 +++
 2 files changed, 51 insertions(+), 48 deletions(-)

diff --git a/model/map-jpa/src/main/java/org/keycloak/models/map/storage/jpa/JpaMapKeycloakTransaction.java b/model/map-jpa/src/main/java/org/keycloak/models/map/storage/jpa/JpaMapKeycloakTransaction.java
index 54c131e3e9..7c10bb7d37 100644
--- a/model/map-jpa/src/main/java/org/keycloak/models/map/storage/jpa/JpaMapKeycloakTransaction.java
+++ b/model/map-jpa/src/main/java/org/keycloak/models/map/storage/jpa/JpaMapKeycloakTransaction.java
@@ -16,6 +16,7 @@
  */
 package org.keycloak.models.map.storage.jpa;
 
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.LinkedList;
 import java.util.List;
@@ -25,7 +26,6 @@ import java.util.UUID;
 import java.util.stream.Stream;
 import javax.persistence.EntityManager;
 import javax.persistence.LockModeType;
-import javax.persistence.Parameter;
 import javax.persistence.PersistenceException;
 import javax.persistence.TypedQuery;
 import javax.persistence.criteria.CriteriaBuilder;
@@ -38,11 +38,9 @@ import javax.persistence.criteria.Selection;
 
 import org.hibernate.Session;
 import org.hibernate.internal.SessionImpl;
-import org.hibernate.query.spi.QueryImplementor;
 import org.jboss.logging.Logger;
 import org.keycloak.common.util.Time;
 import org.keycloak.models.KeycloakSession;
-import org.keycloak.models.ModelException;
 import org.keycloak.models.map.common.AbstractEntity;
 import org.keycloak.models.map.common.ExpirableEntity;
 import org.keycloak.models.map.common.StringKeyConverter;
@@ -134,6 +132,26 @@ public abstract class JpaMapKeycloakTransaction<RE extends JpaRootEntity, E exte
     @Override
     @SuppressWarnings("unchecked")
     public Stream<E> read(QueryParameters<M> queryParameters) {
+        Map<QueryCacheKey, List<RE>> cache = getQueryCache();
+        QueryCacheKey queryCacheKey = new QueryCacheKey(queryParameters, modelType);
+        if (!LockObjectsForModification.isEnabled(this.session, modelType)) {
+            List<RE> previousResult = cache.get(queryCacheKey);
+            SessionImpl session = (SessionImpl) em.unwrap(Session.class);
+            // only do dirty checking if there is a previously cached result that would match the query
+            if (previousResult != null) {
+                // if the session is dirty, data has been modified, and the cache must not be used
+                // check if there are queued actions already, as this allows us to skip the expensive dirty check
+                if (!session.getActionQueue().areInsertionsOrDeletionsQueued() && session.getActionQueue().numberOfUpdates() == 0 && session.getActionQueue().numberOfCollectionUpdates() == 0 &&
+                        !session.isDirty()) {
+                    logger.tracef("tx %d: cache hit for %s/%s%s", hashCode(), queryParameters, queryCacheKey, getShortStackTrace());
+                    return closing(previousResult.stream()).map(this::mapToEntityDelegateUnique);
+                } else {
+                    logger.tracef("tx %d: cache ignored due to dirty session", hashCode());
+                }
+            }
+        }
+        logger.tracef("tx %d: cache miss for %s/%s%s", hashCode(), queryParameters, queryCacheKey, getShortStackTrace());
+
         JpaModelCriteriaBuilder mcb = queryParameters.getModelCriteriaBuilder()
                 .flashToModelCriteriaBuilder(createJpaModelCriteriaBuilder());
 
@@ -170,27 +188,6 @@ public abstract class JpaMapKeycloakTransaction<RE extends JpaRootEntity, E exte
 
         TypedQuery<RE> emQuery = paginateQuery(em.createQuery(query), queryParameters.getOffset(), queryParameters.getLimit());
 
-        Map<QueryCacheKey, List<RE>> cache = getQueryCache();
-        QueryCacheKey queryCacheKey = new QueryCacheKey(emQuery, modelType);
-        if (!LockObjectsForModification.isEnabled(this.session, modelType)) {
-            List<RE> previousResult = cache.get(queryCacheKey);
-            //noinspection resource
-            SessionImpl session = (SessionImpl) em.unwrap(Session.class);
-            // only do dirty checking if there is a previously cached result that would match the query
-            if (previousResult != null) {
-                // if the session is dirty, data has been modified, and the cache must not be used
-                // check if there are queued actions already, as this allows us to skip the expensive dirty check
-                if (!session.getActionQueue().areInsertionsOrDeletionsQueued() && session.getActionQueue().numberOfUpdates() == 0 && session.getActionQueue().numberOfCollectionUpdates() == 0 &&
-                        !session.isDirty()) {
-                    logger.tracef("tx %d: cache hit for %s/%s%s", hashCode(), queryParameters, queryCacheKey, getShortStackTrace());
-                    return closing(previousResult.stream()).map(this::mapToEntityDelegateUnique);
-                } else {
-                    logger.tracef("tx %d: cache ignored due to dirty session", hashCode());
-                }
-            }
-        }
-        logger.tracef("tx %d: cache miss for %s/%s%s", hashCode(), queryParameters, queryCacheKey, getShortStackTrace());
-
         if (LockObjectsForModification.isEnabled(session, modelType)) {
             emQuery = emQuery.setLockMode(LockModeType.PESSIMISTIC_WRITE);
         }
@@ -209,11 +206,10 @@ public abstract class JpaMapKeycloakTransaction<RE extends JpaRootEntity, E exte
     }
 
     private Map<QueryCacheKey, List<RE>> getQueryCache() {
-        //noinspection resource,unchecked
+        //noinspection unchecked
         Map<QueryCacheKey, List<RE>> cache = (Map<QueryCacheKey, List<RE>>) em.unwrap(Session.class).getProperties().get(JPA_MAP_CACHE);
         if (cache == null) {
             cache = new HashMap<>();
-            //noinspection resource
             em.unwrap(Session.class).setProperty(JPA_MAP_CACHE, cache);
         }
         return cache;
@@ -345,24 +341,17 @@ public abstract class JpaMapKeycloakTransaction<RE extends JpaRootEntity, E exte
 
     private static class QueryCacheKey {
         private final String queryString;
-        private final Integer queryMaxResults;
-        private final Integer queryFirstResult;
-        private final HashMap<String, Object> queryParameters;
+        private final Integer queryLimit;
+        private final Integer queryOffset;
         private final Class<?> modelType;
+        private final List<? extends QueryParameters.OrderBy<?>> queryOrderBy;
 
-        public QueryCacheKey(TypedQuery<?> emQuery, Class<?> modelType) {
+        public QueryCacheKey(QueryParameters<?> query, Class<?> modelType) {
             // copy over all fields from the query that relevant for caching
-            QueryImplementor<?> query = emQuery.unwrap(QueryImplementor.class);
-            this.queryString = query.getQueryString();
-            this.queryParameters = new HashMap<>();
-            for (Parameter<?> parameter : query.getParameters()) {
-                if (parameter.getName() == null) {
-                    throw new ModelException("Can't prepare query for caching as parameter doesn't have a name");
-                }
-                this.queryParameters.put(parameter.getName(), query.getParameterValue(parameter.getName()));
-            }
-            this.queryMaxResults = emQuery.getMaxResults();
-            this.queryFirstResult = emQuery.getFirstResult();
+            this.queryString = query.getModelCriteriaBuilder().toString();
+            this.queryLimit = query.getLimit();
+            this.queryOffset = query.getOffset();
+            this.queryOrderBy = new ArrayList<>(query.getOrderBy());
             this.modelType = modelType;
         }
 
@@ -372,24 +361,24 @@ public abstract class JpaMapKeycloakTransaction<RE extends JpaRootEntity, E exte
             if (o == null || getClass() != o.getClass()) return false;
             QueryCacheKey that = (QueryCacheKey) o;
             return Objects.equals(queryString, that.queryString)
-                    && Objects.equals(queryMaxResults, that.queryMaxResults)
-                    && Objects.equals(queryFirstResult, that.queryFirstResult)
-                    && Objects.equals(queryParameters, that.queryParameters)
+                    && Objects.equals(queryLimit, that.queryLimit)
+                    && Objects.equals(queryOffset, that.queryOffset)
+                    && Objects.equals(queryOrderBy, that.queryOrderBy)
                     && Objects.equals(modelType, that.modelType);
         }
 
         @Override
         public int hashCode() {
-            return Objects.hash(queryString, queryMaxResults, queryFirstResult, queryParameters, modelType);
+            return Objects.hash(queryString, queryLimit, queryOffset, queryOrderBy, modelType);
         }
 
         @Override
         public String toString() {
             return "QueryCacheKey{" +
                     "queryString='" + queryString + '\'' +
-                    ", queryMaxResults=" + queryMaxResults +
-                    ", queryFirstResult=" + queryFirstResult +
-                    ", queryParameters=" + queryParameters +
+                    ", queryMaxResults=" + queryLimit +
+                    ", queryFirstResult=" + queryOffset +
+                    ", queryOrderBy=" + queryOrderBy +
                     ", modelType=" + modelType.getName() +
                     '}';
         }
diff --git a/model/map/src/main/java/org/keycloak/models/map/storage/QueryParameters.java b/model/map/src/main/java/org/keycloak/models/map/storage/QueryParameters.java
index d01d209578..f520e5c1df 100644
--- a/model/map/src/main/java/org/keycloak/models/map/storage/QueryParameters.java
+++ b/model/map/src/main/java/org/keycloak/models/map/storage/QueryParameters.java
@@ -5,6 +5,7 @@ import org.keycloak.storage.SearchableModelField;
 
 import java.util.LinkedList;
 import java.util.List;
+import java.util.Objects;
 
 import static org.keycloak.models.map.storage.QueryParameters.Order.ASCENDING;
 
@@ -146,5 +147,18 @@ public class QueryParameters<M> {
         public Order getOrder() {
             return order;
         }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (o == null || getClass() != o.getClass()) return false;
+            OrderBy<?> orderBy = (OrderBy<?>) o;
+            return Objects.equals(modelField, orderBy.modelField) && order == orderBy.order;
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(modelField, order);
+        }
     }
 }