Add support for JTA transaction scoped entity managers

quarkusio · Sep 19, 2018 · 5676bb3 · 5676bb3
1 parent 2553d8e
commit 5676bb3
Show file tree

Hide file tree

Showing 12 changed files with 567 additions and 33 deletions.
diff --git a/examples/jpa-strict/src/main/resources/META-INF/persistence.xml b/examples/jpa-strict/src/main/resources/META-INF/persistence.xml
@@ -4,7 +4,7 @@
              http://xmlns.jcp.org/xml/ns/persistence/persistence_2_1.xsd"
              version="2.1">
 
-    <persistence-unit name="templatePU" transaction-type="RESOURCE_LOCAL">
+    <persistence-unit name="templatePU" transaction-type="JTA">
 
         <description>Hibernate test case template Persistence Unit</description>
 

diff --git a/examples/strict/src/main/java/org/jboss/shamrock/example/jpa/JPATestEMInjectionEndpoint.java b/examples/strict/src/main/java/org/jboss/shamrock/example/jpa/JPATestEMInjectionEndpoint.java
@@ -16,6 +16,7 @@
 import javax.servlet.http.HttpServlet;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
+import javax.transaction.UserTransaction;
 
 /**
  * Various tests for the JPA integration.
@@ -26,6 +27,9 @@ public class JPATestEMInjectionEndpoint extends HttpServlet {
     @Inject
     private EntityManager em;
 
+    @Inject
+    private UserTransaction transaction;
+
     @Override
     protected void doGet(HttpServletRequest req, HttpServletResponse resp) throws IOException {
         try {
@@ -43,16 +47,14 @@ public void testStoreLoadOnJPA() throws Exception {
 
     }
 
-    private void doStuffWithHibernate() {
-        EntityTransaction transaction = em.getTransaction();
+    private void doStuffWithHibernate() throws Exception {
         transaction.begin();
 
         persistNewPerson(em);
 
         listExistingPersons(em);
 
         transaction.commit();
-        em.close();
     }
 
     private static void listExistingPersons(EntityManager em) {

diff --git a/examples/strict/src/main/resources/META-INF/persistence.xml b/examples/strict/src/main/resources/META-INF/persistence.xml
@@ -4,7 +4,7 @@
              http://xmlns.jcp.org/xml/ns/persistence/persistence_2_1.xsd"
              version="2.1">
 
-    <persistence-unit name="templatePU" transaction-type="RESOURCE_LOCAL">
+    <persistence-unit name="templatePU" transaction-type="JTA">
 
         <description>Hibernate test case template Persistence Unit</description>
 

diff --git a/ext/gizmo/src/main/java/org/jboss/protean/gizmo/FieldCreator.java b/ext/gizmo/src/main/java/org/jboss/protean/gizmo/FieldCreator.java
@@ -1,7 +1,7 @@
 package org.jboss.protean.gizmo;
 
 
-public interface FieldCreator extends MemberCreator<FieldCreator> {
+public interface FieldCreator extends MemberCreator<FieldCreator>,AnnotatedElement {
 
     FieldDescriptor getFieldDescriptor();
 

diff --git a/ext/gizmo/src/main/java/org/jboss/protean/gizmo/FieldCreatorImpl.java b/ext/gizmo/src/main/java/org/jboss/protean/gizmo/FieldCreatorImpl.java
@@ -1,11 +1,18 @@
 package org.jboss.protean.gizmo;
 
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+import org.objectweb.asm.AnnotationVisitor;
 import org.objectweb.asm.ClassWriter;
+import org.objectweb.asm.FieldVisitor;
 import org.objectweb.asm.Opcodes;
 
 class FieldCreatorImpl implements FieldCreator {
 
     private final FieldDescriptor fieldDescriptor;
+    private final List<AnnotationCreatorImpl> annotations = new ArrayList<>();
 
     private int modifiers;
 
@@ -32,7 +39,21 @@ public FieldCreator setModifiers(int modifiers) {
 
     @Override
     public void write(ClassWriter file) {
-        file.visitField(modifiers, fieldDescriptor.getName(), fieldDescriptor.getType(), null, null);
+        FieldVisitor fieldVisitor = file.visitField(modifiers, fieldDescriptor.getName(), fieldDescriptor.getType(), null, null);
+        for(AnnotationCreatorImpl annotation : annotations) {
+            AnnotationVisitor av = fieldVisitor.visitAnnotation(DescriptorUtils.extToInt(annotation.getAnnotationType()), true);
+            for(Map.Entry<String, Object> e : annotation.getValues().entrySet()) {
+                av.visit(e.getKey(), e.getValue());
+            }
+            av.visitEnd();
+        }
+        fieldVisitor.visitEnd();
     }
 
+    @Override
+    public AnnotationCreator addAnnotation(String annotationType) {
+        AnnotationCreatorImpl ac = new AnnotationCreatorImpl(annotationType);
+        annotations.add(ac);
+        return ac;
+    }
 }
diff --git a/jpa/deployment/src/main/java/org/jboss/shamrock/jpa/HibernateReflectiveNeeds.java b/jpa/deployment/src/main/java/org/jboss/shamrock/jpa/HibernateReflectiveNeeds.java
@@ -35,6 +35,10 @@ private void registerAll() {
         simpleConstructor(org.hibernate.resource.transaction.backend.jdbc.internal.JdbcResourceLocalTransactionCoordinatorBuilderImpl.class);
         simpleConstructor(org.hibernate.id.enhanced.SequenceStyleGenerator.class);
         simpleConstructor(org.hibernate.boot.model.naming.ImplicitNamingStrategyJpaCompliantImpl.class);
+        simpleConstructor(org.hibernate.resource.transaction.backend.jta.internal.JtaTransactionCoordinatorBuilderImpl.class);
+        processorContext.addReflectiveClass(true, false, com.arjuna.ats.jta.UserTransaction.class.getName());
+        processorContext.addReflectiveClass(true, false, com.arjuna.ats.jta.TransactionManager.class.getName());
+
         //FIXME following is not Hibernate specific?
         simpleConstructor("com.sun.xml.internal.stream.events.XMLEventFactoryImpl");
         //ANTLR tokens:

diff --git a/jpa/deployment/src/main/java/org/jboss/shamrock/jpa/JpaJandexScavenger.java b/jpa/deployment/src/main/java/org/jboss/shamrock/jpa/JpaJandexScavenger.java
@@ -1,6 +1,7 @@
 package org.jboss.shamrock.jpa;
 
 import java.io.IOException;
+import java.io.Serializable;
 import java.util.Collection;
 import java.util.HashSet;
 import java.util.Set;
@@ -19,6 +20,7 @@
 import org.jboss.shamrock.deployment.ArchiveContext;
 import org.jboss.shamrock.deployment.ProcessorContext;
 import org.jboss.shamrock.deployment.codegen.BytecodeRecorder;
+import org.jboss.shamrock.jpa.runtime.JPADeploymentTemplate;
 
 /**
  * Scan the Jandex index to find JPA entities (and embeddables supporting entity models).
@@ -120,7 +122,7 @@ private static void addClassHierarchyToReflectiveList(DomainObjectSet collector,
         }
         ClassInfo classInfo = index.getClassByName(className);
         if (classInfo == null) {
-            if (className == ClassType.OBJECT_TYPE.name()) {
+            if (className == ClassType.OBJECT_TYPE.name() || className.toString().equals(Serializable.class.getName())) {
                 return;
             }
             else {

diff --git a/jpa/deployment/src/main/java/org/jboss/shamrock/jpa/cdi/HibernateCdiResourceProcessor.java b/jpa/deployment/src/main/java/org/jboss/shamrock/jpa/cdi/HibernateCdiResourceProcessor.java
@@ -15,12 +15,16 @@
 import javax.persistence.Persistence;
 import javax.persistence.PersistenceContext;
 import javax.persistence.PersistenceUnit;
+import javax.transaction.TransactionManager;
+import javax.transaction.TransactionSynchronizationRegistry;
 
 import org.jboss.jandex.AnnotationInstance;
 import org.jboss.jandex.AnnotationTarget;
+import org.jboss.jandex.AnnotationValue;
 import org.jboss.jandex.DotName;
 import org.jboss.protean.gizmo.ClassCreator;
 import org.jboss.protean.gizmo.ClassOutput;
+import org.jboss.protean.gizmo.FieldCreator;
 import org.jboss.protean.gizmo.FieldDescriptor;
 import org.jboss.protean.gizmo.MethodCreator;
 import org.jboss.protean.gizmo.MethodDescriptor;
@@ -30,6 +34,8 @@
 import org.jboss.shamrock.deployment.BeanDeployment;
 import org.jboss.shamrock.deployment.ProcessorContext;
 import org.jboss.shamrock.deployment.ResourceProcessor;
+import org.jboss.shamrock.jpa.runtime.cdi.SystemEntityManager;
+import org.jboss.shamrock.jpa.runtime.cdi.TransactionScopedEntityManager;
 
 public class HibernateCdiResourceProcessor implements ResourceProcessor {
 
@@ -45,6 +51,7 @@ public class HibernateCdiResourceProcessor implements ResourceProcessor {
 
     @Override
     public void process(ArchiveContext archiveContext, ProcessorContext processorContext) throws Exception {
+
         Set<String> knownUnitNames = new HashSet<>();
         Set<String> knownContextNames = new HashSet<>();
         scanForAnnotations(archiveContext, knownUnitNames, PERSISTENCE_UNIT);
@@ -57,7 +64,7 @@ public void process(ArchiveContext archiveContext, ProcessorContext processorCon
         Set<String> allKnownNames = new HashSet<>(knownUnitNames);
         allKnownNames.addAll(knownContextNames);
 
-        for (String name : knownContextNames) {
+        for (String name : allKnownNames) {
             String className = getClass().getName() + "$$EMFProducer-" + name;
             AtomicReference<byte[]> bytes = new AtomicReference<>();
             try (ClassCreator creator = new ClassCreator(new InMemoryClassOutput(bytes, processorContext), className, null, Object.class.getName())) {
@@ -80,27 +87,60 @@ public void process(ArchiveContext archiveContext, ProcessorContext processorCon
         }
 
 
-        for (String name : knownUnitNames) {
+        for (String name : knownContextNames) {
             String className = getClass().getName() + "$$EMProducer-" + name;
             AtomicReference<byte[]> bytes = new AtomicReference<>();
 
             //we need to know if transactions are present or not
-//            if (processorContext.isCapabilityPresent("transactions")) {
-//
-//            } else {
+            //TODO: this should be based on if a PU is JTA enabled or not
+            if (processorContext.isCapabilityPresent("transactions")) {
+                try (ClassCreator creator = new ClassCreator(new InMemoryClassOutput(bytes, processorContext), className, null, Object.class.getName())) {
+
+                    creator.addAnnotation(Dependent.class);
+
+                    FieldCreator emfField = creator.getFieldCreator("emf", EntityManagerFactory.class);
+                    emfField.addAnnotation(Inject.class);
+                    if (!knownUnitNames.contains(name)) {
+                        emfField.addAnnotation(SystemEntityManager.class);
+                    }
+                    FieldDescriptor emf = emfField.getFieldDescriptor();
+
+
+                    FieldCreator tsrField = creator.getFieldCreator("tsr", TransactionSynchronizationRegistry.class);
+                    tsrField.addAnnotation(Inject.class);
+                    FieldDescriptor tsr = tsrField.getFieldDescriptor();
+
+
+                    FieldCreator tmField = creator.getFieldCreator("tm", TransactionManager.class);
+                    tmField.addAnnotation(Inject.class);
+                    FieldDescriptor tm = tmField.getFieldDescriptor();
+
+                    MethodCreator producer = creator.getMethodCreator("producerMethod", EntityManager.class);
+                    producer.addAnnotation(Produces.class);
+                    producer.addAnnotation(ApplicationScoped.class);
+
+                    ResultHandle emfRh = producer.readInstanceField(emf, producer.getThis());
+                    ResultHandle tsrRh = producer.readInstanceField(tsr, producer.getThis());
+                    ResultHandle tmRh = producer.readInstanceField(tm, producer.getThis());
+
+                    producer.returnValue(producer.newInstance(MethodDescriptor.ofConstructor(TransactionScopedEntityManager.class, TransactionManager.class, TransactionSynchronizationRegistry.class, EntityManagerFactory.class), tmRh, tsrRh, emfRh));
+
+
+                }
+                beanDeployment.addGeneratedBean(className, bytes.get());
+            } else {
                 //if there is no TX support then we just use a super simple approach, and produce a normal EM
                 try (ClassCreator creator = new ClassCreator(new InMemoryClassOutput(bytes, processorContext), className, null, Object.class.getName())) {
 
                     creator.addAnnotation(Dependent.class);
 
-                    FieldDescriptor emf = creator.getFieldCreator("emf", EntityManagerFactory.class).getFieldDescriptor();
-                    MethodCreator setter = creator.getMethodCreator("setEmf", void.class, EntityManagerFactory.class);
-                    setter.writeInstanceField(emf, setter.getThis(), setter.getMethodParam(0));
-                    setter.addAnnotation(Inject.class);
+                    FieldCreator emfField = creator.getFieldCreator("emf", EntityManagerFactory.class);
+                    emfField.addAnnotation(Inject.class);
                     if (!knownUnitNames.contains(name)) {
-                        setter.addAnnotation(SystemEntityManager.class);
+                        emfField.addAnnotation(SystemEntityManager.class);
                     }
-                    setter.returnValue(null);
+                    FieldDescriptor emf = emfField.getFieldDescriptor();
+
 
                     MethodCreator producer = creator.getMethodCreator("producerMethod", EntityManager.class);
                     producer.addAnnotation(Produces.class);
@@ -117,28 +157,32 @@ public void process(ArchiveContext archiveContext, ProcessorContext processorCon
 
                 }
                 beanDeployment.addGeneratedBean(className, bytes.get());
-//            }
+            }
         }
 
     }
 
     private void scanForAnnotations(ArchiveContext archiveContext, Set<String> knownUnitNames, DotName nm) {
         for (AnnotationInstance anno : archiveContext.getCombinedIndex().getAnnotations(nm)) {
+            AnnotationValue unitName = anno.value("unitName");
+            if(unitName == null) {
+                continue;
+            }
             if (anno.target().kind() == AnnotationTarget.Kind.METHOD) {
                 if (anno.target().asMethod().hasAnnotation(PRODUCES)) {
-                    knownUnitNames.add(anno.value("unitName").asString());
+                    knownUnitNames.add(unitName.asString());
                 }
             } else if (anno.target().kind() == AnnotationTarget.Kind.FIELD) {
                 for (AnnotationInstance i : anno.target().asField().annotations()) {
                     if (i.name().equals(PRODUCES)) {
-                        knownUnitNames.add(anno.value("unitName").asString());
+                        knownUnitNames.add(unitName.asString());
                         break;
                     }
                 }
             } else if (anno.target().kind() == AnnotationTarget.Kind.CLASS) {
                 for (AnnotationInstance i : anno.target().asClass().classAnnotations()) {
                     if (i.name().equals(PRODUCES)) {
-                        knownUnitNames.add(anno.value("unitName").asString());
+                        knownUnitNames.add(unitName.asString());
                         break;
                     }
                 }

diff --git a/...s/shamrock/jpa/JPADeploymentTemplate.java → ...ck/jpa/runtime/JPADeploymentTemplate.java b/...s/shamrock/jpa/JPADeploymentTemplate.java → ...ck/jpa/runtime/JPADeploymentTemplate.java
@@ -1,4 +1,4 @@
-package org.jboss.shamrock.jpa;
+package org.jboss.shamrock.jpa.runtime;
 
 import org.hibernate.protean.Hibernate;
 

diff --git a/...shamrock/jpa/cdi/SystemEntityManager.java → .../jpa/runtime/cdi/SystemEntityManager.java b/...shamrock/jpa/cdi/SystemEntityManager.java → .../jpa/runtime/cdi/SystemEntityManager.java
@@ -1,4 +1,4 @@
-package org.jboss.shamrock.jpa.cdi;
+package org.jboss.shamrock.jpa.runtime.cdi;
 
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;