Skip to content

Commit 0f147cc

Browse files
sguilhenahus1
authored and
hmlnarik
committed
Enlist JPA transaction in JpaMapStorageProvider.getStorage
Closes keycloak#11230 Co-authored-by: Alexander Schwartz <[email protected]>
1 parent 0ce5dfc commit 0f147cc

File tree

7 files changed

+245
-29
lines changed

7 files changed

+245
-29
lines changed

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

Lines changed: 3 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -166,24 +166,17 @@ public long delete(QueryParameters<M> queryParameters) {
166166

167167
@Override
168168
public void begin() {
169-
logger.tracef("tx %d: begin", hashCode());
170-
em.getTransaction().begin();
169+
// no-op: rely on JPA transaction enlisted by the JPA storage provider.
171170
}
172171

173172
@Override
174173
public void commit() {
175-
try {
176-
logger.tracef("tx %d: commit", hashCode());
177-
em.getTransaction().commit();
178-
} catch (PersistenceException e) {
179-
throw PersistenceExceptionConverter.convert(e.getCause() != null ? e.getCause() : e);
180-
}
174+
// no-op: rely on JPA transaction enlisted by the JPA storage provider.
181175
}
182176

183177
@Override
184178
public void rollback() {
185-
logger.tracef("tx %d: rollback", hashCode());
186-
em.getTransaction().rollback();
179+
// no-op: rely on JPA transaction enlisted by the JPA storage provider.
187180
}
188181

189182
@Override

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

Lines changed: 12 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,8 @@
1818

1919
import javax.persistence.EntityManager;
2020

21-
import org.jboss.logging.Logger;
2221
import org.keycloak.models.KeycloakSession;
23-
import org.keycloak.models.UserLoginFailureModel;
22+
import org.keycloak.models.KeycloakTransaction;
2423
import org.keycloak.models.map.common.AbstractEntity;
2524
import org.keycloak.models.map.storage.MapKeycloakTransaction;
2625
import org.keycloak.models.map.storage.MapStorage;
@@ -29,18 +28,16 @@
2928

3029
public class JpaMapStorageProvider implements MapStorageProvider {
3130

32-
private static final Logger logger = Logger.getLogger(JpaMapStorageProvider.class);
33-
34-
private final String SESSION_TX_PREFIX = "jpa-map-tx-";
35-
3631
private final JpaMapStorageProviderFactory factory;
3732
private final KeycloakSession session;
3833
private final EntityManager em;
34+
private final String sessionTxKey;
3935

40-
public JpaMapStorageProvider(JpaMapStorageProviderFactory factory, KeycloakSession session, EntityManager em) {
36+
public JpaMapStorageProvider(JpaMapStorageProviderFactory factory, KeycloakSession session, EntityManager em, String sessionTxKey) {
4137
this.factory = factory;
4238
this.session = session;
4339
this.em = em;
40+
this.sessionTxKey = sessionTxKey;
4441
}
4542

4643
@Override
@@ -51,21 +48,19 @@ public void close() {
5148
@Override
5249
@SuppressWarnings("unchecked")
5350
public <V extends AbstractEntity, M> MapStorage<V, M> getStorage(Class<M> modelType, Flag... flags) {
54-
if (modelType == UserLoginFailureModel.class) {
55-
logger.warn("Enabling JPA storage for user login failures will result in testsuite failures until GHI #11230 is resolved");
51+
// validate and update the schema for the storage.
52+
this.factory.validateAndUpdateSchema(this.session, modelType);
53+
// create the JPA transaction and enlist it if needed.
54+
if (session.getAttribute(this.sessionTxKey) == null) {
55+
KeycloakTransaction jpaTransaction = new JpaTransactionWrapper(em.getTransaction());
56+
session.getTransactionManager().enlist(jpaTransaction);
57+
session.setAttribute(this.sessionTxKey, jpaTransaction);
5658
}
57-
factory.validateAndUpdateSchema(session, modelType);
5859
return new MapStorage<V, M>() {
5960
@Override
6061
public MapKeycloakTransaction<V, M> createTransaction(KeycloakSession session) {
61-
MapKeycloakTransaction<V, M> sessionTx = session.getAttribute(SESSION_TX_PREFIX + modelType.hashCode(), MapKeycloakTransaction.class);
62-
if (sessionTx == null) {
63-
sessionTx = factory.createTransaction(modelType, em);
64-
session.setAttribute(SESSION_TX_PREFIX + modelType.hashCode(), sessionTx);
65-
}
66-
return sessionTx;
62+
return factory.createTransaction(modelType, em);
6763
}
6864
};
6965
}
70-
7166
}

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

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import java.util.Map;
2929
import java.util.Set;
3030
import java.util.concurrent.ConcurrentHashMap;
31+
import java.util.concurrent.atomic.AtomicInteger;
3132
import java.util.function.Function;
3233

3334
import javax.naming.InitialContext;
@@ -94,6 +95,7 @@
9495
import org.keycloak.models.map.storage.jpa.clientscope.entity.JpaClientScopeEntity;
9596
import org.keycloak.models.map.storage.jpa.group.JpaGroupMapKeycloakTransaction;
9697
import org.keycloak.models.map.storage.jpa.group.entity.JpaGroupEntity;
98+
import org.keycloak.models.map.storage.jpa.hibernate.listeners.JpaAutoFlushListener;
9799
import org.keycloak.models.map.storage.jpa.hibernate.listeners.JpaEntityVersionListener;
98100
import org.keycloak.models.map.storage.jpa.hibernate.listeners.JpaOptimisticLockingListener;
99101
import org.keycloak.models.map.storage.jpa.loginFailure.JpaUserLoginFailureMapKeycloakTransaction;
@@ -114,13 +116,17 @@ public class JpaMapStorageProviderFactory implements
114116
EnvironmentDependentProviderFactory {
115117

116118
public static final String PROVIDER_ID = "jpa-map-storage";
119+
private static final String SESSION_TX_PREFIX = "jpa-map-tx-";
120+
private static final AtomicInteger ENUMERATOR = new AtomicInteger(0);
117121
private static final Logger logger = Logger.getLogger(JpaMapStorageProviderFactory.class);
118122

119123
public static final String HIBERNATE_DEFAULT_SCHEMA = "hibernate.default_schema";
120124

121125
private volatile EntityManagerFactory emf;
122126
private final Set<Class<?>> validatedModels = ConcurrentHashMap.newKeySet();
123127
private Config.Scope config;
128+
private final String sessionProviderKey;
129+
private final String sessionTxKey;
124130

125131
public final static DeepCloner CLONER = new DeepCloner.Builder()
126132
//auth-session
@@ -163,14 +169,26 @@ public class JpaMapStorageProviderFactory implements
163169
MODEL_TO_TX.put(UserLoginFailureModel.class, JpaUserLoginFailureMapKeycloakTransaction::new);
164170
}
165171

172+
public JpaMapStorageProviderFactory() {
173+
int index = ENUMERATOR.getAndIncrement();
174+
this.sessionProviderKey = PROVIDER_ID + "-" + index;
175+
this.sessionTxKey = SESSION_TX_PREFIX + index;
176+
}
177+
166178
public MapKeycloakTransaction createTransaction(Class<?> modelType, EntityManager em) {
167179
return MODEL_TO_TX.get(modelType).apply(em);
168180
}
169181

170182
@Override
171183
public MapStorageProvider create(KeycloakSession session) {
172184
lazyInit();
173-
return new JpaMapStorageProvider(this, session, emf.createEntityManager());
185+
// check the session for a cached provider before creating a new one.
186+
JpaMapStorageProvider provider = session.getAttribute(this.sessionProviderKey, JpaMapStorageProvider.class);
187+
if (provider == null) {
188+
provider = new JpaMapStorageProvider(this, session, emf.createEntityManager(), this.sessionTxKey);
189+
session.setAttribute(this.sessionProviderKey, provider);
190+
}
191+
return provider;
174192
}
175193

176194
@Override
@@ -259,6 +277,9 @@ public void integrate(Metadata metadata, SessionFactoryImplementor sessionFactor
259277
eventListenerRegistry.appendListeners(EventType.PRE_INSERT, JpaEntityVersionListener.INSTANCE);
260278
eventListenerRegistry.appendListeners(EventType.PRE_UPDATE, JpaEntityVersionListener.INSTANCE);
261279
eventListenerRegistry.appendListeners(EventType.PRE_DELETE, JpaEntityVersionListener.INSTANCE);
280+
281+
// replace auto-flush listener
282+
eventListenerRegistry.setListeners(EventType.AUTO_FLUSH, JpaAutoFlushListener.INSTANCE);
262283
}
263284

264285
@Override

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

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
/*
2+
* Copyright 2022 Red Hat, Inc. and/or its affiliates
3+
* and other contributors as indicated by the @author tags.
4+
*
5+
* Licensed under the Apache License, Version 2.0 (the "License");
6+
* you may not use this file except in compliance with the License.
7+
* You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
package org.keycloak.models.map.storage.jpa;
18+
19+
import javax.persistence.EntityTransaction;
20+
import javax.persistence.PersistenceException;
21+
22+
import org.jboss.logging.Logger;
23+
import org.keycloak.connections.jpa.PersistenceExceptionConverter;
24+
import org.keycloak.models.KeycloakTransaction;
25+
26+
/**
27+
* Wraps an {@link EntityTransaction} as a {@link KeycloakTransaction} so it can be enlisted in {@link org.keycloak.models.KeycloakTransactionManager}.
28+
*
29+
* @author <a href="mailto:[email protected]">Stefan Guilhen</a>
30+
*/
31+
public class JpaTransactionWrapper implements KeycloakTransaction {
32+
33+
private static final Logger logger = Logger.getLogger(JpaTransactionWrapper.class);
34+
35+
private final EntityTransaction transaction;
36+
37+
public JpaTransactionWrapper(EntityTransaction transaction) {
38+
this.transaction = transaction;
39+
}
40+
41+
@Override
42+
public void begin() {
43+
logger.tracef("tx %d: begin", hashCode());
44+
this.transaction.begin();
45+
}
46+
47+
@Override
48+
public void commit() {
49+
try {
50+
logger.tracef("tx %d: commit", hashCode());
51+
this.transaction.commit();
52+
} catch(PersistenceException pe) {
53+
throw PersistenceExceptionConverter.convert(pe.getCause() != null ? pe.getCause() : pe);
54+
}
55+
}
56+
57+
@Override
58+
public void rollback() {
59+
logger.tracef("tx %d: rollback", hashCode());
60+
this.transaction.rollback();
61+
}
62+
63+
@Override
64+
public void setRollbackOnly() {
65+
this.transaction.setRollbackOnly();
66+
}
67+
68+
@Override
69+
public boolean getRollbackOnly() {
70+
return this.transaction.getRollbackOnly();
71+
}
72+
73+
@Override
74+
public boolean isActive() {
75+
return this.transaction.isActive();
76+
}
77+
}

‎model/map-jpa/src/main/java/org/keycloak/models/map/storage/jpa/hibernate/listeners/JpaAutoFlushListener.java‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
/*
2+
* Copyright 2022. Red Hat, Inc. and/or its affiliates
3+
* and other contributors as indicated by the @author tags.
4+
*
5+
* Licensed under the Apache License, Version 2.0 (the "License");
6+
* you may not use this file except in compliance with the License.
7+
* You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
18+
package org.keycloak.models.map.storage.jpa.hibernate.listeners;
19+
20+
import org.hibernate.FlushMode;
21+
import org.hibernate.HibernateException;
22+
import org.hibernate.engine.spi.ActionQueue;
23+
import org.hibernate.event.internal.DefaultAutoFlushEventListener;
24+
import org.hibernate.event.spi.AutoFlushEvent;
25+
import org.hibernate.event.spi.EventSource;
26+
import org.hibernate.internal.CoreMessageLogger;
27+
import org.jboss.logging.Logger;
28+
29+
/**
30+
* Extends Hibernate's {@link DefaultAutoFlushEventListener} to always flush queued inserts to allow correct handling
31+
* of orphans of that entities in the same transactions.
32+
* If they wouldn't be flushed, they won't be orphaned (at least not in Hibernate 5.3.24.Final).
33+
* This class copies over all functionality of the base class that can't be overwritten via inheritance.
34+
*/
35+
public class JpaAutoFlushListener extends DefaultAutoFlushEventListener {
36+
37+
public static final JpaAutoFlushListener INSTANCE = new JpaAutoFlushListener();
38+
39+
private static final CoreMessageLogger LOG = Logger.getMessageLogger(CoreMessageLogger.class, DefaultAutoFlushEventListener.class.getName());
40+
41+
public void onAutoFlush(AutoFlushEvent event) throws HibernateException {
42+
final EventSource source = event.getSession();
43+
try {
44+
source.getEventListenerManager().partialFlushStart();
45+
46+
if (flushMightBeNeeded(source)) {
47+
// Need to get the number of collection removals before flushing to executions
48+
// (because flushing to executions can add collection removal actions to the action queue).
49+
final ActionQueue actionQueue = source.getActionQueue();
50+
final int oldSize = actionQueue.numberOfCollectionRemovals();
51+
flushEverythingToExecutions(event);
52+
if (flushIsReallyNeeded(event, source)) {
53+
LOG.trace("Need to execute flush");
54+
event.setFlushRequired(true);
55+
56+
// note: performExecutions() clears all collectionXxxxtion
57+
// collections (the collection actions) in the session
58+
performExecutions(source);
59+
postFlush(source);
60+
61+
postPostFlush(source);
62+
63+
if (source.getFactory().getStatistics().isStatisticsEnabled()) {
64+
source.getFactory().getStatistics().flush();
65+
}
66+
} else {
67+
LOG.trace("Don't need to execute flush");
68+
event.setFlushRequired(false);
69+
actionQueue.clearFromFlushNeededCheck(oldSize);
70+
}
71+
}
72+
} finally {
73+
source.getEventListenerManager().partialFlushEnd(
74+
event.getNumberOfEntitiesProcessed(),
75+
event.getNumberOfEntitiesProcessed()
76+
);
77+
}
78+
}
79+
80+
private boolean flushIsReallyNeeded(AutoFlushEvent event, final EventSource source) {
81+
return source.getHibernateFlushMode() == FlushMode.ALWAYS
82+
// START OF FIX for auto-flush-mode on inserts that might later be deleted in same transaction
83+
|| source.getActionQueue().numberOfInsertions() > 0
84+
// END OF FIX
85+
|| source.getActionQueue().areTablesToBeUpdated(event.getQuerySpaces());
86+
}
87+
88+
private boolean flushMightBeNeeded(final EventSource source) {
89+
return !source.getHibernateFlushMode().lessThan(FlushMode.AUTO)
90+
&& source.getDontFlushFromFind() == 0
91+
&& (source.getPersistenceContext().getNumberOfManagedEntities() > 0 ||
92+
source.getPersistenceContext().getCollectionEntries().size() > 0);
93+
}
94+
95+
}

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

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,10 @@
1717

1818
package org.keycloak.models.map.storage.jpa.liquibase.updater;
1919

20+
import liquibase.database.Database;
21+
import liquibase.database.DatabaseFactory;
22+
import liquibase.database.core.CockroachDatabase;
23+
import liquibase.database.jvm.JdbcConnection;
2024
import org.keycloak.models.map.storage.jpa.liquibase.connection.MapLiquibaseConnectionProvider;
2125
import java.io.File;
2226
import java.io.FileWriter;
@@ -186,7 +190,9 @@ private Liquibase getLiquibase(Class modelType, Connection connection, String de
186190
if (modelName == null) {
187191
throw new IllegalStateException("Cannot find changlelog for modelClass " + modelType.getName());
188192
}
189-
String changelog = "META-INF/jpa-" + modelName + "-changelog.xml";
193+
Database database = DatabaseFactory.getInstance().findCorrectDatabaseImplementation(new JdbcConnection(connection));
194+
// if database is cockroachdb, use the aggregate changelog (see GHI #11230).
195+
String changelog = database instanceof CockroachDatabase ? "META-INF/jpa-aggregate-changelog.xml" : "META-INF/jpa-" + modelName + "-changelog.xml";
190196
return liquibaseProvider.getLiquibaseForCustomUpdate(connection, defaultSchema, changelog, this.getClass().getClassLoader(), "databasechangelog");
191197
}
192198

‎model/map-jpa/src/main/resources/META-INF/jpa-aggregate-changelog.xml‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
<?xml version="1.0" encoding="UTF-8"?>
2+
<!--
3+
Copyright 2022 Red Hat, Inc. and/or its affiliates
4+
and other contributors as indicated by the @author tags.
5+
6+
Licensed under the Apache License, Version 2.0 (the "License");
7+
you may not use this file except in compliance with the License.
8+
You may obtain a copy of the License at
9+
10+
http://www.apache.org/licenses/LICENSE-2.0
11+
12+
Unless required by applicable law or agreed to in writing, software
13+
distributed under the License is distributed on an "AS IS" BASIS,
14+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
See the License for the specific language governing permissions and
16+
limitations under the License.
17+
-->
18+
19+
20+
<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">
21+
<!-- aggregate changelog for cockroachdb -->
22+
<include file="META-INF/jpa-auth-sessions-changelog.xml"/>
23+
<include file="META-INF/jpa-client-scopes-changelog.xml"/>
24+
<include file="META-INF/jpa-clients-changelog.xml"/>
25+
<include file="META-INF/jpa-groups-changelog.xml"/>
26+
<include file="META-INF/jpa-realms-changelog.xml"/>
27+
<include file="META-INF/jpa-roles-changelog.xml"/>
28+
<include file="META-INF/jpa-user-login-failures-changelog.xml"/>
29+
</databaseChangeLog>

0 commit comments

Comments
 (0)