001/*
002 * Copyright 2012 Atteo.
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 *
008 *      http://www.apache.org/licenses/LICENSE-2.0
009 *
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016package org.atteo.moonshine.jpa;
017
018import javax.inject.Inject;
019import javax.persistence.EntityManager;
020import javax.persistence.EntityManagerFactory;
021import javax.persistence.EntityTransaction;
022import javax.persistence.TransactionRequiredException;
023import javax.persistence.criteria.CriteriaBuilder;
024import javax.persistence.metamodel.Metamodel;
025import javax.transaction.RollbackException;
026import javax.transaction.Synchronization;
027import javax.transaction.SystemException;
028import javax.transaction.Transaction;
029import javax.transaction.TransactionManager;
030
031public class TransactionScopedEntityManager extends DelegatingEntityManager {
032
033    @Inject
034    private EntityManagerFactory factory;
035
036    @Inject
037    private TransactionManager transactionManager;
038
039    private final ThreadLocal<EntityManager> entityManagerHolder = new ThreadLocal<>();
040
041    @Override
042    protected EntityManager getEntityManager() {
043        EntityManager entityManager = entityManagerHolder.get();
044
045        if (entityManager == null) {
046            try {
047                Transaction transaction = transactionManager.getTransaction();
048                if (transaction == null) {
049                    throw new TransactionRequiredException("Not in transaction. Initiate transaction in JTA.");
050                }
051                entityManager = factory.createEntityManager();
052
053                final EntityManager entityManagerToClose = entityManager;
054                entityManagerHolder.set(entityManager);
055                transaction.registerSynchronization(new Synchronization() {
056                    @Override
057                    public void beforeCompletion() {
058                    }
059
060                    @Override
061                    public void afterCompletion(int status) {
062                        if (entityManagerHolder.get() != entityManagerToClose) {
063                            throw new RuntimeException("Synchronization called from different thread");
064                        }
065                        if (entityManagerToClose != null) {
066                            entityManagerToClose.close();
067                        }
068                        entityManagerHolder.set(null);
069                    }
070                });
071            } catch (SystemException | RollbackException e) {
072                throw new RuntimeException(e);
073            }
074        }
075
076        return entityManager;
077    }
078
079    @Override
080    public boolean isOpen() {
081        try {
082            return super.isOpen();
083        } catch (TransactionRequiredException e) {
084            return false;
085        }
086    }
087
088    @Override
089    public void close() {
090        throw new IllegalStateException("Cannot close container managed entity manager");
091    }
092
093    @Override
094    public EntityTransaction getTransaction() {
095        throw new IllegalStateException("Not allowed to create transaction on shared EntityManager");
096    }
097
098    @Override
099    public void joinTransaction() {
100        throw new IllegalStateException("Not allowed to create transaction on shared EntityManager");
101    }
102
103    @Override
104    public EntityManagerFactory getEntityManagerFactory() {
105        return factory;
106    }
107
108    @Override
109    public Metamodel getMetamodel() {
110        return factory.getMetamodel();
111    }
112
113    @Override
114    public CriteriaBuilder getCriteriaBuilder() {
115        return factory.getCriteriaBuilder();
116    }
117
118    @Override
119    public String toString() {
120        return "Shared EntityManager proxy for target factory [" + factory + "]";
121    }
122}