Dropwizard unit of work

pom.xml

@@ -78,6 +78,7 @@
         <sonar.host.url>https://sonarcloud.io</sonar.host.url>
 
         <jdbi.version>2.78</jdbi.version>
+        <org.reflections.version>0.10.2</org.reflections.version>
     </properties>
 
     <dependencies>
@@ -146,6 +147,21 @@
             <groupId>org.eclipse.jetty</groupId>
             <artifactId>jetty-util</artifactId>
         </dependency>
+        <dependency>
+            <groupId>org.reflections</groupId>
+            <artifactId>reflections</artifactId>
+            <version>${org.reflections.version}</version>
+            <exclusions>
+                <exclusion>
+                    <groupId>com.google.guava</groupId>
+                    <artifactId>guava</artifactId>
+                </exclusion>
+            </exclusions>
+        </dependency>
+        <dependency>
+            <groupId>org.glassfish.jersey.core</groupId>
+            <artifactId>jersey-server</artifactId>
+        </dependency>
 
         <dependency>
             <groupId>org.junit.jupiter</groupId>
@@ -202,7 +218,7 @@
                     <execution>
                         <id>prepare-agent</id>
                         <goals>
-                              <goal>prepare-agent</goal>
+                            <goal>prepare-agent</goal>
                         </goals>
                     </execution>
                     <execution>

src/main/java/io/dropwizard/jdbi/unitofwork/JdbiUnitOfWork.java

@@ -0,0 +1,22 @@
+package io.dropwizard.jdbi.unitofwork;
+
+import java.lang.annotation.Documented;
+import java.lang.annotation.Retention;
+import java.lang.annotation.Target;
+
+import static java.lang.annotation.ElementType.METHOD;
+import static java.lang.annotation.RetentionPolicy.RUNTIME;
+
+/**
+ * When annotating a Jersey resource method, wraps the method in a Jdbi transaction context
+ * associated with a valid handle.
+ * <br><br>
+ * A transaction will automatically {@code begin} before the resource method is invoked,
+ * {@code commit} if the method returned without throwing any exception and {@code rollback}
+ * if an exception was thrown.
+ */
+@Target(METHOD)
+@Retention(RUNTIME)
+@Documented
+public @interface JdbiUnitOfWork {
+}

src/main/java/io/dropwizard/jdbi/unitofwork/README.md

@@ -0,0 +1,85 @@
+## @JdbiUnitOfWork - Unit of Work Support
+
+Provides a `Unit of Work` annotation for a Jdbi backed Dropwizard service for wrapping resource methods in a transaction
+context
+
+- [`Dropwizard`](https://github.com/dropwizard/dropwizard) provides a very
+  slick [`@UnitOfWork`](https://www.dropwizard.io/en/latest/manual/hibernate.html) annotation that wraps a transaction
+  context around resource methods annotated with this annotation. This is very useful for wrapping multiple calls in a
+  single database transaction all of which will succeed or roll back atomically.
+
+
+- However this support is only available for `Hibernate`. This module provides support for a `Jdbi`backend
+
+## Features
+
+- `transactionality` across multiple datasources when called from a request thread
+- `transactionality` across multiple datasources across `multiple threads`
+- `excluding` selectively, certain set of URI's from transaction contexts, such as `ELB`, `Health Checks` etc
+- `Http GET` methods are excluded from transaction by default.
+- `Http POST` methods are wrapped around in a transaction only when annotated with `@JdbiUnitOfWork`
+
+## Usage
+
+- Add the dependency to your `pom.xml`
+
+- Construct a `JdbiUnitOfWorkProvider` from the DBI instance.
+
+  ```java
+  JdbiUnitOfWorkProvider provider = JdbiUnitOfWorkProvider.withDefault(dbi); // most common
+               or
+  JdbiUnitOfWorkProvider provider = JdbiUnitOfWorkProvider.withLinked(dbi); // most common
+  ```
+
+  If you are using Guice, you can bind the instance
+  ```
+  bind(JdbiUnitOfWorkProvider.class).toInstance(provider);
+  ```
+
+<br>
+
+- Provide the list of package where the SQL Objects / DAO (to be attached) are located. Classes with Jdbi
+  annotations `@SqlQuery` or `@SqlUpdate` or `@SqlBatch` or `@SqlCall` will be picked automatically.
+
+  <br>
+
+  Use `JdbiUnitOfWorkProvider` to generate the proxies. You can also register the classes one by one.
+
+  ```java
+
+  // class level
+  SampleDao dao = (SampleDao) provider.getWrappedInstanceForDaoClass(SampleDao.class);
+  // use the proxies and pass it as they were normal instances
+  resource = new SampleResource(dao);
+
+  // package level
+  List<String> daoPackages = Lists.newArrayList("<fq-package-name>", "fq-package-name-2", ...);
+  Map<? extends Class, Object> proxies = unitOfWorkProvider.getWrappedInstanceForDaoPackage(daoPackages);
+  // use the proxies and pass it as they were normal instances
+  resource = ...new SampleResource((SampleDao)proxies.get(SampleDao.class))
+  ```
+
+<br>
+
+- Finally, we need to register the event listener with the Jersey Environment using the constructed provider
+  ```
+  environment.jersey().register(new JdbiUnitOfWorkApplicationEventListener(provider, new HashSet<>()));;
+  ```
+  In case you'd like to exclude certain URI paths from being monitored, you can pass them into exclude paths;
+  ```
+  Set<String> excludePaths = new HashSet<>();
+  environment.jersey().register(new JdbiUnitOfWorkApplicationEventListener(handleManager, excludePaths));
+  ```
+
+<br>
+
+- Start annotating resource methods with `@JdbiUnitOfWork` and you're good to go.
+    ```java
+    @POST
+    @Path("/")
+    @JdbiUnitOfWork
+    public RequestResponse createRequest() {
+          ..do stateful work (across multiple Dao's)
+          return response
+    }
+    ```

src/main/java/io/dropwizard/jdbi/unitofwork/core/DefaultJdbiHandleManager.java

@@ -0,0 +1,39 @@
+package io.dropwizard.jdbi.unitofwork.core;
+
+import org.skife.jdbi.v2.DBI;
+import org.skife.jdbi.v2.Handle;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * This implementation gets a new handle each time it is invoked. It simulates the default
+ * behaviour of creating new handles each time the dao method is invoked.
+ * <br><br>
+ * It can be used to service requests which interact with only a single method in a single handle.
+ * This is a lightweight implementation suitable for testing, such as with embedded databases.
+ * Any serious application should not be using this as it may quickly leak / run out of handles
+ *
+ * @apiNote Not suitable for requests spanning multiple Dbi as the handle returned is different
+ * This implementation, therefore, does not support thread factory creation.
+ */
+public class DefaultJdbiHandleManager implements JdbiHandleManager {
+
+    private final Logger log = LoggerFactory.getLogger(DefaultJdbiHandleManager.class);
+    private final DBI dbi;
+
+    public DefaultJdbiHandleManager(DBI dbi) {
+        this.dbi = dbi;
+    }
+
+    @Override
+    public Handle get() {
+        Handle handle = dbi.open();
+        log.debug("handle [{}] : Thread Id [{}]", handle.hashCode(), Thread.currentThread().getId());
+        return handle;
+    }
+
+    @Override
+    public void clear() {
+        log.debug("No Op");
+    }
+}

src/main/java/io/dropwizard/jdbi/unitofwork/core/JdbiHandleManager.java

@@ -0,0 +1,55 @@
+package io.dropwizard.jdbi.unitofwork.core;
+
+import org.skife.jdbi.v2.Handle;
+
+import java.util.concurrent.ThreadFactory;
+
+/**
+ * A {@link JdbiHandleManager} is used to provide the lifecycle of a {@link Handle} with respect
+ * to a given scope. A scope may be session based, request based or may be invoked on every run.
+ */
+public interface JdbiHandleManager {
+
+    /**
+     * Provide a way to get a Jdbi handle, a wrapped connection to the underlying database
+     *
+     * @return a valid handle tied with a specific scope
+     */
+    Handle get();
+
+    /**
+     * Provide a way to clear the handle rendering it useless for the other methods
+     */
+    void clear();
+
+    /**
+     * Provide a thread factory for the caller with some identity represented by the
+     * {@link #getConversationId()}. This can be used by the caller to create multiple threads,
+     * say, using {@link java.util.concurrent.ExecutorService}. The {@link JdbiHandleManager} can
+     * then use the thread factory to identify and manage handle use across multiple threads.
+     *
+     * @return a thread factory used to safely create multiple threads
+     * @throws UnsupportedOperationException by default. Implementations overriding this method
+     *                                       must ensure that the conversation id is unique
+     */
+    default ThreadFactory createThreadFactory() {
+        throw new UnsupportedOperationException("Thread factory creation is not supported");
+    }
+
+    /**
+     * Provide a unique identifier for the conversation with a handle. No two identifiers
+     * should co exist at once during the application lifecycle or else handle corruption
+     * or misuse might occur.
+     * <br><br>
+     * This can be relied upon by the {@link #createThreadFactory()} to reuse handles across
+     * multiple threads spawned off a request thread.
+     *
+     * @return a unique identifier applicable to a scope
+     * @implNote hashcode can not be relied upon for providing a unique identifier due to the
+     * possibility of collision. Instead opt for a monotonically increasing counter, such as
+     * the thread id.
+     */
+    default String getConversationId() {
+        return String.valueOf(Thread.currentThread().getId());
+    }
+}

src/main/java/io/dropwizard/jdbi/unitofwork/core/JdbiUnitOfWorkProvider.java

@@ -0,0 +1,129 @@
+package io.dropwizard.jdbi.unitofwork.core;
+
+import com.google.common.collect.Sets;
+import com.google.common.reflect.Reflection;
+import org.reflections.Reflections;
+import org.reflections.scanners.Scanners;
+import org.skife.jdbi.v2.DBI;
+import org.skife.jdbi.v2.sqlobject.SqlBatch;
+import org.skife.jdbi.v2.sqlobject.SqlCall;
+import org.skife.jdbi.v2.sqlobject.SqlQuery;
+import org.skife.jdbi.v2.sqlobject.SqlUpdate;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.lang.reflect.Method;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+@SuppressWarnings({"UnstableApiUsage", "rawtypes", "unchecked"})
+public class JdbiUnitOfWorkProvider {
+
+    private final Logger log = LoggerFactory.getLogger(JdbiUnitOfWorkProvider.class);
+    private final JdbiHandleManager handleManager;
+
+    private JdbiUnitOfWorkProvider(JdbiHandleManager handleManager) {
+        this.handleManager = handleManager;
+    }
+
+    public static JdbiUnitOfWorkProvider withDefault(DBI dbi) {
+        JdbiHandleManager handleManager = new RequestScopedJdbiHandleManager(dbi);
+        return new JdbiUnitOfWorkProvider(handleManager);
+    }
+
+    public static JdbiUnitOfWorkProvider withLinked(DBI dbi) {
+        JdbiHandleManager handleManager = new LinkedRequestScopedJdbiHandleManager(dbi);
+        return new JdbiUnitOfWorkProvider(handleManager);
+    }
+
+    public JdbiHandleManager getHandleManager() {
+        return handleManager;
+    }
+
+    /**
+     * getWrappedInstanceForDaoClass generates a proxy instance of the dao class for which
+     * the jdbi unit of work aspect would be wrapped around with.
+     * <p>
+     * Note: It is recommended to use {@link JdbiUnitOfWorkProvider#getWrappedInstanceForDaoPackage(List)} instead
+     * as passing a list of packages is easier than passing each instance individually.
+     * <p>
+     * This method however may be used in case the classpath scanning is disabled.
+     * If the original class is null or contains no relevant JDBI annotations, this method throws an
+     * exception
+     *
+     * @param daoClass the DAO class for which a proxy needs to be created fo
+     * @return the wrapped instance ready to be passed around
+     */
+    public Object getWrappedInstanceForDaoClass(Class daoClass) {
+        if (daoClass == null) {
+            throw new IllegalArgumentException("DAO Class cannot be null");
+        }
+        boolean atLeastOneJdbiMethod = false;
+        for (Method method : daoClass.getDeclaredMethods()) {
+            if (method.getDeclaringClass() == daoClass) {
+                atLeastOneJdbiMethod = method.getAnnotation(SqlQuery.class) != null;
+                atLeastOneJdbiMethod = atLeastOneJdbiMethod || method.getAnnotation(SqlUpdate.class) != null;
+                atLeastOneJdbiMethod = atLeastOneJdbiMethod || method.getAnnotation(SqlUpdate.class) != null;
+                atLeastOneJdbiMethod = atLeastOneJdbiMethod || method.getAnnotation(SqlBatch.class) != null;
+                atLeastOneJdbiMethod = atLeastOneJdbiMethod || method.getAnnotation(SqlCall.class) != null;
+            }
+        }
+        if (!atLeastOneJdbiMethod) {
+            throw new IllegalArgumentException(String.format("Class [%s] has no method annotated with a Jdbi SQL Object", daoClass.getSimpleName()));
+        }
+
+        log.info("Binding class [{}] with proxy handler [{}] ", daoClass.getSimpleName(), handleManager.getClass().getSimpleName());
+        ManagedHandleInvocationHandler handler = new ManagedHandleInvocationHandler<>(handleManager, daoClass);
+        Object proxiedInstance = Reflection.newProxy(daoClass, handler);
+        return daoClass.cast(proxiedInstance);
+    }
+
+    /**
+     * getWrappedInstanceForDaoPackage generates a map where every DAO class identified
+     * through the given list of packages is mapped to its initialised proxy instance
+     * the jdbi unit of work aspect would be wrapped around with.
+     * <p>
+     * In case classpath scanning is disabled, use {@link JdbiUnitOfWorkProvider#getWrappedInstanceForDaoClass(Class)}
+     * <p>
+     * If the original package list is null, this method throws an exception
+     *
+     * @param daoPackages the list of packages that contain the DAO classes
+     * @return the map mapping dao classes to its initialised proxies
+     */
+    public Map<? extends Class, Object> getWrappedInstanceForDaoPackage(List<String> daoPackages) {
+        if (daoPackages == null) {
+            throw new IllegalArgumentException("DAO Class package list cannot be null");
+        }
+
+        Set<? extends Class<?>> allDaoClasses = daoPackages.stream()
+            .map(this::getDaoClassesForPackage)
+            .flatMap(Collection::stream)
+            .collect(Collectors.toSet());
+
+        Map<Class, Object> classInstanceMap = new HashMap<>();
+        for (Class klass : allDaoClasses) {
+            log.info("Binding class [{}] with proxy handler [{}] ", klass.getSimpleName(), handleManager.getClass().getSimpleName());
+            Object instance = getWrappedInstanceForDaoClass(klass);
+            classInstanceMap.put(klass, instance);
+        }
+        return classInstanceMap;
+    }
+
+    private Set<? extends Class<?>> getDaoClassesForPackage(String pkg) {
+        Set<Method> daoClasses = new HashSet<>();
+
+        Sets.SetView<Method> union = Sets.union(daoClasses, new Reflections(pkg, Scanners.MethodsAnnotated).getMethodsAnnotatedWith(SqlQuery.class));
+        union = Sets.union(union, new Reflections(pkg, Scanners.MethodsAnnotated).getMethodsAnnotatedWith(SqlUpdate.class));
+        union = Sets.union(union, new Reflections(pkg, Scanners.MethodsAnnotated).getMethodsAnnotatedWith(SqlBatch.class));
+        union = Sets.union(union, new Reflections(pkg, Scanners.MethodsAnnotated).getMethodsAnnotatedWith(SqlCall.class));
+
+        return union.stream()
+            .map(Method::getDeclaringClass)
+            .collect(Collectors.toSet());
+    }
+}