From 0944487390685b5a6059472d20f20509350f87e1 Mon Sep 17 00:00:00 2001
From: Scott Leberknight <174812+sleberknight@users.noreply.github.com>
Date: Thu, 8 Feb 2024 16:50:36 -0500
Subject: [PATCH] Code cleanup: UnitOfWorkInvoker and UnitOfWorkInvokerFactory
 (#98)

* Update obsolete javadoc comment in UnitOfWorkInvoker. Dropwizard made
UnitOfWorkAspect public a long time ago, in version 1.1.0.
* Modify UnitOfWorkInvoker#invoke to use try-with-resources for the
Hibernate Session, and wrap the Map#get call with requireNonNull since
the containsKey check returned true. Remove redundant RuntimeException
type parameter in call to rethrow.
* Use var in UnitOfWorkInvokerFactory for the ImmutableMap.Builder to
reduce code verbosity
* Refactor logic in UnitOfWorkInvokerFactory so that it doesn't reassign
a local variable. Instead, if the "unitOfWorkMethods" map is empty,
return the rootInvoker. This inverts the logic and is clearer, since
re-assigning variables is generally not a good practice.
---
 .../jakarta/xml/ws/UnitOfWorkInvoker.java     | 25 ++++++++++++-------
 .../xml/ws/UnitOfWorkInvokerFactory.java      | 11 +++-----
 2 files changed, 20 insertions(+), 16 deletions(-)

diff --git a/dropwizard-jakarta-xml-ws/src/main/java/org/kiwiproject/dropwizard/jakarta/xml/ws/UnitOfWorkInvoker.java b/dropwizard-jakarta-xml-ws/src/main/java/org/kiwiproject/dropwizard/jakarta/xml/ws/UnitOfWorkInvoker.java
index 682cfdf..31e5605 100644
--- a/dropwizard-jakarta-xml-ws/src/main/java/org/kiwiproject/dropwizard/jakarta/xml/ws/UnitOfWorkInvoker.java
+++ b/dropwizard-jakarta-xml-ws/src/main/java/org/kiwiproject/dropwizard/jakarta/xml/ws/UnitOfWorkInvoker.java
@@ -1,5 +1,7 @@
 package org.kiwiproject.dropwizard.jakarta.xml.ws;
 
+import static java.util.Objects.requireNonNull;
+
 import com.google.common.collect.ImmutableMap;
 import io.dropwizard.hibernate.UnitOfWork;
 import org.apache.cxf.message.Exchange;
@@ -12,7 +14,15 @@
 
 /**
  * Wraps underlying invoker in a Hibernate session. Code in this class is based on Dropwizard's UnitOfWorkApplication
- * listener and UnitOfWorkAspect. We don't use UnitOfWorkAspect here because it is declared package private.
+ * listener and UnitOfWorkAspect.
+ * <p>
+ * <strong>Background information:</strong>
+ * The javadocs used to state "We don't use UnitOfWorkAspect here because it is declared package private." This
+ * was true up until Dropwizard 1.1.0, which made UnitOfWorkAspect public in
+ * <a href="https://github.com/dropwizard/dropwizard/pull/1661">this PR</a>. See
+ * <a href="https://github.com/kiwiproject/dropwizard-jakarta-xml-ws/discussions/91">this discussion</a> which
+ * proposes to change this class to use UnitOfWorkAspect directly.
+ * 
  *
  * @see io.dropwizard.hibernate.UnitOfWorkAspect
  * @see io.dropwizard.hibernate.UnitOfWorkApplicationListener
@@ -34,14 +44,12 @@ public UnitOfWorkInvoker(Invoker underlying, ImmutableMap<String, UnitOfWork> un
     public Object invoke(Exchange exchange, Object o) {
 
         Object result;
-        String methodname = this.getTargetMethod(exchange).getName();
-
-        if (unitOfWorkMethods.containsKey(methodname)) {
+        String methodName = this.getTargetMethod(exchange).getName();
 
-            final Session session = sessionFactory.openSession();
-            UnitOfWork unitOfWork = unitOfWorkMethods.get(methodname);
+        if (unitOfWorkMethods.containsKey(methodName)) {
 
-            try {
+            try (var session = sessionFactory.openSession()) {
+                UnitOfWork unitOfWork = requireNonNull(unitOfWorkMethods.get(methodName));
                 configureSession(session, unitOfWork);
                 ManagedSessionContext.bind(session);
                 beginTransaction(session, unitOfWork);
@@ -51,11 +59,10 @@ public Object invoke(Exchange exchange, Object o) {
                     return result;
                 } catch (Exception e) {
                     rollbackTransaction(session, unitOfWork);
-                    this.<RuntimeException>rethrow(e); // unchecked rethrow
+                    this.rethrow(e); // unchecked rethrow
                     return null; // avoid compiler warning
                 }
             } finally {
-                session.close();
                 ManagedSessionContext.unbind(sessionFactory);
             }
         } else {
diff --git a/dropwizard-jakarta-xml-ws/src/main/java/org/kiwiproject/dropwizard/jakarta/xml/ws/UnitOfWorkInvokerFactory.java b/dropwizard-jakarta-xml-ws/src/main/java/org/kiwiproject/dropwizard/jakarta/xml/ws/UnitOfWorkInvokerFactory.java
index d22c183..cea3c8c 100644
--- a/dropwizard-jakarta-xml-ws/src/main/java/org/kiwiproject/dropwizard/jakarta/xml/ws/UnitOfWorkInvokerFactory.java
+++ b/dropwizard-jakarta-xml-ws/src/main/java/org/kiwiproject/dropwizard/jakarta/xml/ws/UnitOfWorkInvokerFactory.java
@@ -14,8 +14,7 @@ public class UnitOfWorkInvokerFactory {
      */
     public Invoker create(Object service, Invoker rootInvoker, SessionFactory sessionFactory) {
 
-        ImmutableMap.Builder<String, UnitOfWork> unitOfWorkMethodsBuilder =
-                new ImmutableMap.Builder<>();
+        var unitOfWorkMethodsBuilder = new ImmutableMap.Builder<String, UnitOfWork>();
 
         for (Method m : service.getClass().getMethods()) {
             if (m.isAnnotationPresent(UnitOfWork.class)) {
@@ -24,13 +23,11 @@ public Invoker create(Object service, Invoker rootInvoker, SessionFactory sessio
         }
         ImmutableMap<String, UnitOfWork> unitOfWorkMethods = unitOfWorkMethodsBuilder.build();
 
-        Invoker invoker = rootInvoker;
-
-        if (!unitOfWorkMethods.isEmpty()) {
-            invoker = new UnitOfWorkInvoker(invoker, unitOfWorkMethods, sessionFactory);
+        if (unitOfWorkMethods.isEmpty()) {
+            return rootInvoker;
         }
 
-        return invoker;
+        return new UnitOfWorkInvoker(rootInvoker, unitOfWorkMethods, sessionFactory);
     }
 
 }