KEYCLOAK-18842: deleteExpiredClientSessions very slow on MariaDB

model/jpa/src/main/java/org/keycloak/connections/jpa/DefaultJpaConnectionProviderFactory.java

@@ -25,6 +25,7 @@ import org.keycloak.ServerStartupError;
 import org.keycloak.common.util.StackUtil;
 import org.keycloak.common.util.StringPropertyReplacer;
 import org.keycloak.connections.jpa.updater.JpaUpdaterProvider;
+import org.keycloak.connections.jpa.updater.liquibase.conn.LiquibaseConnectionProvider;
 import org.keycloak.connections.jpa.util.JpaUtils;
 import org.keycloak.migration.MigrationModelManager;
 import org.keycloak.models.KeycloakSession;
@@ -54,6 +55,8 @@ import java.util.Collection;
 import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.Map;
+import liquibase.Liquibase;
+import liquibase.exception.LiquibaseException;
 
 /**
  * @author <a href="mailto:sthorger@redhat.com">Stian Thorgersen</a>
@@ -82,6 +85,10 @@ public class DefaultJpaConnectionProviderFactory implements JpaConnectionProvide
         logger.trace("Create JpaConnectionProvider");
         lazyInit(session);
 
+        return new DefaultJpaConnectionProvider(createEntityManager(session));
+    }
+
+    private EntityManager createEntityManager(KeycloakSession session) {
         EntityManager em;
         if (!jtaEnabled) {
             logger.trace("enlisting EntityManager in JpaKeycloakTransaction");
@@ -91,8 +98,24 @@ public class DefaultJpaConnectionProviderFactory implements JpaConnectionProvide
             em = emf.createEntityManager(SynchronizationType.SYNCHRONIZED);
         }
         em = PersistenceExceptionConverter.create(session, em);
-        if (!jtaEnabled) session.getTransactionManager().enlist(new JpaKeycloakTransaction(em));
+        if (!jtaEnabled) {
-        return new DefaultJpaConnectionProvider(em);
+            session.getTransactionManager().enlist(new JpaKeycloakTransaction(em));
+        }
+        return em;
+    }
+
+    private void addSpecificNamedQueries(KeycloakSession session, Connection connection) {
+        LiquibaseConnectionProvider liquibaseProvider = session.getProvider(LiquibaseConnectionProvider.class);
+        EntityManager em = null;
+        try {
+            Liquibase liquibase = liquibaseProvider.getLiquibase(connection, this.getSchema());
+            em = createEntityManager(session);
+            JpaUtils.addSpecificNamedQueries(em, liquibase.getDatabase().getShortName());
+        } catch (LiquibaseException e) {
+            throw new IllegalStateException(e);
+        } finally {
+            JpaUtils.closeEntityManager(em);
+        }
     }
 
     @Override
@@ -210,6 +233,7 @@ public class DefaultJpaConnectionProviderFactory implements JpaConnectionProvide
                             classLoaders.add(getClass().getClassLoader());
                             properties.put(AvailableSettings.CLASSLOADERS, classLoaders);
                             emf = JpaUtils.createEntityManagerFactory(session, unitName, properties, jtaEnabled);
+                            addSpecificNamedQueries(session, connection);
                             logger.trace("EntityManagerFactory created");
 
                             if (globalStatsInterval != -1) {

model/jpa/src/main/java/org/keycloak/connections/jpa/util/JpaUtils.java

@@ -17,6 +17,10 @@
 
 package org.keycloak.connections.jpa.util;
 
+import org.jboss.logging.Logger;
+import org.hibernate.engine.query.spi.sql.NativeSQLQueryReturn;
+import org.hibernate.engine.query.spi.sql.NativeSQLQuerySpecification;
+import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.jpa.boot.internal.ParsedPersistenceXmlDescriptor;
 import org.hibernate.jpa.boot.internal.PersistenceXmlParser;
 import org.hibernate.jpa.boot.spi.Bootstrap;
@@ -27,9 +31,14 @@ import org.keycloak.models.KeycloakSession;
 import javax.persistence.EntityManager;
 import javax.persistence.EntityManagerFactory;
 import javax.persistence.spi.PersistenceUnitTransactionType;
+import java.io.IOException;
+import java.io.InputStream;
+import java.net.URL;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Properties;
 import java.util.Set;
 
 /**
@@ -38,6 +47,9 @@ import java.util.Set;
 public class JpaUtils {
 
     public static final String HIBERNATE_DEFAULT_SCHEMA = "hibernate.default_schema";
+    public static final String QUERY_NATIVE_SUFFIX = "[native]";
+    public static final String QUERY_JPQL_SUFFIX = "[jpql]";
+    private static final Logger logger = Logger.getLogger(JpaUtils.class);
 
     public static String getTableNameForNativeQuery(String tableName, EntityManager em) {
         String schema = (String) em.getEntityManagerFactory().getProperties().get(HIBERNATE_DEFAULT_SCHEMA);
@@ -93,4 +105,124 @@ public class JpaUtils {
         return "DATABASECHANGELOG_" + upperCased.substring(0, Math.min(10, upperCased.length()));
     }
 
+    /**
+     * Loads the URL as a properties file.
+     * @param url The url to load, it can be null
+     * @return A properties file with the url loaded or null
+     */
+    private static Properties loadSqlProperties(URL url) {
+        if (url == null) {
+            return null;
+        }
+        Properties props = new Properties();
+        try (InputStream is = url.openStream()) {
+            props.load(is);
+        } catch (IOException e) {
+            throw new IllegalStateException(e);
+        }
+        return props;
+    }
+
+    /**
+     * Returns the name of the query in the queries file. It searches for the
+     * three possible forms: name[native], name[jpql] or name.
+     * @param name The name of the query to search
+     * @param queries The properties file with the queries
+     * @return The key with the query found or null if not found
+     */
+    private static String getQueryFromProperties(String name, Properties queries) {
+        if (queries == null) {
+            return null;
+        }
+        String nameFull = name + QUERY_NATIVE_SUFFIX;
+        if (queries.containsKey(nameFull)) {
+            return nameFull;
+        }
+        nameFull = name + QUERY_JPQL_SUFFIX;
+        if (queries.containsKey(nameFull)) {
+            return nameFull;
+        }
+        nameFull = name;
+        if (queries.containsKey(nameFull)) {
+            return nameFull;
+        }
+        return null;
+    }
+
+    /**
+     * Returns the query name but removing the suffix.
+     * @param name The query name as it is on the key
+     * @return The name without the suffix
+     */
+    private static String getQueryShortName(String name) {
+        if (name.endsWith(QUERY_NATIVE_SUFFIX)) {
+            return name.substring(0, name.length() - QUERY_NATIVE_SUFFIX.length());
+        } else if (name.endsWith(QUERY_JPQL_SUFFIX)) {
+            return name.substring(0, name.length() - QUERY_JPQL_SUFFIX.length());
+        } else {
+            return name;
+        }
+    }
+
+    /**
+     * Method that adds the different query variants for the database.
+     * The method loads the queries specified in the files
+     * <em>META-INF/queries-{dbType}.properties</em> and the default
+     * <em>META-INF/queries-default.properties</em>. At least the default file
+     * should exist inside the jar file. The default file contains all the
+     * needed queries and the specific one can overload all or some of them for
+     * that database type.
+     * @param em The entity manager to use
+     * @param databaseType The database type as managed in
+     *                     <a href="https://www.liquibase.org/get-started/databases">liquibase</a>.
+     */
+    public static void addSpecificNamedQueries(EntityManager em, String databaseType) {
+        final SessionFactoryImplementor sfi = em.getEntityManagerFactory().unwrap(SessionFactoryImplementor.class);
+        URL specificUrl = JpaUtils.class.getClassLoader().getResource("META-INF/queries-" + databaseType + ".properties");
+        URL defaultUrl = JpaUtils.class.getClassLoader().getResource("META-INF/queries-default.properties");
+        if (defaultUrl == null) {
+            throw new IllegalStateException("META-INF/queries-default.properties was not found in the classpath");
+        }
+        Properties specificQueries = loadSqlProperties(specificUrl);
+        Properties defaultQueries = loadSqlProperties(defaultUrl);
+
+        for (String queryNameFull : defaultQueries.stringPropertyNames()) {
+            String querySql;
+            String queryName = getQueryShortName(queryNameFull);
+            String specificQueryNameFull = getQueryFromProperties(queryName, specificQueries);
+            if (specificQueryNameFull != null) {
+                // the query is redefined in the specific database file => use it
+                queryNameFull = specificQueryNameFull;
+                querySql = specificQueries.getProperty(queryNameFull);
+            } else {
+                // use the default query sql
+                querySql = defaultQueries.getProperty(queryNameFull);
+            }
+            boolean isNative = queryNameFull.endsWith(QUERY_NATIVE_SUFFIX);
+
+            logger.tracef("adding query from properties files native=%b %s:%s", isNative, queryName, querySql);
+            if (isNative) {
+                NativeSQLQuerySpecification spec = new NativeSQLQuerySpecification(querySql, new NativeSQLQueryReturn[0], Collections.emptySet());
+                sfi.getQueryPlanCache().getNativeSQLQueryPlan(spec);
+                em.getEntityManagerFactory().addNamedQuery(queryName, em.createNativeQuery(querySql));
+            } else {
+                sfi.getQueryPlanCache().getHQLQueryPlan(querySql, false, Collections.emptyMap());
+                em.getEntityManagerFactory().addNamedQuery(queryName, em.createQuery(querySql));
+            }
+        }
+    }
+
+    /**
+     * Helper to close the entity manager.
+     * @param em The entity manager to close
+     */
+    public static void closeEntityManager(EntityManager em) {
+        if (em != null) {
+            try {
+                em.close();
+            } catch (Exception e) {
+                logger.warn("Failed to close entity manager", e);
+            }
+        }
+    }
 }

model/jpa/src/main/java/org/keycloak/models/jpa/session/PersistentClientSessionEntity.java

@@ -36,7 +36,9 @@ import java.io.Serializable;
         @NamedQuery(name="deleteClientSessionsByClientStorageProvider", query="delete from PersistentClientSessionEntity sess where sess.clientStorageProvider = :clientStorageProvider"),
         @NamedQuery(name="deleteClientSessionsByUser", query="delete from PersistentClientSessionEntity sess where sess.userSessionId IN (select u.userSessionId from PersistentUserSessionEntity u where u.userId = :userId)"),
         @NamedQuery(name="deleteClientSessionsByUserSession", query="delete from PersistentClientSessionEntity sess where sess.userSessionId = :userSessionId and sess.offline = :offline"),
-        @NamedQuery(name="deleteExpiredClientSessions", query="delete from PersistentClientSessionEntity sess where sess.userSessionId IN (select u.userSessionId from PersistentUserSessionEntity u where u.realmId = :realmId AND u.offline = :offline AND u.lastSessionRefresh < :lastSessionRefresh)"),
+        // KEYCLOAK-18842: The deleteExpiredClientSessions performs very slow in MySQL/MariaDB databases
+        //                 It is removed from here and added manually in JpaUtils to give a native implementation if needed
+        //@NamedQuery(name="deleteExpiredClientSessions", query="delete from PersistentClientSessionEntity sess where sess.userSessionId IN (select u.userSessionId from PersistentUserSessionEntity u where u.realmId = :realmId AND u.offline = :offline AND u.lastSessionRefresh < :lastSessionRefresh)"),
         @NamedQuery(name="findClientSessionsByUserSession", query="select sess from PersistentClientSessionEntity sess where sess.userSessionId=:userSessionId and sess.offline = :offline"),
         @NamedQuery(name="findClientSessionsOrderedById", query="select sess from PersistentClientSessionEntity sess where sess.offline = :offline and sess.userSessionId >= :fromSessionId and sess.userSessionId <= :toSessionId order by sess.userSessionId"),
         @NamedQuery(name="findClientSessionsCountByClient", query="select count(sess) from PersistentClientSessionEntity sess where sess.offline = :offline and sess.clientId = :clientId")

model/jpa/src/main/resources/META-INF/queries-default.properties

@@ -0,0 +1,10 @@
+# properties file to define all default queries that are loaded separately
+# in a properties file. These queries can be overloaded with a
+# specific file for each database type. Queries are defined in the form:
+#     name[type]=sql
+# type can be native (for native queries) or jpql (jpql syntax)
+# if no type is defined jpql is the default
+
+deleteExpiredClientSessions=delete from PersistentClientSessionEntity sess where sess.userSessionId IN (\
+ select u.userSessionId from PersistentUserSessionEntity u \
+ where u.realmId = :realmId AND u.offline = :offline AND u.lastSessionRefresh < :lastSessionRefresh)

model/jpa/src/main/resources/META-INF/queries-mariadb.properties

@@ -0,0 +1,9 @@
+# properties file to define queries that can be overloaded for
+# this specific database. Queries are defined in the form:
+#     name[type]=sql
+# type can be native (for native queries) or jpql (jpql syntax)
+# if no type is defined jpql is the default
+
+deleteExpiredClientSessions[native]=delete c from OFFLINE_CLIENT_SESSION c join OFFLINE_USER_SESSION u \
+ where c.USER_SESSION_ID = u.USER_SESSION_ID and u.REALM_ID = :realmId and u.OFFLINE_FLAG = :offline \
+ and u.LAST_SESSION_REFRESH < :lastSessionRefresh

model/jpa/src/main/resources/META-INF/queries-mysql.properties

@@ -0,0 +1,9 @@
+# properties file to define queries that can be overloaded for
+# this specific database. Queries are defined in the form:
+#     name[type]=sql
+# type can be native (for native queries) or jpql (jpql syntax)
+# if no type is defined jpql is the default
+
+deleteExpiredClientSessions[native]=delete c from OFFLINE_CLIENT_SESSION c join OFFLINE_USER_SESSION u \
+ where c.USER_SESSION_ID = u.USER_SESSION_ID and u.REALM_ID = :realmId and u.OFFLINE_FLAG = :offline \
+ and u.LAST_SESSION_REFRESH < :lastSessionRefresh

quarkus/runtime/src/main/java/org/keycloak/connections/jpa/QuarkusJpaConnectionProviderFactory.java

@@ -43,12 +43,16 @@ import javax.transaction.Transaction;
 
 import com.fasterxml.jackson.core.type.TypeReference;
 import io.quarkus.runtime.Quarkus;
+import liquibase.Liquibase;
+import liquibase.exception.LiquibaseException;
 import org.hibernate.internal.SessionFactoryImpl;
 import org.jboss.logging.Logger;
 import org.keycloak.Config;
 import org.keycloak.ServerStartupError;
 import org.keycloak.common.Version;
 import org.keycloak.connections.jpa.updater.JpaUpdaterProvider;
+import org.keycloak.connections.jpa.updater.liquibase.conn.LiquibaseConnectionProvider;
+import org.keycloak.connections.jpa.util.JpaUtils;
 import org.keycloak.exportimport.ExportImportManager;
 import org.keycloak.migration.MigrationModelManager;
 import org.keycloak.migration.ModelVersion;
@@ -110,6 +114,20 @@ public final class QuarkusJpaConnectionProviderFactory implements JpaConnectionP
         this.config = config;
     }
 
+    private void addSpecificNamedQueries(KeycloakSession session, Connection connection) {
+        LiquibaseConnectionProvider liquibaseProvider = session.getProvider(LiquibaseConnectionProvider.class);
+        EntityManager em = null;
+        try {
+            Liquibase liquibase = liquibaseProvider.getLiquibase(connection, this.getSchema());
+            em = createEntityManager(session);
+            JpaUtils.addSpecificNamedQueries(em, liquibase.getDatabase().getShortName());
+        } catch (LiquibaseException e) {
+            throw new IllegalStateException(e);
+        } finally {
+            JpaUtils.closeEntityManager(em);
+        }
+    }
+
     @Override
     public void postInit(KeycloakSessionFactory factory) {
         this.factory = factory;
@@ -126,6 +144,7 @@ public final class QuarkusJpaConnectionProviderFactory implements JpaConnectionP
 
         try (Connection connection = getConnection()) {
             createOperationalInfo(connection);
+            addSpecificNamedQueries(session, connection);
             initSchema = createOrUpdateSchema(getSchema(), connection, session);
         } catch (SQLException cause) {
             throw new RuntimeException("Failed to update database.", cause);