Avoid optimistic locking queries on CockroachDB to avoid rolling back transactions

Closes #16976

model/map-jpa/src/main/java/org/keycloak/models/map/storage/jpa/EventListenerIntegrator.java

@@ -18,10 +18,13 @@
 package org.keycloak.models.map.storage.jpa;
 
 import org.hibernate.boot.Metadata;
+import org.hibernate.dialect.CockroachDialect;
+import org.hibernate.engine.OptimisticLockStyle;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.event.service.spi.EventListenerRegistry;
 import org.hibernate.event.spi.EventType;
 import org.hibernate.integrator.spi.Integrator;
+import org.hibernate.mapping.RootClass;
 import org.hibernate.service.spi.SessionFactoryServiceRegistry;
 import org.keycloak.models.map.storage.jpa.hibernate.listeners.JpaAutoFlushListener;
 import org.keycloak.models.map.storage.jpa.hibernate.listeners.JpaEntityVersionListener;
@@ -43,9 +46,23 @@ public class EventListenerIntegrator implements Integrator {
             final EventListenerRegistry eventListenerRegistry =
                     sessionFactoryServiceRegistry.getService(EventListenerRegistry.class);
 
-            eventListenerRegistry.appendListeners(EventType.PRE_INSERT, JpaOptimisticLockingListener.INSTANCE);
+            if (metadata.getDatabase().getDialect() instanceof CockroachDialect) {
-            eventListenerRegistry.appendListeners(EventType.PRE_UPDATE, JpaOptimisticLockingListener.INSTANCE);
+                // CockroachDB will always use serializable transactions, therefore no optimistic locking is necessary
-            eventListenerRegistry.appendListeners(EventType.PRE_DELETE, JpaOptimisticLockingListener.INSTANCE);
+                metadata.getEntityBindings().forEach(persistentClass -> {
+                    if (persistentClass instanceof RootClass) {
+                        RootClass root = (RootClass) persistentClass;
+                        root.setOptimisticLockStyle(OptimisticLockStyle.NONE);
+                        root.setVersion(null);
+                        root.setDeclaredVersion(null);
+                    }
+                });
+                // If the version column should be updated with an incrementing number on each change in the future,
+                // implement a listener similar to JpaOptimisticLockingListener to increment it.
+            } else {
+                eventListenerRegistry.appendListeners(EventType.PRE_INSERT, JpaOptimisticLockingListener.INSTANCE);
+                eventListenerRegistry.appendListeners(EventType.PRE_UPDATE, JpaOptimisticLockingListener.INSTANCE);
+                eventListenerRegistry.appendListeners(EventType.PRE_DELETE, JpaOptimisticLockingListener.INSTANCE);
+            }
 
             eventListenerRegistry.appendListeners(EventType.PRE_INSERT, JpaEntityVersionListener.INSTANCE);
             eventListenerRegistry.appendListeners(EventType.PRE_UPDATE, JpaEntityVersionListener.INSTANCE);

model/map-jpa/src/main/java/org/keycloak/models/map/storage/jpa/authSession/JpaRootAuthenticationSessionMapStorage.java

@@ -17,31 +17,22 @@
 package org.keycloak.models.map.storage.jpa.authSession;
 
 import jakarta.persistence.EntityManager;
-import jakarta.persistence.Query;
 import jakarta.persistence.criteria.CriteriaBuilder;
 import jakarta.persistence.criteria.Root;
 import jakarta.persistence.criteria.Selection;
 
-import org.hibernate.Session;
-import org.hibernate.query.NativeQuery;
 import org.keycloak.models.KeycloakSession;
-import org.keycloak.models.ModelException;
 import org.keycloak.models.map.authSession.MapRootAuthenticationSessionEntity;
 import org.keycloak.models.map.authSession.MapRootAuthenticationSessionEntityDelegate;
 import static org.keycloak.models.map.storage.jpa.Constants.CURRENT_SCHEMA_VERSION_AUTH_SESSION;
 
-import org.keycloak.models.map.common.StringKeyConverter;
 import org.keycloak.models.map.storage.jpa.JpaMapStorage;
 import org.keycloak.models.map.storage.jpa.JpaModelCriteriaBuilder;
 import org.keycloak.models.map.storage.jpa.JpaRootEntity;
 import org.keycloak.models.map.storage.jpa.authSession.delegate.JpaRootAuthenticationSessionDelegateProvider;
-import org.keycloak.models.map.storage.jpa.authSession.entity.JpaAuthenticationSessionEntity;
 import org.keycloak.models.map.storage.jpa.authSession.entity.JpaRootAuthenticationSessionEntity;
 import org.keycloak.sessions.RootAuthenticationSessionModel;
 
-import java.sql.Connection;
-import java.util.UUID;
-
 public class JpaRootAuthenticationSessionMapStorage extends JpaMapStorage<JpaRootAuthenticationSessionEntity, MapRootAuthenticationSessionEntity, RootAuthenticationSessionModel> {
 
     public JpaRootAuthenticationSessionMapStorage(KeycloakSession session, EntityManager em) {
@@ -75,39 +66,4 @@ public class JpaRootAuthenticationSessionMapStorage extends JpaMapStorage<JpaRoo
         return new MapRootAuthenticationSessionEntityDelegate(new JpaRootAuthenticationSessionDelegateProvider(original, em));
     }
 
-    @Override
-    public boolean delete(String key) {
-        int isolationLevel = em.unwrap(Session.class).doReturningWork(Connection::getTransactionIsolation);
-        if (isolationLevel == Connection.TRANSACTION_SERIALIZABLE) {
-            // If the isolation level is SERIALIZABLE, there is no need to apply the optimistic locking, as the database with its serializable checks
-            // takes care that no-one has modified or deleted the row sind the transaction started. On CockroachDB, using optimistic locking with the added
-            // version column in a delete-statement will cause a table lock, which will lead to deadlock.
-            // As a workaround, this is using a native query instead, without including the version for optimistic locking.
-            if (key == null) return false;
-            UUID uuid = StringKeyConverter.UUIDKey.INSTANCE.fromStringSafe(key);
-            if (uuid == null) return false;
-            removeFromCache(key);
-            // will throw an exception if the entity doesn't exist in the Hibernate session or in the database.
-            JpaRootAuthenticationSessionEntity rootAuth = em.getReference(JpaRootAuthenticationSessionEntity.class, uuid);
-            // will use cascading delete to all child entities
-            //noinspection JpaQueryApiInspection
-            Query deleteById =
-                    em.createNamedQuery("deleteRootAuthenticationSessionByIdNoOptimisticLocking");
-            deleteById.unwrap(NativeQuery.class).addSynchronizedQuerySpace(JpaRootAuthenticationSessionEntity.TABLE_NAME)
-                    .addSynchronizedQuerySpace(JpaAuthenticationSessionEntity.TABLE_NAME);
-            deleteById.setParameter("id", key);
-            int deleteCount = deleteById.executeUpdate();
-            rootAuth.getAuthenticationSessions().forEach(e -> em.detach(e));
-            em.detach(rootAuth);
-            if (deleteCount == 1) {
-                return true;
-            } else if (deleteCount == 0) {
-                throw new ModelException("Unable to find root authentication session");
-            } else {
-                throw new ModelException("Deleted " + deleteCount + " root authentication session when expecting to delete one");
-            }
-        } else {
-            return super.delete(key);
-        }
-    }
 }

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

@@ -4,5 +4,3 @@
 #     name[type]=sql
 # type can be native (for native queries) or jpql (jpql syntax)
 # if no type is defined jpql is the default
-
-deleteRootAuthenticationSessionByIdNoOptimisticLocking[native]=delete from ${schemaprefix}kc_auth_root_session where id = :id