JPA entity enhancement without the Maven plugin

Author: Sanne

jpa/deployment/pom.xml

@@ -24,6 +24,17 @@
             <groupId>org.hibernate</groupId>
             <artifactId>hibernate-core</artifactId>
         </dependency>
+        <dependency>
+            <groupId>junit</groupId>
+            <artifactId>junit</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.jboss.protean.gizmo</groupId>
+            <artifactId>gizmo</artifactId>
+            <type>test-jar</type>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
 

jpa/deployment/src/main/java/org/jboss/shamrock/jpa/HibernateEntityEnhancer.java

@@ -0,0 +1,94 @@
+package org.jboss.shamrock.jpa;
+
+import java.util.Set;
+import java.util.function.Function;
+
+import org.hibernate.bytecode.enhance.spi.DefaultEnhancementContext;
+import org.hibernate.bytecode.enhance.spi.Enhancer;
+import org.hibernate.bytecode.spi.BytecodeProvider;
+
+import org.objectweb.asm.ClassReader;
+import org.objectweb.asm.ClassVisitor;
+import org.objectweb.asm.ClassWriter;
+import org.objectweb.asm.Opcodes;
+
+/**
+ * Used to transform bytecode by registering to org.jboss.shamrock.deployment.ProcessorContext#addByteCodeTransformer(java.util.function.Function)
+ * this function adapts the Shamrock bytecode transformer API - which uses ASM - to use the Entity Enhancement API of
+ * Hibernate ORM, which exposes a simple byte array.
+ *
+ * @author Sanne Grinovero  <sanne@hibernate.org>
+ */
+public final class HibernateEntityEnhancer implements Function<String, Function<ClassVisitor, ClassVisitor>> {
+
+    private final Enhancer enhancer;
+    private final Set<String> classnameWhitelist;
+
+    @Deprecated //shouldn't use this version, but will do for a while
+    public HibernateEntityEnhancer() {
+        this(null);
+    }
+
+    public HibernateEntityEnhancer(Set<String> classnameWhitelist) {
+        this.classnameWhitelist = classnameWhitelist;
+        BytecodeProvider provider = new org.hibernate.bytecode.internal.bytebuddy.BytecodeProviderImpl();
+        this.enhancer = provider.getEnhancer(new DefaultEnhancementContext());
+    }
+
+    @Override
+    public Function<ClassVisitor, ClassVisitor> apply(String classname) {
+        if (classnameWhitelist == null || classnameWhitelist.contains(classname))
+            return new HibernateTransformingVisitorFunction(classname);
+        else
+            return null;
+    }
+
+    /**
+     * Having to convert a ClassVisitor into another, this allows visitor chaining: the returned ClassVisitor needs to
+     * refer to the previous ClassVisitor in the chain to forward input events (optionally transformed).
+     */
+    private class HibernateTransformingVisitorFunction implements Function<ClassVisitor, ClassVisitor> {
+
+        private final String className;
+
+        public HibernateTransformingVisitorFunction(String className) {
+            this.className = className;
+        }
+
+        @Override
+        public ClassVisitor apply(ClassVisitor outputClassVisitor) {
+            return new HibernateEnhancingClassVisitor(className, outputClassVisitor);
+        }
+    }
+
+    private class HibernateEnhancingClassVisitor extends ClassVisitor {
+
+        private final String className;
+        private final ClassVisitor outputClassVisitor;
+
+        public HibernateEnhancingClassVisitor(String className, ClassVisitor outputClassVisitor) {
+            super(Opcodes.ASM6, new ClassWriter(0));
+            this.className = className;
+            this.outputClassVisitor = outputClassVisitor;
+        }
+
+        public void visitEnd() {
+            super.visitEnd();
+            final ClassWriter writer = (ClassWriter) this.cv; //safe cast: cv is the the ClassWriter instance we passed to the super constructor
+            //We need to convert the nice Visitor chain into a plain byte array to adapt to the Hibernate ORM
+            //enhancement API:
+            final byte[] inputBytes = writer.toByteArray();
+            final byte[] transformedBytes = hibernateEnhancement(className, inputBytes);
+            //Then re-convert the transformed bytecode to not interrupt the visitor chain:
+            ClassReader cr = new ClassReader(transformedBytes);
+            cr.accept(outputClassVisitor, 0);
+        }
+
+    }
+
+    private byte[] hibernateEnhancement(final String className, final byte[] originalBytes) {
+        final byte[] enhanced = enhancer.enhance(className, originalBytes);
+        return enhanced == null ? originalBytes : enhanced;
+    }
+
+}

jpa/deployment/src/main/java/org/jboss/shamrock/jpa/JPAAnnotationProcessor.java

@@ -27,6 +27,9 @@ public class JPAAnnotationProcessor implements ResourceProcessor {
 
     @Override
     public void process(ArchiveContext archiveContext, ProcessorContext processorContext) throws Exception {
+
+        enhanceEntities(archiveContext, processorContext);
+
         // Hibernate specific reflective classes
         processorContext.addReflectiveClass(false, false, org.hibernate.jpa.HibernatePersistenceProvider.class.getName());
         processorContext.addReflectiveClass(false, false, org.hibernate.persister.entity.SingleTableEntityPersister.class.getName());
@@ -47,6 +50,11 @@ public void process(ArchiveContext archiveContext, ProcessorContext processorCon
         }
     }
 
+    private void enhanceEntities(ArchiveContext archiveContext, ProcessorContext processorContext) {
+        //TODO make sure we pass the list of known entities as a white-list filter so to avoid processing all classes
+        processorContext.addByteCodeTransformer(new HibernateEntityEnhancer());
+    }
+
     private void enlistReturnType(ProcessorContext processorContext, IndexView index) {
         Collection<AnnotationInstance> annotations = index.getAnnotations(EMBEDDED);
         if (annotations != null && annotations.size() > 0) {

jpa/deployment/src/test/java/org/jboss/shamrock/jpa/enhancer/Address.java

@@ -0,0 +1,41 @@
+package org.jboss.shamrock.jpa.enhancer;
+
+import javax.persistence.Entity;
+import javax.persistence.GeneratedValue;
+import javax.persistence.GenerationType;
+import javax.persistence.Id;
+
+@Entity
+public class Address {
+
+    private long id;
+    private String street;
+
+    public Address() {
+    }
+
+    public Address(String street) {
+        this.street = street;
+    }
+
+    @Id @GeneratedValue(strategy = GenerationType.SEQUENCE, generator="addressSeq")
+    public long getId() {
+        return id;
+    }
+
+    public void setId(long id) {
+        this.id = id;
+    }
+
+    public String getStreet() {
+        return street;
+    }
+
+    public void setStreet(String name) {
+        this.street = name;
+    }
+
+    public void describeFully(StringBuilder sb) {
+        sb.append( "Address with id=" ).append( id ).append( ", street='" ).append( street ).append( "'" );
+    }
+}

jpa/deployment/src/test/java/org/jboss/shamrock/jpa/enhancer/HibernateEntityEnhancerTest.java

@@ -0,0 +1,57 @@
+package org.jboss.shamrock.jpa.enhancer;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.Set;
+
+import org.hibernate.engine.spi.ManagedEntity;
+import org.hibernate.engine.spi.PersistentAttributeInterceptable;
+import org.hibernate.engine.spi.SelfDirtinessTracker;
+import org.jboss.shamrock.jpa.HibernateEntityEnhancer;
+
+import org.junit.Assert;
+import org.junit.Test;
+import org.objectweb.asm.ClassReader;
+import org.objectweb.asm.ClassVisitor;
+import org.objectweb.asm.ClassWriter;
+
+import org.jboss.protean.gizmo.TestClassLoader;
+
+/**
+ * Verifies the HibernateEntityEnhancer actually does enhance the entity class
+ *
+ * @author Sanne Grinovero  <sanne@hibernate.org>
+ */
+public class HibernateEntityEnhancerTest {
+
+    @Test
+    public void testBytecodeEnhancement() throws IOException, ClassNotFoundException {
+
+        Assert.assertFalse(isEnhanced(Address.class));
+
+        final String className = Address.class.getName();
+
+        ClassReader classReader = new ClassReader(className);
+        ClassWriter writer = new ClassWriter(classReader, ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS);
+        ClassVisitor visitor = writer;
+        HibernateEntityEnhancer hibernateEntityEnhancer = new HibernateEntityEnhancer();
+        visitor = hibernateEntityEnhancer.apply(className).apply(visitor);
+        classReader.accept(visitor, 0);
+        final byte[] modifiedBytecode = writer.toByteArray();
+
+        TestClassLoader cl = new TestClassLoader(getClass().getClassLoader());
+        cl.write(className, modifiedBytecode);
+        final Class<?> modifiedClass = cl.loadClass(className);
+        Assert.assertTrue(isEnhanced(modifiedClass));
+    }
+
+    private boolean isEnhanced(final Class<?> modifiedClass) {
+        Set<Class> interfaces = new HashSet<Class>(Arrays.asList(modifiedClass.getInterfaces()));
+        //Assert it now implements these three interfaces:
+        return interfaces.contains(ManagedEntity.class) &&
+              interfaces.contains(PersistentAttributeInterceptable.class) &&
+              interfaces.contains(SelfDirtinessTracker.class);
+    }
+
+}