#107 - Polishing.

Extract method references for better readability. Add missing tests and update documentation.
Add delay for transactional MySql tests to avoid failures due to potentially delayed transaction id storage within the database.

Original Pull Request: #107

src/main/java/org/springframework/data/r2dbc/function/connectionfactory/ConnectionFactoryTransactionManager.java

@@ -5,7 +5,7 @@
  * 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
+ *      https://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,

src/main/java/org/springframework/data/r2dbc/function/connectionfactory/ConnectionFactoryUtils.java

@@ -17,14 +17,9 @@ package org.springframework.data.r2dbc.function.connectionfactory;
 
 import io.r2dbc.spi.Connection;
 import io.r2dbc.spi.ConnectionFactory;
-import reactor.core.publisher.Mono;
-import reactor.util.function.Tuple2;
-import reactor.util.function.Tuples;
-
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.reactivestreams.Publisher;
-
 import org.springframework.core.Ordered;
 import org.springframework.dao.DataAccessResourceFailureException;
 import org.springframework.lang.Nullable;
@@ -32,6 +27,9 @@ import org.springframework.transaction.NoTransactionException;
 import org.springframework.transaction.reactive.TransactionSynchronization;
 import org.springframework.transaction.reactive.TransactionSynchronizationManager;
 import org.springframework.util.Assert;
+import reactor.core.publisher.Mono;
+import reactor.util.function.Tuple2;
+import reactor.util.function.Tuples;
 
 /**
  * Helper class that provides static methods for obtaining R2DBC Connections from a
@@ -41,6 +39,7 @@ import org.springframework.util.Assert;
  * objects. Can also be used directly in application code.
  *
  * @author Mark Paluch
+ * @author Christoph Strobl
  */
 public abstract class ConnectionFactoryUtils {
 
@@ -52,7 +51,6 @@ public abstract class ConnectionFactoryUtils {
 	private static final Log logger = LogFactory.getLog(ConnectionFactoryUtils.class);
 
 	private ConnectionFactoryUtils() {
-
 	}
 
 	/**
@@ -93,14 +91,20 @@ public abstract class ConnectionFactoryUtils {
 			if (conHolder != null && (conHolder.hasConnection() || conHolder.isSynchronizedWithTransaction())) {
 				conHolder.requested();
 				if (!conHolder.hasConnection()) {
-					logger.debug("Fetching resumed R2DBC Connection from ConnectionFactory");
+
+					if (logger.isDebugEnabled()) {
+						logger.debug("Fetching resumed R2DBC Connection from ConnectionFactory");
+					}
 					return fetchConnection(connectionFactory).doOnNext(conHolder::setConnection);
 				}
 				return Mono.just(conHolder.getConnection());
 			}
 			// Else we either got no holder or an empty thread-bound holder here.
 
-			logger.debug("Fetching R2DBC Connection from ConnectionFactory");
+			if (logger.isDebugEnabled()) {
+				logger.debug("Fetching R2DBC Connection from ConnectionFactory");
+			}
+
 			Mono<Connection> con = fetchConnection(connectionFactory);
 
 			if (synchronizationManager.isSynchronizationActive()) {
@@ -124,11 +128,9 @@ public abstract class ConnectionFactoryUtils {
 						if (holderToUse != conHolder) {
 							synchronizationManager.bindResource(connectionFactory, holderToUse);
 						}
-
 					}).onErrorResume(e -> {
 						// Unexpected exception from external delegation call -> close Connection and rethrow.
 						return releaseConnection(it, connectionFactory).then(Mono.error(e));
-
 					});
 				});
 			}
@@ -158,7 +160,9 @@ public abstract class ConnectionFactoryUtils {
 
 		if (synchronization.isSynchronizationActive()) {
 
-			logger.debug("Registering transaction synchronization for R2DBC Connection");
+			if (logger.isDebugEnabled()) {
+				logger.debug("Registering transaction synchronization for R2DBC Connection");
+			}
 
 			TransactionResources txContext = synchronization.getCurrentTransaction();
 			ConnectionFactory resource = txContext.getResource(ConnectionFactory.class);
@@ -166,7 +170,9 @@ public abstract class ConnectionFactoryUtils {
 			Mono<Tuple2<Connection, ConnectionFactory>> attachNewConnection = Mono
 					.defer(() -> Mono.from(connectionFactory.create()).map(it -> {
 
-						logger.debug("Fetching new R2DBC Connection from ConnectionFactory");
+						if (logger.isDebugEnabled()) {
+							logger.debug("Fetching new R2DBC Connection from ConnectionFactory");
+						}
 
 						SingletonConnectionFactory s = new SingletonConnectionFactory(connectionFactory.getMetadata(), it);
 						txContext.registerResource(ConnectionFactory.class, s);
@@ -174,14 +180,9 @@ public abstract class ConnectionFactoryUtils {
 						return Tuples.of(it, connectionFactory);
 					}));
 
-			return Mono.justOrEmpty(resource).flatMap(factory -> {
-
-				logger.debug("Fetching resumed R2DBC Connection from ConnectionFactory");
-
-				return Mono.from(factory.create())
-						.map(connection -> Tuples.<Connection, ConnectionFactory> of(connection, factory));
-
-			}).switchIfEmpty(attachNewConnection);
+			return Mono.justOrEmpty(resource)
+					.flatMap(ConnectionFactoryUtils::createConnection)
+					.switchIfEmpty(attachNewConnection);
 		}
 
 		return Mono.empty();
@@ -199,7 +200,7 @@ public abstract class ConnectionFactoryUtils {
 	private static Mono<Connection> fetchConnection(ConnectionFactory connectionFactory) {
 
 		Publisher<? extends Connection> con = connectionFactory.create();
-		if (con == null) {
+		if (con == null) { // TODO: seriously why would it do that?
 			throw new IllegalStateException("ConnectionFactory returned null from getConnection(): " + connectionFactory);
 		}
 		return Mono.from(con);
@@ -211,7 +212,7 @@ public abstract class ConnectionFactoryUtils {
 	 *
 	 * @param con the {@link io.r2dbc.spi.Connection} to close if necessary.
 	 * @param connectionFactory the {@link ConnectionFactory} that the Connection was obtained from (may be
-	 *          {@literal null}).
+	 * {@literal null}).
 	 * @see #getConnection
 	 */
 	public static Mono<Void> releaseConnection(@Nullable io.r2dbc.spi.Connection con,
@@ -227,7 +228,7 @@ public abstract class ConnectionFactoryUtils {
 	 *
 	 * @param con the {@link io.r2dbc.spi.Connection} to close if necessary.
 	 * @param connectionFactory the {@link ConnectionFactory} that the Connection was obtained from (may be
-	 *          {@literal null}).
+	 * {@literal null}).
 	 * @see #doGetConnection
 	 */
 	public static Mono<Void> doReleaseConnection(@Nullable io.r2dbc.spi.Connection con,
@@ -247,12 +248,16 @@ public abstract class ConnectionFactoryUtils {
 
 				SingletonConnectionFactory factory = (SingletonConnectionFactory) connectionFactory;
 
-				logger.debug("Releasing R2DBC Connection");
+				if (logger.isDebugEnabled()) {
+					logger.debug("Releasing R2DBC Connection");
+				}
 
 				return factory.close(con);
 			}
 
-			logger.debug("Closing R2DBC Connection");
+			if (logger.isDebugEnabled()) {
+				logger.debug("Closing R2DBC Connection");
+			}
 
 			return Mono.from(con.close());
 		});
@@ -267,6 +272,7 @@ public abstract class ConnectionFactoryUtils {
 	 * @throws DataAccessResourceFailureException if the attempt to get a {@link io.r2dbc.spi.Connection} failed
 	 */
 	public static Mono<Void> closeConnection(Connection connection, ConnectionFactory connectionFactory) {
+
 		return doCloseConnection(connection, connectionFactory)
 				.onErrorMap(e -> new DataAccessResourceFailureException("Failed to obtain R2DBC Connection", e));
 	}
@@ -295,10 +301,10 @@ public abstract class ConnectionFactoryUtils {
 	 * Obtain the currently {@link ReactiveTransactionSynchronization} from the current subscriber
 	 * {@link reactor.util.context.Context}.
 	 *
+	 * @throws NoTransactionException if no active {@link ReactiveTransactionSynchronization} is associated with the
+	 * current subscription.
 	 * @see Mono#subscriberContext()
 	 * @see ReactiveTransactionSynchronization
-	 * @throws NoTransactionException if no active {@link ReactiveTransactionSynchronization} is associated with the
-	 *           current subscription.
 	 */
 	public static Mono<ReactiveTransactionSynchronization> currentReactiveTransactionSynchronization() {
 
@@ -312,10 +318,10 @@ public abstract class ConnectionFactoryUtils {
 	 * Obtain the currently active {@link ReactiveTransactionSynchronization} from the current subscriber
 	 * {@link reactor.util.context.Context}.
 	 *
+	 * @throws NoTransactionException if no active {@link ReactiveTransactionSynchronization} is associated with the
+	 * current subscription.
 	 * @see Mono#subscriberContext()
 	 * @see ReactiveTransactionSynchronization
-	 * @throws NoTransactionException if no active {@link ReactiveTransactionSynchronization} is associated with the
-	 *           current subscription.
 	 */
 	public static Mono<ReactiveTransactionSynchronization> currentActiveReactiveTransactionSynchronization() {
 
@@ -341,16 +347,8 @@ public abstract class ConnectionFactoryUtils {
 						return true;
 					}
 					return false;
-				}).map(it -> connectionFactory).onErrorResume(NoTransactionException.class, e -> {
-
-					return currentActiveReactiveTransactionSynchronization().map(synchronization -> {
-
-						TransactionResources currentSynchronization = synchronization.getCurrentTransaction();
-						return currentSynchronization.getResource(ConnectionFactory.class);
-					}).switchIfEmpty(Mono.error(new DataAccessResourceFailureException(
-							"Cannot extract ConnectionFactory from current TransactionContext!")));
-
-				});
+				}).map(it -> connectionFactory) //
+				.onErrorResume(NoTransactionException.class, ConnectionFactoryUtils::obtainDefaultConnectionFactory);
 	}
 
 	/**
@@ -359,7 +357,7 @@ public abstract class ConnectionFactoryUtils {
 	 *
 	 * @param conHolder the {@link ConnectionHolder} for the held Connection (potentially a proxy)
 	 * @param passedInCon the {@link Connection} passed-in by the user (potentially a target {@link Connection} without
-	 *          proxy)
+	 * proxy)
 	 * @return whether the given Connections are equal
 	 * @see #getTargetConnection
 	 */
@@ -402,6 +400,7 @@ public abstract class ConnectionFactoryUtils {
 	 * @see #CONNECTION_SYNCHRONIZATION_ORDER
 	 */
 	private static int getConnectionSynchronizationOrder(ConnectionFactory connectionFactory) {
+
 		int order = CONNECTION_SYNCHRONIZATION_ORDER;
 		ConnectionFactory current = connectionFactory;
 		while (current instanceof DelegatingConnectionFactory) {
@@ -411,6 +410,37 @@ public abstract class ConnectionFactoryUtils {
 		return order;
 	}
 
+	/**
+	 * @param e
+	 * @return an {@link Mono#error(Throwable) error} if not transaction present.
+	 */
+	private static Mono<? extends ConnectionFactory> obtainDefaultConnectionFactory(NoTransactionException e) {
+
+		return currentActiveReactiveTransactionSynchronization().map(synchronization -> {
+
+			TransactionResources currentSynchronization = synchronization.getCurrentTransaction();
+			return currentSynchronization.getResource(ConnectionFactory.class);
+		}).switchIfEmpty(Mono.error(new DataAccessResourceFailureException(
+				"Cannot extract ConnectionFactory from current TransactionContext!")));
+	}
+
+	/**
+	 * Create a {@link Connection} via the given {@link ConnectionFactory#create() factory} and return a {@link Tuple2} associating the
+	 * {@link Connection} with its creating {@link ConnectionFactory}.
+	 *
+	 * @param factory must not be {@literal null}.
+	 * @return never {@literal null}
+	 */
+	private static Mono<Tuple2<Connection, ConnectionFactory>> createConnection(ConnectionFactory factory) {
+
+		if (logger.isDebugEnabled()) {
+			logger.debug("Fetching resumed R2DBC Connection from ConnectionFactory");
+		}
+
+		return Mono.from(factory.create())
+				.map(connection -> Tuples.of(connection, factory));
+	}
+
 	/**
 	 * Callback for resource cleanup at the end of a non-native R2DBC transaction.
 	 */
@@ -465,7 +495,6 @@ public abstract class ConnectionFactoryUtils {
 				}).then();
 			}
 			return Mono.empty();
-
 		}
 
 		@Override
@@ -510,17 +539,14 @@ public abstract class ConnectionFactoryUtils {
 									// Reset the ConnectionHolder: It might remain bound to the context.
 									this.connectionHolder.setConnection(null);
 								});
-
 					}
 
 					return Mono.empty();
 				});
-
 			}
 
 			this.connectionHolder.reset();
 			return Mono.empty();
 		}
 	}
-
 }

src/main/java/org/springframework/data/r2dbc/function/connectionfactory/ConnectionHandle.java

@@ -5,7 +5,7 @@
  * 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
+ *      https://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,
@@ -38,8 +38,8 @@ public interface ConnectionHandle {
 	 * <p>
 	 * The default implementation is empty, assuming that the lifecycle of the connection is managed externally.
 	 *
-	 * @param con the R2DBC Connection to release
+	 * @param connection the R2DBC Connection to release
 	 */
-	default void releaseConnection(Connection con) {}
-
+	default void releaseConnection(Connection connection) {
+	}
 }

src/main/java/org/springframework/data/r2dbc/function/connectionfactory/ConnectionHolder.java

@@ -5,7 +5,7 @@
  * 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
+ *      https://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,
@@ -17,7 +17,6 @@ package org.springframework.data.r2dbc.function.connectionfactory;
 
 import io.r2dbc.spi.Connection;
 import io.r2dbc.spi.ConnectionFactory;
-
 import org.springframework.lang.Nullable;
 import org.springframework.transaction.support.ResourceHolderSupport;
 import org.springframework.util.Assert;
@@ -30,6 +29,8 @@ import org.springframework.util.Assert;
  * <p>
  * Note: This is an SPI class, not intended to be used by applications.
  *
+ * @author Mark Paluch
+ * @author Christoph Strobl
  * @see ConnectionFactoryTransactionManager
  * @see ConnectionFactoryUtils
  */
@@ -39,7 +40,7 @@ public class ConnectionHolder extends ResourceHolderSupport {
 
 	@Nullable private Connection currentConnection;
 
-	private boolean transactionActive = false;
+	private boolean transactionActive;
 
 	/**
 	 * Create a new ConnectionHolder for the given R2DBC {@link Connection}, wrapping it with a
@@ -50,7 +51,7 @@ public class ConnectionHolder extends ResourceHolderSupport {
 	 * @see #ConnectionHolder(Connection, boolean)
 	 */
 	public ConnectionHolder(Connection connection) {
-		this.connectionHandle = new SimpleConnectionHandle(connection);
+		this(connection, false);
 	}
 
 	/**
@@ -62,7 +63,8 @@ public class ConnectionHolder extends ResourceHolderSupport {
 	 * @see SimpleConnectionHandle
 	 */
 	public ConnectionHolder(Connection connection, boolean transactionActive) {
-		this(connection);
+
+		this.connectionHandle = new SimpleConnectionHandle(connection);
 		this.transactionActive = transactionActive;
 	}
 
@@ -127,6 +129,7 @@ public class ConnectionHolder extends ResourceHolderSupport {
 	 * @see #released()
 	 */
 	public Connection getConnection() {
+
 		Assert.notNull(this.connectionHandle, "Active Connection is required");
 		if (this.currentConnection == null) {
 			this.currentConnection = this.connectionHandle.getConnection();
@@ -152,7 +155,7 @@ public class ConnectionHolder extends ResourceHolderSupport {
 		}
 	}
 
-	/* 
+	/*
 	 * (non-Javadoc)
 	 * @see org.springframework.transaction.support.ResourceHolderSupport#clear()
 	 */

src/main/java/org/springframework/data/r2dbc/function/connectionfactory/ConnectionProxy.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2018-2019 the original author or authors.
+ * Copyright 2019 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -19,12 +19,13 @@ import io.r2dbc.spi.Connection;
 import io.r2dbc.spi.Wrapped;
 
 /**
- * Subinterface of {@link Connection} to be implemented by Connection proxies. Allows access to the underlying target
+ * Sub interface of {@link Connection} to be implemented by Connection proxies. Allows access to the underlying target
  * Connection.
  * <p/>
  * This interface can be checked when there is a need to cast to a native R2DBC {@link Connection}.
  *
  * @author Mark Paluch
+ * @author Christoph Strobl
  */
 public interface ConnectionProxy extends Connection, Wrapped<Connection> {
 
@@ -34,6 +35,7 @@ public interface ConnectionProxy extends Connection, Wrapped<Connection> {
 	 * This will typically be the native driver {@link Connection} or a wrapper from a connection pool.
 	 *
 	 * @return the underlying Connection (never {@literal null})
+	 * @throws IllegalStateException in case the connection has already been closed.
 	 */
 	Connection getTargetConnection();
 }

src/main/java/org/springframework/data/r2dbc/function/connectionfactory/DelegatingConnectionFactory.java

@@ -5,7 +5,7 @@
  * 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
+ *      https://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,
@@ -19,7 +19,6 @@ import io.r2dbc.spi.Connection;
 import io.r2dbc.spi.ConnectionFactory;
 import io.r2dbc.spi.ConnectionFactoryMetadata;
 import io.r2dbc.spi.Wrapped;
-
 import org.reactivestreams.Publisher;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
@@ -43,7 +42,7 @@ public class DelegatingConnectionFactory implements ConnectionFactory, Wrapped<C
 		this.targetConnectionFactory = targetConnectionFactory;
 	}
 
-	/* 
+	/*
 	 * (non-Javadoc)
 	 * @see io.r2dbc.spi.ConnectionFactory#create()
 	 */
@@ -52,13 +51,6 @@ public class DelegatingConnectionFactory implements ConnectionFactory, Wrapped<C
 		return targetConnectionFactory.create();
 	}
 
-	/**
-	 * Obtain the target {@link ConnectionFactory} for actual use (never {@code null}).
-	 */
-	protected ConnectionFactory obtainTargetConnectionFactory() {
-		return getTargetConnectionFactory();
-	}
-
 	/**
 	 * Return the target {@link ConnectionFactory} that this {@link ConnectionFactory} should delegate to.
 	 */
@@ -67,7 +59,7 @@ public class DelegatingConnectionFactory implements ConnectionFactory, Wrapped<C
 		return this.targetConnectionFactory;
 	}
 
-	/* 
+	/*
 	 * (non-Javadoc)
 	 * @see io.r2dbc.spi.ConnectionFactory#getMetadata()
 	 */
@@ -76,7 +68,7 @@ public class DelegatingConnectionFactory implements ConnectionFactory, Wrapped<C
 		return obtainTargetConnectionFactory().getMetadata();
 	}
 
-	/* 
+	/*
 	 * (non-Javadoc)
 	 * @see io.r2dbc.spi.Wrapped#unwrap()
 	 */
@@ -84,4 +76,11 @@ public class DelegatingConnectionFactory implements ConnectionFactory, Wrapped<C
 	public ConnectionFactory unwrap() {
 		return obtainTargetConnectionFactory();
 	}
+
+	/**
+	 * Obtain the target {@link ConnectionFactory} for actual use (never {@code null}).
+	 */
+	protected ConnectionFactory obtainTargetConnectionFactory() {
+		return getTargetConnectionFactory();
+	}
 }

src/main/java/org/springframework/data/r2dbc/function/connectionfactory/R2dbcTransactionObjectSupport.java

@@ -5,7 +5,7 @@
  * 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
+ *      https://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,
@@ -16,13 +16,12 @@
 package org.springframework.data.r2dbc.function.connectionfactory;
 
 import io.r2dbc.spi.IsolationLevel;
-
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 
 /**
  * Convenient base class for R2DBC-aware transaction objects. Can contain a {@link ConnectionHolder} with a R2DBC
- * {@link Connection}.
+ * {@link io.r2dbc.spi.Connection}.
  *
  * @author Mark Paluch
  * @see ConnectionFactoryTransactionManager

src/main/java/org/springframework/data/r2dbc/function/connectionfactory/SimpleConnectionHandle.java

@@ -5,7 +5,7 @@
  * 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
+ *      https://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,
@@ -23,8 +23,9 @@ import org.springframework.util.Assert;
  * Simple implementation of the {@link ConnectionHandle} interface, containing a given R2DBC Connection.
  *
  * @author Mark Paluch
+ * @author Christoph Strobl
  */
-public class SimpleConnectionHandle implements ConnectionHandle {
+class SimpleConnectionHandle implements ConnectionHandle {
 
 	private final Connection connection;
 
@@ -33,7 +34,8 @@ public class SimpleConnectionHandle implements ConnectionHandle {
 	 *
 	 * @param connection the R2DBC Connection
 	 */
-	public SimpleConnectionHandle(Connection connection) {
+	SimpleConnectionHandle(Connection connection) {
+
 		Assert.notNull(connection, "Connection must not be null");
 		this.connection = connection;
 	}

src/main/java/org/springframework/data/r2dbc/function/connectionfactory/TransactionAwareConnectionFactoryProxy.java

@@ -5,7 +5,7 @@
  * 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
+ *      https://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,
@@ -15,18 +15,21 @@
  */
 package org.springframework.data.r2dbc.function.connectionfactory;
 
-import io.r2dbc.spi.Connection;
-import io.r2dbc.spi.ConnectionFactory;
-import io.r2dbc.spi.Wrapped;
-import reactor.core.publisher.Mono;
+import static org.springframework.util.ReflectionUtils.*;
 
 import java.lang.reflect.InvocationHandler;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.lang.reflect.Proxy;
 
+import io.r2dbc.spi.Connection;
+import io.r2dbc.spi.ConnectionFactory;
+import io.r2dbc.spi.Wrapped;
 import org.springframework.data.r2dbc.function.DatabaseClient;
 import org.springframework.lang.Nullable;
+import org.springframework.util.ReflectionUtils;
+import reactor.core.publisher.Mono;
+import reactor.util.function.Tuple2;
 
 /**
  * Proxy for a target R2DBC {@link ConnectionFactory}, adding awareness of Spring-managed transactions.
@@ -54,6 +57,7 @@ import org.springframework.lang.Nullable;
  * the native R2DBC Connection.
  *
  * @author Mark Paluch
+ * @author Christoph Strobl
  * @see ConnectionFactory#create
  * @see Connection#close
  * @see ConnectionFactoryUtils#doGetConnection
@@ -65,6 +69,7 @@ public class TransactionAwareConnectionFactoryProxy extends DelegatingConnection
 	 * Create a new {@link TransactionAwareConnectionFactoryProxy}.
 	 *
 	 * @param targetConnectionFactory the target {@link ConnectionFactory}.
+	 * @throws IllegalArgumentException if given {@link ConnectionFactory} is {@literal null}.
 	 */
 	public TransactionAwareConnectionFactoryProxy(ConnectionFactory targetConnectionFactory) {
 		super(targetConnectionFactory);
@@ -95,19 +100,21 @@ public class TransactionAwareConnectionFactoryProxy extends DelegatingConnection
 	 * @see ConnectionFactoryUtils#doReleaseConnection
 	 */
 	protected Mono<Connection> getTransactionAwareConnectionProxy(ConnectionFactory targetConnectionFactory) {
+		return ConnectionFactoryUtils.getConnection(targetConnectionFactory).map(TransactionAwareConnectionFactoryProxy::proxyConnection);
+	}
 
-		return ConnectionFactoryUtils.getConnection(targetConnectionFactory).map(tuple -> {
-			return (Connection) Proxy.newProxyInstance(ConnectionProxy.class.getClassLoader(),
-					new Class<?>[] { ConnectionProxy.class },
-					new TransactionAwareInvocationHandler(tuple.getT1(), targetConnectionFactory));
-		});
+	private static Connection proxyConnection(Tuple2<Connection, ConnectionFactory> connectionConnectionFactoryTuple) {
+
+		return (Connection) Proxy.newProxyInstance(ConnectionProxy.class.getClassLoader(),
+				new Class<?>[]{ConnectionProxy.class},
+				new TransactionAwareInvocationHandler(connectionConnectionFactoryTuple.getT1(), connectionConnectionFactoryTuple.getT2()));
 	}
 
 	/**
 	 * Invocation handler that delegates close calls on R2DBC Connections to {@link ConnectionFactoryUtils} for being
 	 * aware of context-bound transactions.
 	 */
-	private class TransactionAwareInvocationHandler implements InvocationHandler {
+	private static class TransactionAwareInvocationHandler implements InvocationHandler {
 
 		private final Connection connection;
 
@@ -116,11 +123,12 @@ public class TransactionAwareConnectionFactoryProxy extends DelegatingConnection
 		private boolean closed = false;
 
 		TransactionAwareInvocationHandler(Connection connection, ConnectionFactory targetConnectionFactory) {
+
 			this.connection = connection;
 			this.targetConnectionFactory = targetConnectionFactory;
 		}
 
-		/* 
+		/*
 		 * (non-Javadoc)
 		 * @see java.lang.reflect.InvocationHandler#invoke(java.lang.Object, java.lang.reflect.Method, java.lang.Object[])
 		 */
@@ -128,23 +136,24 @@ public class TransactionAwareConnectionFactoryProxy extends DelegatingConnection
 		@Nullable
 		public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
 
+			if (ReflectionUtils.isObjectMethod(method)) {
+
+				if (ReflectionUtils.isToStringMethod(method)) {
+					return proxyToString(proxy);
+				}
+
+				if (ReflectionUtils.isEqualsMethod(method)) {
+					return (proxy == args[0]);
+				}
+
+				if (ReflectionUtils.isHashCodeMethod(method)) {
+					return System.identityHashCode(proxy);
+				}
+			}
+
 			// Invocation on ConnectionProxy interface coming in...
 			switch (method.getName()) {
-				case "equals":
-					// Only considered as equal when proxies are identical.
-					return (proxy == args[0]);
-				case "hashCode":
-					// Use hashCode of Connection proxy.
-					return System.identityHashCode(proxy);
-				case "toString":
-					// Allow for differentiating between the proxy and the raw Connection.
-					StringBuilder sb = new StringBuilder("Transaction-aware proxy for target Connection ");
-					if (this.connection != null) {
-						sb.append("[").append(this.connection.toString()).append("]");
-					} else {
-						sb.append(" from ConnectionFactory [").append(this.targetConnectionFactory).append("]");
-					}
-					return sb.toString();
+
 				case "unwrap":
 					return this.connection;
 				case "close":
@@ -171,5 +180,17 @@ public class TransactionAwareConnectionFactoryProxy extends DelegatingConnection
 				throw ex.getTargetException();
 			}
 		}
+
+		private String proxyToString(@Nullable Object proxy) {
+
+			// Allow for differentiating between the proxy and the raw Connection.
+			StringBuilder sb = new StringBuilder("Transaction-aware proxy for target Connection ");
+			if (this.connection != null) {
+				sb.append("[").append(this.connection.toString()).append("]");
+			} else {
+				sb.append(" from ConnectionFactory [").append(this.targetConnectionFactory).append("]");
+			}
+			return sb.toString();
+		}
 	}
 }