Merge pull request #31 from scottescue/issue/21/support-serializing-l…

…azy-loaded-entities

Building the ability to serialize lazy loaded entity associations

src/main/java/com/scottescue/dropwizard/entitymanager/EntityManagerBundle.java

@@ -27,6 +27,8 @@ public abstract class EntityManagerBundle<T extends Configuration> implements Co
     private EntityManagerFactory entityManagerFactory;
     private EntityManagerContext entityManagerContext;
     private EntityManager sharedEntityManager;
+    private boolean serializeLazyLoadedEntitiesEnabled = true;
+    private boolean initialized = false;
 
     private final ImmutableList<Class<?>> entities;
     private final EntityManagerFactoryFactory entityManagerFactoryFactory;
@@ -64,20 +66,12 @@ public final void run(T configuration, Environment environment) throws Exception
                         dbConfig.getValidationQuery()));
     }
 
-    private UnitOfWorkApplicationListener registerUnitOfWorkListerIfAbsent(Environment environment) {
-        for (Object singleton : environment.jersey().getResourceConfig().getSingletons()) {
-            if (singleton instanceof UnitOfWorkApplicationListener) {
-                return (UnitOfWorkApplicationListener) singleton;
-            }
-        }
-        final UnitOfWorkApplicationListener listener = new UnitOfWorkApplicationListener();
-        environment.jersey().register(listener);
-        return listener;
-    }
-
     @Override
     public final void initialize(Bootstrap<?> bootstrap) {
-        bootstrap.getObjectMapper().registerModule(createHibernate5Module());
+        Hibernate5Module module = createHibernate5Module();
+        configure(module);
+        bootstrap.getObjectMapper().registerModule(module);
+        initialized = true;
     }
 
     /**
@@ -109,10 +103,33 @@ public EntityManager getSharedEntityManager() {
     }
 
     /**
-     * Override to configure the {@link Hibernate5Module}.
+     * Returns a boolean value indicating whether or not serializing lazy loaded entity associations is enabled.
+     * Serializing Lazy loaded entity associations is enabled by default.
+     *
+     * @return the value indicating whether serializing lazy loaded entity associations is enabled or not
+     */
+    public boolean isSerializeLazyLoadedEntitiesEnabled() {
+        return serializeLazyLoadedEntitiesEnabled;
+    }
+
+    /**
+     * Enables or disables serializing lazy loaded entity associations as determined by the given value.
+     *
+     * <br/><br/><i><strong>Note: </strong>This method should be called before the EntityManagerBundle is added
+     * to the application's {@link Bootstrap}, which initializes the bundle.  Once the bundle is initialized,
+     * any changes to the lazy loading property are ignored.</i>
+     *
+     * @param serializeLazyLoadedEntitiesEnabled the value indicating whether lazy loading is enabled or not
      */
-    protected Hibernate5Module createHibernate5Module() {
-        return new Hibernate5Module();
+    public void setSerializeLazyLoadedEntitiesEnabled(boolean serializeLazyLoadedEntitiesEnabled) {
+        // If the module is already initialized/bootstrapped there's no point in updating this property,
+        // an ObjectMapper has already been created and updating the property value could make the property
+        // out-of-sync with how the ObjectMapper is configured
+        if (initialized) {
+            return;
+        }
+        this.serializeLazyLoadedEntitiesEnabled = serializeLazyLoadedEntitiesEnabled;
+
     }
 
     /**
@@ -123,6 +140,14 @@ protected String name() {
         return DEFAULT_NAME;
     }
 
+    /**
+     * Override to configure Jackson's {@link Hibernate5Module}.
+     *
+     * @param module the Hibernate5Module object
+     */
+    protected void configure(Hibernate5Module module) {
+    }
+
     /**
      * Override to configure the JPA persistence unit.
      *
@@ -138,4 +163,23 @@ ImmutableList<Class<?>> getEntities() {
     EntityManagerContext getEntityManagerContext() {
         return this.entityManagerContext;
     }
+
+    private UnitOfWorkApplicationListener registerUnitOfWorkListerIfAbsent(Environment environment) {
+        for (Object singleton : environment.jersey().getResourceConfig().getSingletons()) {
+            if (singleton instanceof UnitOfWorkApplicationListener) {
+                return (UnitOfWorkApplicationListener) singleton;
+            }
+        }
+        final UnitOfWorkApplicationListener listener = new UnitOfWorkApplicationListener();
+        environment.jersey().register(listener);
+        return listener;
+    }
+
+    private Hibernate5Module createHibernate5Module() {
+        Hibernate5Module module = new Hibernate5Module();
+        if (serializeLazyLoadedEntitiesEnabled) {
+            module.enable(Hibernate5Module.Feature.FORCE_LAZY_LOADING);
+        }
+        return module;
+    }
 }

src/main/java/com/scottescue/dropwizard/entitymanager/EntityManagerContext.java

@@ -6,7 +6,9 @@
 import javax.persistence.EntityManagerFactory;
 import javax.persistence.PersistenceException;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.Map;
+import java.util.function.Consumer;
 
 /**
  * Represents the notion of contextual {@link EntityManager} instances managed by
@@ -80,6 +82,24 @@ static EntityManager unbind(EntityManagerFactory factory) {
         return existing;
     }
 
+    /**
+     * Unbinds all EntityManagers, regardless of EntityManagerFactory, currently associated with the context.
+     *
+     * @param function the function to apply to each EntityManager removed
+     */
+    static void unBindAll(Consumer<EntityManager> function) {
+        final Map<EntityManagerFactory,EntityManager> entityManagerMap = entityManagerMap(false);
+        if ( entityManagerMap != null ) {
+            Iterator<EntityManager> iterator = entityManagerMap.values().iterator();
+            while (iterator.hasNext()) {
+                EntityManager entityManager = iterator.next();
+                function.accept(entityManager);
+                iterator.remove();
+            }
+            doCleanup();
+        }
+    }
+
     @VisibleForTesting
     static synchronized Map<EntityManagerFactory,EntityManager> entityManagerMap(boolean createMap) {
         Map<EntityManagerFactory,EntityManager> entityManagerMap = CONTEXT_TL.get();

src/main/java/com/scottescue/dropwizard/entitymanager/UnitOfWorkApplicationListener.java

@@ -8,6 +8,7 @@
 import org.glassfish.jersey.server.monitoring.RequestEvent;
 import org.glassfish.jersey.server.monitoring.RequestEventListener;
 
+import javax.persistence.EntityManager;
 import javax.persistence.EntityManagerFactory;
 import javax.ws.rs.ext.Provider;
 import java.lang.reflect.Method;
@@ -70,18 +71,21 @@ public UnitOfWorkEventListener(Map<Method, UnitOfWork> methodMap,
 
         @Override
         public void onEvent(RequestEvent event) {
-            if (event.getType() == RequestEvent.Type.RESOURCE_METHOD_START) {
+            final RequestEvent.Type eventType = event.getType();
+            if (eventType == RequestEvent.Type.RESOURCE_METHOD_START) {
                 UnitOfWork unitOfWork = methodMap.get(event.getUriInfo()
                         .getMatchedResourceMethod().getInvocable().getDefinitionMethod());
                 unitOfWorkAspect.beforeStart(unitOfWork);
-            } else if (event.getType() == RequestEvent.Type.RESP_FILTERS_START) {
+            } else if (eventType == RequestEvent.Type.RESP_FILTERS_START) {
                 try {
                     unitOfWorkAspect.afterEnd();
                 } catch (Exception e) {
                     throw new MappableException(e);
                 }
-            } else if (event.getType() == RequestEvent.Type.ON_EXCEPTION) {
+            } else if (eventType == RequestEvent.Type.ON_EXCEPTION) {
                 unitOfWorkAspect.onError();
+            } else if (eventType == RequestEvent.Type.FINISHED) {
+                EntityManagerContext.unBindAll(EntityManager::close);
             }
         }
     }

src/main/java/com/scottescue/dropwizard/entitymanager/UnitOfWorkAspect.java

@@ -23,15 +23,14 @@ class UnitOfWorkAspect {
     // Context variables
     private UnitOfWork unitOfWork;
     private EntityManager entityManager;
-    private EntityManagerFactory entityManagerFactory;
 
     public void beforeStart(UnitOfWork unitOfWork) {
         if (unitOfWork == null) {
             return;
         }
         this.unitOfWork = unitOfWork;
 
-        entityManagerFactory = entityManagerFactories.get(unitOfWork.value());
+        EntityManagerFactory entityManagerFactory = entityManagerFactories.get(unitOfWork.value());
         if (entityManagerFactory == null) {
             // If the user didn't specify the name of a entityManager factory,
             // and we have only one registered, we can assume that it's the right one.
@@ -47,9 +46,7 @@ public void beforeStart(UnitOfWork unitOfWork) {
             EntityManagerContext.bind(entityManager);
             beginTransaction();
         } catch (Throwable th) {
-            entityManager.close();
             entityManager = null;
-            EntityManagerContext.unbind(entityManagerFactory);
             throw th;
         }
     }
@@ -64,12 +61,9 @@ public void afterEnd() {
         } catch (Exception e) {
             rollbackTransaction();
             throw e;
-        } finally {
-            entityManager.close();
-            entityManager = null;
-            EntityManagerContext.unbind(entityManagerFactory);
         }
-
+        // The entityManager should not be closed to let lazy loading work when serializing a response to the client.
+        // If the response is successfully serialized, then the entityManager will be closed by the `onFinish` method
     }
 
     public void onError() {
@@ -80,9 +74,7 @@ public void onError() {
         try {
             rollbackTransaction();
         } finally {
-            entityManager.close();
             entityManager = null;
-            EntityManagerContext.unbind(entityManagerFactory);
         }
     }
 

src/test/java/com/scottescue/dropwizard/entitymanager/AbstractIntegrationTest.java

@@ -0,0 +1,66 @@
+package com.scottescue.dropwizard.entitymanager;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+import io.dropwizard.Application;
+import io.dropwizard.Configuration;
+import io.dropwizard.db.DataSourceFactory;
+import io.dropwizard.jackson.Jackson;
+import io.dropwizard.jersey.jackson.JacksonMessageBodyProvider;
+import io.dropwizard.testing.ConfigOverride;
+import io.dropwizard.testing.DropwizardTestSupport;
+import io.dropwizard.testing.ResourceHelpers;
+import org.glassfish.jersey.client.JerseyClientBuilder;
+import org.junit.After;
+
+import javax.ws.rs.client.Client;
+
+public abstract class AbstractIntegrationTest {
+
+    public static class TestConfiguration extends Configuration {
+        private DataSourceFactory dataSource = new DataSourceFactory();
+
+        TestConfiguration(@JsonProperty("dataSource") DataSourceFactory dataSource) {
+            this.dataSource = dataSource;
+        }
+
+        public DataSourceFactory getDataSource() {
+            return dataSource;
+        }
+    }
+
+    final protected Client client = new JerseyClientBuilder()
+            .register(new JacksonMessageBodyProvider(Jackson.newObjectMapper()))
+            .build();
+
+    private DropwizardTestSupport dropwizardTestSupport;
+
+    @After
+    public void tearDown() {
+        dropwizardTestSupport.after();
+        client.close();
+    }
+
+    protected void setup(Class<? extends Application<TestConfiguration>> applicationClass) {
+        dropwizardTestSupport = new DropwizardTestSupport<>(applicationClass, ResourceHelpers.resourceFilePath("integration-test.yaml"),
+                ConfigOverride.config("dataSource.url", "jdbc:hsqldb:mem:DbTest" + System.nanoTime() + "?hsqldb.translate_dti_types=false"));
+        dropwizardTestSupport.before();
+    }
+
+    protected String getUrlPrefix() {
+        return "http://localhost:" + dropwizardTestSupport.getLocalPort();
+    }
+
+    protected String getUrl(String path) {
+        return getUrlPrefix() + path;
+    }
+
+    @SuppressWarnings({"ThrowableResultOfMethodCallIgnored", "unchecked"})
+    protected static <T extends Throwable> T unwrapThrowable(Class<T> type, Throwable throwable) {
+        Throwable cause = throwable.getCause();
+        if (cause == null) {
+            return null;
+        }
+        return cause.getClass().equals(type) ? (T) cause : unwrapThrowable(type, cause);
+    }
+
+}

src/test/java/com/scottescue/dropwizard/entitymanager/AbstractTestApplication.java

@@ -0,0 +1,65 @@
+package com.scottescue.dropwizard.entitymanager;
+
+import com.google.common.collect.ImmutableList;
+import io.dropwizard.db.PooledDataSourceFactory;
+import io.dropwizard.setup.Bootstrap;
+import io.dropwizard.setup.Environment;
+
+import javax.persistence.EntityManager;
+import javax.persistence.EntityManagerFactory;
+import javax.persistence.EntityTransaction;
+import javax.persistence.PersistenceException;
+
+public abstract class AbstractTestApplication extends io.dropwizard.Application<AbstractIntegrationTest.TestConfiguration> {
+    final EntityManagerBundle<AbstractIntegrationTest.TestConfiguration> entityManagerBundle = new EntityManagerBundle<AbstractIntegrationTest.TestConfiguration>(
+            supportedEntities(),
+            new EntityManagerFactoryFactory(),
+            new SharedEntityManagerFactory()) {
+        @Override
+        public PooledDataSourceFactory getDataSourceFactory(AbstractIntegrationTest.TestConfiguration configuration) {
+            return configuration.getDataSource();
+        }
+    };
+
+    @Override
+    public void initialize(Bootstrap<AbstractIntegrationTest.TestConfiguration> bootstrap) {
+        bootstrap.addBundle(entityManagerBundle);
+        onInitialize(bootstrap);
+    }
+
+    @Override
+    public void run(AbstractIntegrationTest.TestConfiguration configuration, Environment environment) throws Exception {
+        final EntityManagerFactory entityManagerFactory = entityManagerBundle.getEntityManagerFactory();
+        initDatabase(entityManagerFactory);
+
+        environment.jersey().register(new UnitOfWorkApplicationListener("hr-db", entityManagerFactory));
+        onRun(configuration, environment);
+    }
+
+    protected void initDatabase(EntityManagerFactory entityManagerFactory) {
+        final EntityManager entityManager = entityManagerFactory.createEntityManager();
+        final EntityTransaction transaction = entityManager.getTransaction();
+        try {
+            transaction.begin();
+
+            onInitDatabase(entityManager);
+
+            transaction.commit();
+        } catch (PersistenceException e) {
+            if (transaction.isActive()) {
+                transaction.rollback();
+            }
+        } finally {
+            entityManager.close();
+        }
+    }
+
+    protected abstract ImmutableList<Class<?>> supportedEntities();
+
+    protected void onInitialize(Bootstrap<AbstractIntegrationTest.TestConfiguration> bootstrap) {}
+
+    protected void onRun(AbstractIntegrationTest.TestConfiguration configuration, Environment environment) throws Exception {}
+
+    protected void onInitDatabase(EntityManager entityManager) {}
+
+}

src/test/java/com/scottescue/dropwizard/entitymanager/EntityManagerBundleTest.java

@@ -140,7 +140,8 @@ public void registersASessionFactoryHealthCheck() throws Exception {
     @Test
     @SuppressWarnings("unchecked")
     public void registersACustomNameOfHealthCheckAndDBPoolMetrics() throws Exception {
-        final EntityManagerBundle<Configuration> customBundle = new EntityManagerBundle<Configuration>(entities, factory, sharedEntityManagerFactory) {
+        final EntityManagerBundle<Configuration> customBundle = new EntityManagerBundle<Configuration>(entities,
+                factory, sharedEntityManagerFactory) {
             @Override
             public DataSourceFactory getDataSourceFactory(Configuration configuration) {
                 return dbConfig;
@@ -163,4 +164,25 @@ protected String name() {
                 ArgumentCaptor.forClass(EntityManagerFactoryHealthCheck.class);
         verify(healthChecks).register(eq("custom-hibernate"), captor.capture());
     }
+
+    @Test
+    public void serializingLazyLoadedEntitiesConfigChanges() {
+        bundle.setSerializeLazyLoadedEntitiesEnabled(false);
+
+        // Ensure the value IS changed since the bundle has not been initialized
+        assertThat(bundle.isSerializeLazyLoadedEntitiesEnabled()).isFalse();
+    }
+
+    @Test
+    public void ignoresSerializingLazyLoadedEntitiesConfigChangeAfterInit() {
+        ObjectMapper objectMapper = mock(ObjectMapper.class);
+        Bootstrap<?> bootstrap = mock(Bootstrap.class);
+        when(bootstrap.getObjectMapper()).thenReturn(objectMapper);
+
+        bundle.initialize(bootstrap);
+        bundle.setSerializeLazyLoadedEntitiesEnabled(false);
+
+        // Ensure the value IS NOT changed since the bundle was already initialized
+        assertThat(bundle.isSerializeLazyLoadedEntitiesEnabled()).isTrue();
+    }
 }