google#731 Add @RequiresUnitOfWork annotation

extensions/persist/src/com/google/inject/persist/PersistModule.java

@@ -40,14 +40,20 @@ protected final void configure() {
     /*if[AOP]*/
     // wrapping in an if[AOP] just to allow this to compile in NO_AOP -- it won't be used
 
-    // class-level @Transacational
+    // class-level @Transactional and @RequiresUnitOfWork
     bindInterceptor(annotatedWith(Transactional.class), any(), getTransactionInterceptor());
-    // method-level @Transacational
+    bindInterceptor(annotatedWith(RequiresUnitOfWork.class), any(),
+          getRequiresUnitOfWorkInterceptor());
+    // method-level @Transactional and @RequiresUnitOfWork
     bindInterceptor(any(), annotatedWith(Transactional.class), getTransactionInterceptor());
+    bindInterceptor(any(), annotatedWith(RequiresUnitOfWork.class),
+          getRequiresUnitOfWorkInterceptor());
     /*end[AOP]*/
   }
 
   protected abstract void configurePersistence();
 
   protected abstract MethodInterceptor getTransactionInterceptor();
+
+  protected abstract MethodInterceptor getRequiresUnitOfWorkInterceptor();
 }

extensions/persist/src/com/google/inject/persist/RequiresUnitOfWork.java

@@ -0,0 +1,19 @@
+package com.google.inject.persist;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Inherited;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * <p> Any method or class marked with this annotation will require the existence
+ * of a unit of work.
+ * <p>Marking a method {@code @RequiresUnitOfWork} will start a unit of work if none
+ * exists before the method executes and end it if it was started after the method returns.
+ */
+@Target({ ElementType.METHOD, ElementType.TYPE })
+@Retention(RetentionPolicy.RUNTIME)
+@Inherited
+public @interface RequiresUnitOfWork {
+}

extensions/persist/src/com/google/inject/persist/jpa/JpaLocalTxnInterceptor.java

@@ -18,7 +18,6 @@
 
 import com.google.inject.Inject;
 import com.google.inject.persist.Transactional;
-import com.google.inject.persist.UnitOfWork;
 
 import org.aopalliance.intercept.MethodInterceptor;
 import org.aopalliance.intercept.MethodInvocation;
@@ -33,29 +32,19 @@
  */
 class JpaLocalTxnInterceptor implements MethodInterceptor {
 
-  // TODO(gak): Move these args to the cxtor & make these final.
-  @Inject
-  private JpaPersistService emProvider = null;
-
-  @Inject
-  private UnitOfWork unitOfWork = null;
-
   @Transactional
   private static class Internal {}
-
-  // Tracks if the unit of work was begun implicitly by this transaction.
-  private final ThreadLocal<Boolean> didWeStartWork = new ThreadLocal<Boolean>();
-
+
+  // TODO(gak): Move this arg to the cxtor & make this final.
+  @Inject
+  private UnitOfWorkHandler unitOfWorkHandler;
+
   public Object invoke(MethodInvocation methodInvocation) throws Throwable {
 
-    // Should we start a unit of work?
-    if (!emProvider.isWorking()) {
-      emProvider.begin();
-      didWeStartWork.set(true);
-    }
+    unitOfWorkHandler.requireUnitOfWork();
 
     Transactional transactional = readTransactionMetadata(methodInvocation);
-    EntityManager em = this.emProvider.get();
+    EntityManager em = unitOfWorkHandler.getEntityManager();
 
     // Allow 'joining' of transactions if there is an enclosing @Transactional method.
     if (em.getTransaction().isActive()) {
@@ -78,10 +67,8 @@ public Object invoke(MethodInvocation methodInvocation) throws Throwable {
       //propagate whatever exception is thrown anyway
       throw e;
     } finally {
-      // Close the em if necessary (guarded so this code doesn't run unless catch fired).
-      if (null != didWeStartWork.get() && !txn.isActive()) {
-        didWeStartWork.remove();
-        unitOfWork.end();
+      if (!txn.isActive()) {
+        unitOfWorkHandler.endRequireUnitOfWork();
       }
     }
 
@@ -90,11 +77,7 @@ public Object invoke(MethodInvocation methodInvocation) throws Throwable {
     try {
       txn.commit();
     } finally {
-      //close the em if necessary
-      if (null != didWeStartWork.get() ) {
-        didWeStartWork.remove();
-        unitOfWork.end();
-      }
+      unitOfWorkHandler.endRequireUnitOfWork();
     }
 
     //or return result

extensions/persist/src/com/google/inject/persist/jpa/JpaPersistModule.java

@@ -21,13 +21,11 @@
 import com.google.inject.Inject;
 import com.google.inject.Provides;
 import com.google.inject.Singleton;
-import com.google.inject.TypeLiteral;
 import com.google.inject.persist.PersistModule;
 import com.google.inject.persist.PersistService;
 import com.google.inject.persist.UnitOfWork;
 import com.google.inject.persist.finder.DynamicFinder;
 import com.google.inject.persist.finder.Finder;
-import com.google.inject.util.Providers;
 
 import org.aopalliance.intercept.MethodInterceptor;
 import org.aopalliance.intercept.MethodInvocation;
@@ -58,6 +56,7 @@ public JpaPersistModule(String jpaUnit) {
 
   private Map<?,?> properties;
   private MethodInterceptor transactionInterceptor;
+  private MethodInterceptor requiresUnitOfWorkInterceptor;
 
   @Override protected void configurePersistence() {
     bindConstant().annotatedWith(Jpa.class).to(jpaUnit);
@@ -69,9 +68,13 @@ public JpaPersistModule(String jpaUnit) {
     bind(EntityManager.class).toProvider(JpaPersistService.class);
     bind(EntityManagerFactory.class)
         .toProvider(JpaPersistService.EntityManagerFactoryProvider.class);
+    bind(UnitOfWorkHandler.class);
 
     transactionInterceptor = new JpaLocalTxnInterceptor();
     requestInjection(transactionInterceptor);
+
+    requiresUnitOfWorkInterceptor = new RequiresUnitOfWorkInterceptor();
+    requestInjection(requiresUnitOfWorkInterceptor);
 
     // Bind dynamic finders.
     for (Class<?> finder : dynamicFinders) {
@@ -83,6 +86,10 @@ public JpaPersistModule(String jpaUnit) {
     return transactionInterceptor;
   }
 
+  @Override protected MethodInterceptor getRequiresUnitOfWorkInterceptor() {
+    return requiresUnitOfWorkInterceptor;
+  }
+
   @Provides @Jpa Map<?, ?> provideProperties() {
     return properties;
   }

extensions/persist/src/com/google/inject/persist/jpa/RequiresUnitOfWorkInterceptor.java

@@ -0,0 +1,43 @@
+package com.google.inject.persist.jpa;
+
+import java.lang.reflect.Method;
+
+import org.aopalliance.intercept.MethodInterceptor;
+import org.aopalliance.intercept.MethodInvocation;
+
+import com.google.inject.Inject;
+import com.google.inject.persist.RequiresUnitOfWork;
+
+public class RequiresUnitOfWorkInterceptor implements MethodInterceptor {
+  @Inject
+  private UnitOfWorkHandler unitOfWorkHandler;
+
+  public Object invoke(MethodInvocation methodInvocation) throws Throwable {
+    RequiresUnitOfWork annotation = readAnnotationMetadata(methodInvocation);
+    if (annotation == null) {
+      // Avoid creating the unit of work in Object class methods
+      return methodInvocation.proceed();
+
+    } else {
+      unitOfWorkHandler.requireUnitOfWork();
+      try {
+        return methodInvocation.proceed();
+      } finally {
+        unitOfWorkHandler.endRequireUnitOfWork();
+      }
+    }
+  }
+
+  private RequiresUnitOfWork readAnnotationMetadata(MethodInvocation methodInvocation) {
+    RequiresUnitOfWork annotation;
+    Method method = methodInvocation.getMethod();
+
+    // Annotation in method or class
+    annotation = method.getAnnotation(RequiresUnitOfWork.class);
+    if (annotation == null) {
+      annotation = method.getClass().getAnnotation(RequiresUnitOfWork.class);
+    }
+
+    return annotation;
+  }
+}

extensions/persist/src/com/google/inject/persist/jpa/UnitOfWorkHandler.java

@@ -0,0 +1,44 @@
+package com.google.inject.persist.jpa;
+
+import javax.persistence.EntityManager;
+
+import com.google.inject.Inject;
+import com.google.inject.persist.UnitOfWork;
+
+/**
+ * Utility class able to handle easily the start and end of units of work
+ * in a nested context. It starts the unit of work if there isn't any in
+ * the current thread, and ends it if the unit of work was started by this
+ * instance.
+ */
+public class UnitOfWorkHandler {
+  private JpaPersistService emProvider = null;
+  private UnitOfWork unitOfWork = null;
+
+  /** Tracks if the unit of work was begun implicitly by this handler. */
+  private final ThreadLocal<Boolean> didWeStartWork = new ThreadLocal<Boolean>();
+
+  @Inject
+  public UnitOfWorkHandler(JpaPersistService emProvider, UnitOfWork unitOfWork) {
+    this.emProvider = emProvider;
+    this.unitOfWork = unitOfWork;
+  }
+
+  public void requireUnitOfWork() {
+    if (!emProvider.isWorking()) {
+      unitOfWork.begin();
+      didWeStartWork.set(true);
+    }
+  }
+
+  public void endRequireUnitOfWork() {
+    if (didWeStartWork.get() != null) {
+      didWeStartWork.remove();
+      unitOfWork.end();
+    }
+  }
+
+  public EntityManager getEntityManager() {
+    return emProvider.get();
+  }
+}

extensions/persist/test/com/google/inject/persist/jpa/RequiresUnitOfWorkTest.java

@@ -0,0 +1,88 @@
+package com.google.inject.persist.jpa;
+
+import javax.persistence.EntityManager;
+import javax.persistence.EntityManagerFactory;
+
+import junit.framework.TestCase;
+
+import com.google.inject.Guice;
+import com.google.inject.Inject;
+import com.google.inject.Injector;
+import com.google.inject.Provider;
+import com.google.inject.persist.RequiresUnitOfWork;
+
+public class RequiresUnitOfWorkTest extends TestCase {
+  private Injector injector;
+
+  public void setUp() {
+    injector = Guice.createInjector(new JpaPersistModule("testUnit"));
+
+    //startup persistence and end the unit of work, so that no one is running
+    JpaPersistService persistService = injector.getInstance(JpaPersistService.class);
+    persistService.start();
+    persistService.end();
+  }
+
+  public void tearDown() {
+    injector.getInstance(EntityManagerFactory.class).close();
+  }
+
+  public void testRequiresWithNoUnitOfWork() {
+    JpaPersistService persistService = injector.getInstance(JpaPersistService.class);
+    DataAccessObject dataObject = injector.getInstance(DataAccessObject.class);
+
+    // Run operation with no exception
+    dataObject.runOperationInUnitOfWork();
+
+    // Unit of work has been closed
+    assertFalse(persistService.isWorking());
+  }
+
+  public void testRequiresWithUnitOfWork() {
+    JpaPersistService persistService = injector.getInstance(JpaPersistService.class);
+    DataAccessObject dataObject = injector.getInstance(DataAccessObject.class);
+
+    // Run operation with no exception
+    persistService.begin();
+    dataObject.runOperationInUnitOfWork();
+
+    // Unit of work has not been closed
+    assertTrue(persistService.isWorking());
+    persistService.end();
+  }
+
+  public void testRequiresWithNestedUnitsOfWork() {
+    JpaPersistService persistService = injector.getInstance(JpaPersistService.class);
+    NestDataObject dataObject = injector.getInstance(NestDataObject.class);
+
+    // Run operation with no exception
+    dataObject.runOperationInUnitOfWork();
+
+    // Unit of work has not been closed
+    assertFalse(persistService.isWorking());
+  }
+
+
+  public static class DataAccessObject {
+    @Inject Provider<EntityManager> emProvider;
+
+    @RequiresUnitOfWork
+    public void runOperationInUnitOfWork() {
+      emProvider.get()
+          .createQuery("from JpaTestEntity", JpaTestEntity.class)
+          .setMaxResults(1).getResultList();
+    }
+  }
+
+  public static class NestDataObject {
+    @Inject Provider<EntityManager> emProvider;
+    @Inject DataAccessObject dataObject;
+
+    @RequiresUnitOfWork
+    public void runOperationInUnitOfWork() {
+      emProvider.get()
+          .createQuery("from JpaTestEntity", JpaTestEntity.class)
+          .setMaxResults(1).getResultList();
+    }
+  }
+}