Support custom callback handlers

Introduce new interface that is an analog of `TransactionHandler`
but for all callbacks, not just transactions. This interface,
`Handler` is used to invoke the callback passed to JDBI's
`useHandle`, `withHandle`, `inTransaction` and `withTransaction`.
This gives users a chance to global patch callbacks for all
JDBI uses not just transactions.

core/src/main/java/org/jdbi/v3/core/Handler.java

@@ -0,0 +1,35 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jdbi.v3.core;
+
+public interface Handler {
+    Handler STANDARD_HANDLER = new Handler() {
+        @Override
+        public <R, X extends Exception> HandleCallback<R, X> decorate(HandleCallback<R, X> callback) {
+            return callback;
+        }
+    };
+
+    /**
+     * Decorate the given Handle callback
+     *
+     * @param callback a callback which will receive the open handle
+     *
+     * @return the callback decorated as needed
+     *
+     * @throws X any exception thrown by the callback.
+     * @see Jdbi#withHandle(HandleCallback)
+     */
+    <R, X extends Exception> HandleCallback<R, X> decorate(HandleCallback<R, X> callback);
+}

core/src/main/java/org/jdbi/v3/core/Jdbi.java

@@ -61,6 +61,7 @@ public class Jdbi implements Configurable<Jdbi> {
     private final ConnectionFactory connectionFactory;
     private final AtomicReference<TransactionHandler> transactionhandler = new AtomicReference<>(LocalTransactionHandler.binding());
     private final AtomicReference<StatementBuilderFactory> statementBuilderFactory = new AtomicReference<>(DefaultStatementBuilder.FACTORY);
+    private final AtomicReference<Handler> handler = new AtomicReference<>(Handler.STANDARD_HANDLER);
 
     private final CopyOnWriteArrayList<JdbiPlugin> plugins = new CopyOnWriteArrayList<>();
 
@@ -303,6 +304,30 @@ public TransactionHandler getTransactionHandler() {
         return this.transactionhandler.get();
     }
 
+    /**
+     * Specify the {@link Handler} instance to use. This allows overriding
+     * callbacks for {@link #useHandle}, {@link #withHandle}, {@link #useTransaction} and
+     * {@link #inTransaction}. The default version is a pass-through that returns the callback unchanged.
+     *
+     * @param handler The {@link Handler} to use for all {@link #useHandle}, {@link #withHandle},
+     *                {@link #useTransaction} and {@link #inTransaction} from this Jdbi
+     * @return this
+     */
+    public Jdbi setHandler(Handler handler) {
+        Objects.requireNonNull(handler, "null handler");
+        this.handler.set(handler);
+        return this;
+    }
+
+    /**
+     * Returns the {@link Handler}.
+     *
+     * @return the {@link Handler}
+     */
+    public Handler getHandler() {
+        return this.handler.get();
+    }
+
     /**
      * Obtain a Handle to the data source wrapped by this Jdbi instance.
      * You own this expensive resource and are required to close it or
@@ -358,8 +383,10 @@ public <R, X extends Exception> R withHandle(HandleCallback<R, X> callback) thro
 
         HandleSupplier handleSupplier = threadHandleSupplier.get();
 
+        HandleCallback<R, X> decoratedCallback = handler.get().decorate(callback);
+
         if (handleSupplier != null) {
-            return callback.withHandle(handleSupplier.getHandle());
+            return decoratedCallback.withHandle(handleSupplier.getHandle());
         }
 
         try (Handle h = this.open()) {
@@ -368,7 +395,7 @@ public <R, X extends Exception> R withHandle(HandleCallback<R, X> callback) thro
 
             handleSupplier = ConstantHandleSupplier.of(h);
             threadHandleSupplier.set(handleSupplier);
-            return callback.withHandle(h);
+            return decoratedCallback.withHandle(h);
         } finally {
             threadHandleSupplier.remove();
         }

core/src/test/java/org/jdbi/v3/core/TestHandle.java

@@ -14,6 +14,7 @@
 package org.jdbi.v3.core;
 
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
 
 import org.jdbi.v3.core.junit5.H2DatabaseExtension;
 import org.jdbi.v3.core.statement.UnableToCreateStatementException;
@@ -215,6 +216,55 @@ public void testAllNestedOpsSameHandle() {
                     }))));
     }
 
+    @Test
+    public void testCustomHandler() {
+        AtomicReference<HandleCallback<?, ?>> overrideCallback = new AtomicReference<>();
+        AtomicBoolean gotCalled = new AtomicBoolean();
+        Handler testHandler = new Handler() {
+            @SuppressWarnings("unchecked")
+            @Override
+            public <R, X extends Exception> HandleCallback<R, X> decorate(HandleCallback<R, X> callback) {
+                gotCalled.set(true);
+                HandleCallback<?, ?> testCallback = overrideCallback.get();
+                if (testCallback != null) {
+                    return (HandleCallback<R, X>) testCallback;
+                }
+                return callback;
+            }
+        };
+
+        Jdbi jdbi = h2Extension.getJdbi();
+        Handler saveHandler = jdbi.getHandler();
+        jdbi.setHandler(testHandler);
+        try {
+            String result = jdbi.withHandle(handle -> "hey");
+            assertThat(result).isEqualTo("hey");
+            assertThat(gotCalled).isTrue();
+            gotCalled.set(false);
+
+            jdbi.useHandle(handle -> {});
+            assertThat(gotCalled).isTrue();
+            gotCalled.set(false);
+
+            result = jdbi.inTransaction(handle -> "there");
+            assertThat(result).isEqualTo("there");
+            assertThat(gotCalled).isTrue();
+
+            jdbi.useTransaction(handle -> {});
+            assertThat(gotCalled).isTrue();
+            gotCalled.set(false);
+
+            overrideCallback.set(handle -> "this is different");
+            result = jdbi.inTransaction(handle -> "you");
+            assertThat(result).isEqualTo("this is different");
+            assertThat(gotCalled).isTrue();
+            gotCalled.set(false);
+            overrideCallback.set(null);
+        } finally {
+            jdbi.setHandler(saveHandler);
+        }
+    }
+
     static class BoomHandler extends LocalTransactionHandler {
         boolean failTest;
         boolean failRollback;