Custom Spring JPA transaction manager with support for rollback override

CustomJpaTransactionManager.java

/*
 * 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);
    }
  }
}

CustomTransactionStatus.java

/*
 * 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();
    }
  }
}