From 0af874cd54176b2bb35f81bd7cfacb0196caebb3 Mon Sep 17 00:00:00 2001 From: Jordan Zimmerman Date: Mon, 24 Jul 2023 09:05:11 +0100 Subject: [PATCH] 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. --- .../main/java/org/jdbi/v3/core/Handler.java | 35 +++++++++++++ core/src/main/java/org/jdbi/v3/core/Jdbi.java | 31 +++++++++++- .../java/org/jdbi/v3/core/TestHandle.java | 50 +++++++++++++++++++ 3 files changed, 114 insertions(+), 2 deletions(-) create mode 100644 core/src/main/java/org/jdbi/v3/core/Handler.java diff --git a/core/src/main/java/org/jdbi/v3/core/Handler.java b/core/src/main/java/org/jdbi/v3/core/Handler.java new file mode 100644 index 0000000000..e1d09ed808 --- /dev/null +++ b/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 HandleCallback decorate(HandleCallback 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) + */ + HandleCallback decorate(HandleCallback callback); +} diff --git a/core/src/main/java/org/jdbi/v3/core/Jdbi.java b/core/src/main/java/org/jdbi/v3/core/Jdbi.java index bbe7bba268..fcfe218dcc 100644 --- a/core/src/main/java/org/jdbi/v3/core/Jdbi.java +++ b/core/src/main/java/org/jdbi/v3/core/Jdbi.java @@ -61,6 +61,7 @@ public class Jdbi implements Configurable { private final ConnectionFactory connectionFactory; private final AtomicReference transactionhandler = new AtomicReference<>(LocalTransactionHandler.binding()); private final AtomicReference statementBuilderFactory = new AtomicReference<>(DefaultStatementBuilder.FACTORY); + private final AtomicReference handler = new AtomicReference<>(Handler.STANDARD_HANDLER); private final CopyOnWriteArrayList 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 withHandle(HandleCallback callback) thro HandleSupplier handleSupplier = threadHandleSupplier.get(); + HandleCallback 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 withHandle(HandleCallback callback) thro handleSupplier = ConstantHandleSupplier.of(h); threadHandleSupplier.set(handleSupplier); - return callback.withHandle(h); + return decoratedCallback.withHandle(h); } finally { threadHandleSupplier.remove(); } diff --git a/core/src/test/java/org/jdbi/v3/core/TestHandle.java b/core/src/test/java/org/jdbi/v3/core/TestHandle.java index 89fac1e3dd..2e1c094e6d 100644 --- a/core/src/test/java/org/jdbi/v3/core/TestHandle.java +++ b/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> overrideCallback = new AtomicReference<>(); + AtomicBoolean gotCalled = new AtomicBoolean(); + Handler testHandler = new Handler() { + @SuppressWarnings("unchecked") + @Override + public HandleCallback decorate(HandleCallback callback) { + gotCalled.set(true); + HandleCallback testCallback = overrideCallback.get(); + if (testCallback != null) { + return (HandleCallback) 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;