Skip to content

Instantly share code, notes, and snippets.

@serac
Created June 24, 2021 12:20
Show Gist options
  • Save serac/530d578590f07857ae6bdd8e6d8b86ef to your computer and use it in GitHub Desktop.
Save serac/530d578590f07857ae6bdd8e6d8b86ef to your computer and use it in GitHub Desktop.
Custom Spring JPA transaction manager with support for rollback override
/*
* See LICENSE for licensing and NOTICE for copyright.
*/
package edu.vt.middleware.ed.support.spring.tx;
import java.lang.reflect.Field;
import java.util.concurrent.ConcurrentHashMap;
import javax.persistence.EntityManager;
import javax.persistence.EntityManagerFactory;
import edu.vt.middleware.core.annotation.Trivial;
import org.springframework.orm.jpa.EntityManagerHolder;
import org.springframework.orm.jpa.JpaTransactionManager;
import org.springframework.transaction.TransactionDefinition;
import org.springframework.transaction.support.DefaultTransactionStatus;
import org.springframework.transaction.support.TransactionSynchronizationManager;
/**
* Customization of the Spring default JPA transaction manager that uses a custom status class to allow programmatic
* transaction management to override the default rollback of transactions initiated by declarative transaction
* management components, thereby affording additional control of when and how rollback occurs.
*
* @author Marvin S. Addison
* @see CustomTransactionStatus
*/
public class CustomJpaTransactionManager extends JpaTransactionManager {
/** Map used to track tx status objects that have been created for a particular EntityManager/Session. */
private final ConcurrentHashMap<EntityManager, CustomTransactionStatus> statusMap = new ConcurrentHashMap<>();
@Trivial
public CustomJpaTransactionManager() {
super();
}
@Trivial
public CustomJpaTransactionManager(final EntityManagerFactory emf) {
super(emf);
}
@Override
protected DefaultTransactionStatus newTransactionStatus(
final TransactionDefinition definition,
final Object transaction,
final boolean newTransaction,
final boolean newSync,
final boolean debug,
final Object suspendedResources)
{
final boolean actualNewSync = newSync && !TransactionSynchronizationManager.isSynchronizationActive();
final CustomTransactionStatus status = new CustomTransactionStatus(
transaction, newTransaction, actualNewSync, definition.isReadOnly(), debug, suspendedResources);
final EntityManager em = getEntityManager(transaction);
if (em != null) {
final CustomTransactionStatus existing = statusMap.get(em);
if (existing != null) {
status.setSuppressRollback(existing.isSuppressRollback());
}
}
return status;
}
/**
* We override this method because it's called <em>after</em> the transaction is started and the entity
* manager/holder is initialized, which is a requirement to set <code>statusMap</code> entries.
*
* @param status Transaction status.
* @param definition Transaction definition.
*/
@Override
protected void prepareSynchronization(final DefaultTransactionStatus status, final TransactionDefinition definition) {
super.prepareSynchronization(status, definition);
final EntityManager em = getEntityManager(status.getTransaction());
if (em != null) {
// prepareSynchronization is only invoked on a new transaction so it's safe to create a new map entry here
statusMap.put(em, (CustomTransactionStatus) status);
}
}
@Override
protected void doRollback(final DefaultTransactionStatus status) {
if (status instanceof CustomTransactionStatus && ((CustomTransactionStatus) status).isSuppressRollback()) {
return;
}
super.doRollback(status);
}
@Override
protected void doSetRollbackOnly(final DefaultTransactionStatus status) {
if (status instanceof CustomTransactionStatus && ((CustomTransactionStatus) status).isSuppressRollback()) {
return;
}
super.doSetRollbackOnly(status);
}
@Override
protected void doCleanupAfterCompletion(final Object transaction) {
super.doCleanupAfterCompletion(transaction);
final EntityManager em = getEntityManager(transaction);
if (em != null) {
statusMap.remove(em);
}
}
/**
* Tries to get the entity manager bound to the given transaction.
*
* @param transaction JPA platform-specific transaction object.
*
* @return Entity manager or null if none is found.
*/
private EntityManager getEntityManager(final Object transaction) {
try {
final Field field = transaction.getClass().getDeclaredField("entityManagerHolder");
if (!field.trySetAccessible()) {
throw new RuntimeException("Cannot make entityManagerHolder field on transaction object accessible");
}
final EntityManagerHolder holder = (EntityManagerHolder) field.get(transaction);
if (holder != null) {
return holder.getEntityManager();
}
return null;
} catch (NoSuchFieldException | IllegalAccessException e) {
throw new RuntimeException("Error accessing entity manager bound to transaction", e);
}
}
}
/*
* See LICENSE for licensing and NOTICE for copyright.
*/
package edu.vt.middleware.ed.support.spring.tx;
import edu.vt.middleware.core.annotation.Trivial;
import org.springframework.transaction.support.DefaultTransactionStatus;
/**
* Extension of {@link DefaultTransactionStatus} that provides an additional flag that can be set by programmatic
* transaction management components to suppress rollback of transactions initiated by declarative transactional
* components, i.e. {@code @Transactional}, thereby affording more control of rollback behavior.
*
* @author Marvin S. Addison
* @see CustomJpaTransactionManager
*/
public class CustomTransactionStatus extends DefaultTransactionStatus {
/** Flag that can be set to suppress rollback of declarative transactions. */
private boolean suppressRollback;
@Trivial
public CustomTransactionStatus(
final Object transaction,
final boolean newTransaction,
final boolean newSynchronization,
final boolean readOnly,
final boolean debug,
final Object suspendedResources)
{
super(transaction, newTransaction, newSynchronization, readOnly, debug, suspendedResources);
}
@Trivial
public boolean isSuppressRollback() {
return suppressRollback;
}
@Trivial
public void setSuppressRollback(final boolean block) {
this.suppressRollback = block;
}
@Override
public void setRollbackOnly() {
if (!suppressRollback) {
super.setRollbackOnly();
}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment