libre.sh/keycloak-scim
2
0
Fork 0
Code Issues 14 Pull requests Releases Packages Activity Actions 6

Avoid optimistic locking queries on CockroachDB to avoid rolling back transactions

Browse source 
Closes #16976
This commit is contained in:
Alexander Schwartz 2023-05-20 18:09:47 +02:00 committed by Hynek Mlnařík
parent d7d6b83bd6
commit 23683970bb
3 changed files with 20 additions and 49 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -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);
eventListenerRegistry.appendListeners(EventType.PRE_UPDATE, JpaOptimisticLockingListener.INSTANCE);
eventListenerRegistry.appendListeners(EventType.PRE_DELETE, JpaOptimisticLockingListener.INSTANCE);
if (metadata.getDatabase().getDialect() instanceof CockroachDialect) {
// CockroachDB will always use serializable transactions, therefore no optimistic locking is necessary
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);

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

@ -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);
}
}
}

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

@ -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