libre.sh/keycloak-scim
3
0
Fork 0
Code Issues 15 Pull requests Releases Packages Activity Actions

Enlist JPA transaction in JpaMapStorageProvider.getStorage

Browse source 
Closes #11230

Co-authored-by: Alexander Schwartz <aschwart@redhat.com>
This commit is contained in:
Stefan Guilhen 2022-04-13 23:36:03 -03:00 committed by Hynek Mlnařík
parent 0ce5dfc09c
commit 0f147ccdc0
7 changed files with 245 additions and 29 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
model/map-jpa/src/main/java/org/keycloak/models/map/storage/jpa/JpaMapKeycloakTransaction.java
View file

@ -166,24 +166,17 @@ public abstract class JpaMapKeycloakTransaction<RE extends JpaRootEntity, E exte
@Override
public void begin() {
logger.tracef("tx %d: begin", hashCode());
em.getTransaction().begin();
// no-op: rely on JPA transaction enlisted by the JPA storage provider.
}
@Override
public void commit() {
try {
logger.tracef("tx %d: commit", hashCode());
em.getTransaction().commit();
} catch (PersistenceException e) {
throw PersistenceExceptionConverter.convert(e.getCause() != null ? e.getCause() : e);
}
// no-op: rely on JPA transaction enlisted by the JPA storage provider.
}
@Override
public void rollback() {
logger.tracef("tx %d: rollback", hashCode());
em.getTransaction().rollback();
// no-op: rely on JPA transaction enlisted by the JPA storage provider.
}
@Override

29
model/map-jpa/src/main/java/org/keycloak/models/map/storage/jpa/JpaMapStorageProvider.java
View file

@ -18,9 +18,8 @@ package org.keycloak.models.map.storage.jpa;
import javax.persistence.EntityManager;
import org.jboss.logging.Logger;
import org.keycloak.models.KeycloakSession;
import org.keycloak.models.UserLoginFailureModel;
import org.keycloak.models.KeycloakTransaction;
import org.keycloak.models.map.common.AbstractEntity;
import org.keycloak.models.map.storage.MapKeycloakTransaction;
import org.keycloak.models.map.storage.MapStorage;
@ -29,18 +28,16 @@ import org.keycloak.models.map.storage.MapStorageProviderFactory.Flag;
public class JpaMapStorageProvider implements MapStorageProvider {
private static final Logger logger = Logger.getLogger(JpaMapStorageProvider.class);
private final String SESSION_TX_PREFIX = "jpa-map-tx-";
private final JpaMapStorageProviderFactory factory;
private final KeycloakSession session;
private final EntityManager em;
private final String sessionTxKey;
public JpaMapStorageProvider(JpaMapStorageProviderFactory factory, KeycloakSession session, EntityManager em) {
public JpaMapStorageProvider(JpaMapStorageProviderFactory factory, KeycloakSession session, EntityManager em, String sessionTxKey) {
this.factory = factory;
this.session = session;
this.em = em;
this.sessionTxKey = sessionTxKey;
}
@Override
@ -51,21 +48,19 @@ public class JpaMapStorageProvider implements MapStorageProvider {
@Override
@SuppressWarnings("unchecked")
public <V extends AbstractEntity, M> MapStorage<V, M> getStorage(Class<M> modelType, Flag... flags) {
if (modelType == UserLoginFailureModel.class) {
logger.warn("Enabling JPA storage for user login failures will result in testsuite failures until GHI #11230 is resolved");
// validate and update the schema for the storage.
this.factory.validateAndUpdateSchema(this.session, modelType);
// create the JPA transaction and enlist it if needed.
if (session.getAttribute(this.sessionTxKey) == null) {
KeycloakTransaction jpaTransaction = new JpaTransactionWrapper(em.getTransaction());
session.getTransactionManager().enlist(jpaTransaction);
session.setAttribute(this.sessionTxKey, jpaTransaction);
}
factory.validateAndUpdateSchema(session, modelType);
return new MapStorage<V, M>() {
@Override
public MapKeycloakTransaction<V, M> createTransaction(KeycloakSession session) {
MapKeycloakTransaction<V, M> sessionTx = session.getAttribute(SESSION_TX_PREFIX + modelType.hashCode(), MapKeycloakTransaction.class);
if (sessionTx == null) {
sessionTx = factory.createTransaction(modelType, em);
session.setAttribute(SESSION_TX_PREFIX + modelType.hashCode(), sessionTx);
}
return sessionTx;
return factory.createTransaction(modelType, em);
}
};
}
}

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

@ -28,6 +28,7 @@ import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Function;
import javax.naming.InitialContext;
@ -94,6 +95,7 @@ import org.keycloak.models.map.storage.jpa.clientscope.JpaClientScopeMapKeycloak
import org.keycloak.models.map.storage.jpa.clientscope.entity.JpaClientScopeEntity;
import org.keycloak.models.map.storage.jpa.group.JpaGroupMapKeycloakTransaction;
import org.keycloak.models.map.storage.jpa.group.entity.JpaGroupEntity;
import org.keycloak.models.map.storage.jpa.hibernate.listeners.JpaAutoFlushListener;
import org.keycloak.models.map.storage.jpa.hibernate.listeners.JpaEntityVersionListener;
import org.keycloak.models.map.storage.jpa.hibernate.listeners.JpaOptimisticLockingListener;
import org.keycloak.models.map.storage.jpa.loginFailure.JpaUserLoginFailureMapKeycloakTransaction;
@ -114,6 +116,8 @@ public class JpaMapStorageProviderFactory implements
EnvironmentDependentProviderFactory {
public static final String PROVIDER_ID = "jpa-map-storage";
private static final String SESSION_TX_PREFIX = "jpa-map-tx-";
private static final AtomicInteger ENUMERATOR = new AtomicInteger(0);
private static final Logger logger = Logger.getLogger(JpaMapStorageProviderFactory.class);
public static final String HIBERNATE_DEFAULT_SCHEMA = "hibernate.default_schema";
@ -121,6 +125,8 @@ public class JpaMapStorageProviderFactory implements
private volatile EntityManagerFactory emf;
private final Set<Class<?>> validatedModels = ConcurrentHashMap.newKeySet();
private Config.Scope config;
private final String sessionProviderKey;
private final String sessionTxKey;
public final static DeepCloner CLONER = new DeepCloner.Builder()
//auth-session
@ -163,6 +169,12 @@ public class JpaMapStorageProviderFactory implements
MODEL_TO_TX.put(UserLoginFailureModel.class, JpaUserLoginFailureMapKeycloakTransaction::new);
}
public JpaMapStorageProviderFactory() {
int index = ENUMERATOR.getAndIncrement();
this.sessionProviderKey = PROVIDER_ID + "-" + index;
this.sessionTxKey = SESSION_TX_PREFIX + index;
}
public MapKeycloakTransaction createTransaction(Class<?> modelType, EntityManager em) {
return MODEL_TO_TX.get(modelType).apply(em);
}
@ -170,7 +182,13 @@ public class JpaMapStorageProviderFactory implements
@Override
public MapStorageProvider create(KeycloakSession session) {
lazyInit();
return new JpaMapStorageProvider(this, session, emf.createEntityManager());
// check the session for a cached provider before creating a new one.
JpaMapStorageProvider provider = session.getAttribute(this.sessionProviderKey, JpaMapStorageProvider.class);
if (provider == null) {
provider = new JpaMapStorageProvider(this, session, emf.createEntityManager(), this.sessionTxKey);
session.setAttribute(this.sessionProviderKey, provider);
}
return provider;
}
@Override
@ -259,6 +277,9 @@ public class JpaMapStorageProviderFactory implements
eventListenerRegistry.appendListeners(EventType.PRE_INSERT, JpaEntityVersionListener.INSTANCE);
eventListenerRegistry.appendListeners(EventType.PRE_UPDATE, JpaEntityVersionListener.INSTANCE);
eventListenerRegistry.appendListeners(EventType.PRE_DELETE, JpaEntityVersionListener.INSTANCE);
// replace auto-flush listener
eventListenerRegistry.setListeners(EventType.AUTO_FLUSH, JpaAutoFlushListener.INSTANCE);
}
@Override

77
model/map-jpa/src/main/java/org/keycloak/models/map/storage/jpa/JpaTransactionWrapper.java Normal file
View file

@ -0,0 +1,77 @@
/*
* Copyright 2022 Red Hat, Inc. and/or its affiliates
* and other contributors as indicated by the @author tags.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.keycloak.models.map.storage.jpa;
import javax.persistence.EntityTransaction;
import javax.persistence.PersistenceException;
import org.jboss.logging.Logger;
import org.keycloak.connections.jpa.PersistenceExceptionConverter;
import org.keycloak.models.KeycloakTransaction;
/**
* Wraps an {@link EntityTransaction} as a {@link KeycloakTransaction} so it can be enlisted in {@link org.keycloak.models.KeycloakTransactionManager}.
*
* @author <a href="mailto:sguilhen@redhat.com">Stefan Guilhen</a>
*/
public class JpaTransactionWrapper implements KeycloakTransaction {
private static final Logger logger = Logger.getLogger(JpaTransactionWrapper.class);
private final EntityTransaction transaction;
public JpaTransactionWrapper(EntityTransaction transaction) {
this.transaction = transaction;
}
@Override
public void begin() {
logger.tracef("tx %d: begin", hashCode());
this.transaction.begin();
}
@Override
public void commit() {
try {
logger.tracef("tx %d: commit", hashCode());
this.transaction.commit();
} catch(PersistenceException pe) {
throw PersistenceExceptionConverter.convert(pe.getCause() != null ? pe.getCause() : pe);
}
}
@Override
public void rollback() {
logger.tracef("tx %d: rollback", hashCode());
this.transaction.rollback();
}
@Override
public void setRollbackOnly() {
this.transaction.setRollbackOnly();
}
@Override
public boolean getRollbackOnly() {
return this.transaction.getRollbackOnly();
}
@Override
public boolean isActive() {
return this.transaction.isActive();
}
}

95
model/map-jpa/src/main/java/org/keycloak/models/map/storage/jpa/hibernate/listeners/JpaAutoFlushListener.java Normal file
View file

@ -0,0 +1,95 @@
/*
* Copyright 2022. Red Hat, Inc. and/or its affiliates
* and other contributors as indicated by the @author tags.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.keycloak.models.map.storage.jpa.hibernate.listeners;
import org.hibernate.FlushMode;
import org.hibernate.HibernateException;
import org.hibernate.engine.spi.ActionQueue;
import org.hibernate.event.internal.DefaultAutoFlushEventListener;
import org.hibernate.event.spi.AutoFlushEvent;
import org.hibernate.event.spi.EventSource;
import org.hibernate.internal.CoreMessageLogger;
import org.jboss.logging.Logger;
/**
* Extends Hibernate's {@link DefaultAutoFlushEventListener} to always flush queued inserts to allow correct handling
* of orphans of that entities in the same transactions.
* If they wouldn't be flushed, they won't be orphaned (at least not in Hibernate 5.3.24.Final).
* This class copies over all functionality of the base class that can't be overwritten via inheritance.
*/
public class JpaAutoFlushListener extends DefaultAutoFlushEventListener {
public static final JpaAutoFlushListener INSTANCE = new JpaAutoFlushListener();
private static final CoreMessageLogger LOG = Logger.getMessageLogger(CoreMessageLogger.class, DefaultAutoFlushEventListener.class.getName());
public void onAutoFlush(AutoFlushEvent event) throws HibernateException {
final EventSource source = event.getSession();
try {
source.getEventListenerManager().partialFlushStart();
if (flushMightBeNeeded(source)) {
// Need to get the number of collection removals before flushing to executions
// (because flushing to executions can add collection removal actions to the action queue).
final ActionQueue actionQueue = source.getActionQueue();
final int oldSize = actionQueue.numberOfCollectionRemovals();
flushEverythingToExecutions(event);
if (flushIsReallyNeeded(event, source)) {
LOG.trace("Need to execute flush");
event.setFlushRequired(true);
// note: performExecutions() clears all collectionXxxxtion
// collections (the collection actions) in the session
performExecutions(source);
postFlush(source);
postPostFlush(source);
if (source.getFactory().getStatistics().isStatisticsEnabled()) {
source.getFactory().getStatistics().flush();
}
} else {
LOG.trace("Don't need to execute flush");
event.setFlushRequired(false);
actionQueue.clearFromFlushNeededCheck(oldSize);
}
}
} finally {
source.getEventListenerManager().partialFlushEnd(
event.getNumberOfEntitiesProcessed(),
event.getNumberOfEntitiesProcessed()
);
}
}
private boolean flushIsReallyNeeded(AutoFlushEvent event, final EventSource source) {
return source.getHibernateFlushMode() == FlushMode.ALWAYS
// START OF FIX for auto-flush-mode on inserts that might later be deleted in same transaction
|| source.getActionQueue().numberOfInsertions() > 0
// END OF FIX
|| source.getActionQueue().areTablesToBeUpdated(event.getQuerySpaces());
}
private boolean flushMightBeNeeded(final EventSource source) {
return !source.getHibernateFlushMode().lessThan(FlushMode.AUTO)
&& source.getDontFlushFromFind() == 0
&& (source.getPersistenceContext().getNumberOfManagedEntities() > 0 ||
source.getPersistenceContext().getCollectionEntries().size() > 0);
}
}

8
model/map-jpa/src/main/java/org/keycloak/models/map/storage/jpa/liquibase/updater/MapJpaLiquibaseUpdaterProvider.java
View file

@ -17,6 +17,10 @@
package org.keycloak.models.map.storage.jpa.liquibase.updater;
import liquibase.database.Database;
import liquibase.database.DatabaseFactory;
import liquibase.database.core.CockroachDatabase;
import liquibase.database.jvm.JdbcConnection;
import org.keycloak.models.map.storage.jpa.liquibase.connection.MapLiquibaseConnectionProvider;
import java.io.File;
import java.io.FileWriter;
@ -186,7 +190,9 @@ public class MapJpaLiquibaseUpdaterProvider implements MapJpaUpdaterProvider {
if (modelName == null) {
throw new IllegalStateException("Cannot find changlelog for modelClass " + modelType.getName());
}
String changelog = "META-INF/jpa-" + modelName + "-changelog.xml";
Database database = DatabaseFactory.getInstance().findCorrectDatabaseImplementation(new JdbcConnection(connection));
// if database is cockroachdb, use the aggregate changelog (see GHI #11230).
String changelog = database instanceof CockroachDatabase ? "META-INF/jpa-aggregate-changelog.xml" : "META-INF/jpa-" + modelName + "-changelog.xml";
return liquibaseProvider.getLiquibaseForCustomUpdate(connection, defaultSchema, changelog, this.getClass().getClassLoader(), "databasechangelog");
}

29
model/map-jpa/src/main/resources/META-INF/jpa-aggregate-changelog.xml Normal file
View file

@ -0,0 +1,29 @@
<?xml version="1.0" encoding="UTF-8"?>
<!--
Copyright 2022 Red Hat, Inc. and/or its affiliates
and other contributors as indicated by the @author tags.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
-->
<databaseChangeLog xmlns="http://www.liquibase.org/xml/ns/dbchangelog" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://www.liquibase.org/xml/ns/dbchangelog http://www.liquibase.org/xml/ns/dbchangelog/dbchangelog-3.1.xsd">
<!-- aggregate changelog for cockroachdb -->
<include file="META-INF/jpa-auth-sessions-changelog.xml"/>
<include file="META-INF/jpa-client-scopes-changelog.xml"/>
<include file="META-INF/jpa-clients-changelog.xml"/>
<include file="META-INF/jpa-groups-changelog.xml"/>
<include file="META-INF/jpa-realms-changelog.xml"/>
<include file="META-INF/jpa-roles-changelog.xml"/>
<include file="META-INF/jpa-user-login-failures-changelog.xml"/>
</databaseChangeLog>