#204 - Add SingleConnectionConnectionFactory.

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

@@ -19,7 +19,8 @@ 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 reactor.core.publisher.Mono;
+
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 
@@ -47,8 +48,8 @@ public class DelegatingConnectionFactory implements ConnectionFactory, Wrapped<C
 	 * @see io.r2dbc.spi.ConnectionFactory#create()
 	 */
 	@Override
-	public Publisher<? extends Connection> create() {
-		return targetConnectionFactory.create();
+	public Mono<? extends Connection> create() {
+		return Mono.from(targetConnectionFactory.create());
 	}
 
 	/**

src/main/java/org/springframework/data/r2dbc/connectionfactory/SingleConnectionConnectionFactory.java

@@ -0,0 +1,291 @@
+/*
+ * 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.
+ * You may obtain a copy of the License at
+ *
+ *      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,
+ * 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.springframework.data.r2dbc.connectionfactory;
+
+import io.r2dbc.spi.Connection;
+import io.r2dbc.spi.ConnectionFactories;
+import io.r2dbc.spi.ConnectionFactory;
+import io.r2dbc.spi.ConnectionFactoryMetadata;
+import reactor.core.publisher.Mono;
+
+import java.lang.reflect.InvocationHandler;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.lang.reflect.Proxy;
+import java.util.concurrent.atomic.AtomicReference;
+
+import org.reactivestreams.Publisher;
+
+import org.springframework.beans.factory.DisposableBean;
+import org.springframework.lang.Nullable;
+import org.springframework.util.Assert;
+
+/**
+ * Implementation of {@link SmartConnectionFactory} that wraps a single R2DBC Connection which is not closed after use.
+ * Obviously, this is not multi-threading capable.
+ * <p>
+ * Note that at shutdown, someone should close the underlying Connection via the {@code close()} method. Client code
+ * will never call close on the Connection handle if it is SmartDataSource-aware (e.g. uses
+ * {@link ConnectionFactoryUtils#releaseConnection(io.r2dbc.spi.Connection, ConnectionFactory)}).
+ * <p>
+ * If client code will call {@link #close()} in the assumption of a pooled Connection, like when using persistence
+ * tools, set "suppressClose" to "true". This will return a close-suppressing proxy instead of the physical Connection.
+ * <p>
+ * This is primarily intended for testing. For example, it enables easy testing outside an application server, for code
+ * that expects to work on a {@link ConnectionFactory}.
+ *
+ * @author Mark Paluch
+ * @see #create()
+ * @see io.r2dbc.spi.Connection#close()
+ * @see ConnectionFactoryUtils#releaseConnection(io.r2dbc.spi.Connection, ConnectionFactory)
+ */
+public class SingleConnectionConnectionFactory extends DelegatingConnectionFactory
+		implements SmartConnectionFactory, DisposableBean {
+
+	/** Create a close-suppressing proxy?. */
+	private boolean suppressClose;
+
+	/** Override auto-commit state?. */
+	private @Nullable Boolean autoCommit;
+
+	/** Wrapped Connection. */
+	private final AtomicReference<Connection> target = new AtomicReference<>();
+
+	/** Proxy Connection. */
+	private @Nullable Connection connection;
+
+	private final Mono<? extends Connection> connectionEmitter;
+
+	/**
+	 * Constructor for bean-style configuration.
+	 */
+	public SingleConnectionConnectionFactory(ConnectionFactory targetConnectionFactory) {
+		super(targetConnectionFactory);
+		this.connectionEmitter = super.create().cache();
+	}
+
+	/**
+	 * Create a new {@link SingleConnectionConnectionFactory} using a R2DBC connection URL.
+	 *
+	 * @param url the R2DBC URL to use for accessing {@link ConnectionFactory} discovery.
+	 * @param suppressClose if the returned {@link Connection} should be a close-suppressing proxy or the physical
+	 *          {@link Connection}.
+	 * @see ConnectionFactories#get(String)
+	 */
+	public SingleConnectionConnectionFactory(String url, boolean suppressClose) {
+		super(ConnectionFactories.get(url));
+		this.suppressClose = suppressClose;
+		this.connectionEmitter = super.create().cache();
+	}
+
+	/**
+	 * Create a new {@link SingleConnectionConnectionFactory} with a given {@link Connection} and
+	 * {@link ConnectionFactoryMetadata}.
+	 *
+	 * @param target underlying target {@link Connection}.
+	 * @param metadata {@link ConnectionFactory} metadata to be associated with this {@link ConnectionFactory}.
+	 * @param suppressClose if the {@link Connection} should be wrapped with a {@link Connection} that suppresses
+	 *          {@code close()} calls (to allow for normal {@link #close()} usage in applications that expect a pooled
+	 *          {@link Connection} but do not know our {@link SmartConnectionFactory} interface).
+	 */
+	public SingleConnectionConnectionFactory(Connection target, ConnectionFactoryMetadata metadata,
+			boolean suppressClose) {
+		super(new ConnectionFactory() {
+			@Override
+			public Publisher<? extends Connection> create() {
+				return Mono.just(target);
+			}
+
+			@Override
+			public ConnectionFactoryMetadata getMetadata() {
+				return metadata;
+			}
+		});
+		Assert.notNull(target, "Connection must not be null");
+		Assert.notNull(metadata, "ConnectionFactoryMetadata must not be null");
+		this.target.set(target);
+		this.connectionEmitter = Mono.just(target);
+		this.suppressClose = suppressClose;
+		this.connection = (suppressClose ? getCloseSuppressingConnectionProxy(target) : target);
+	}
+
+	/**
+	 * Set whether the returned {@link Connection} should be a close-suppressing proxy or the physical {@link Connection}.
+	 */
+	public void setSuppressClose(boolean suppressClose) {
+		this.suppressClose = suppressClose;
+	}
+
+	/**
+	 * Return whether the returned {@link Connection} will be a close-suppressing proxy or the physical
+	 * {@link Connection}.
+	 */
+	protected boolean isSuppressClose() {
+		return this.suppressClose;
+	}
+
+	/**
+	 * Set whether the returned {@link Connection}'s "autoCommit" setting should be overridden.
+	 */
+	public void setAutoCommit(boolean autoCommit) {
+		this.autoCommit = autoCommit;
+	}
+
+	/**
+	 * Return whether the returned {@link Connection}'s "autoCommit" setting should be overridden.
+	 *
+	 * @return the "autoCommit" value, or {@code null} if none to be applied
+	 */
+	@Nullable
+	protected Boolean getAutoCommitValue() {
+		return this.autoCommit;
+	}
+
+	@Override
+	public Mono<? extends Connection> create() {
+
+		Connection connection = this.target.get();
+
+		return connectionEmitter.map(it -> {
+
+			if (connection == null) {
+				this.target.compareAndSet(connection, it);
+				this.connection = (isSuppressClose() ? getCloseSuppressingConnectionProxy(it) : it);
+			}
+
+			return this.connection;
+		}).flatMap(this::prepareConnection);
+	}
+
+	/**
+	 * This is a single Connection: Do not close it when returning to the "pool".
+	 */
+	@Override
+	public boolean shouldClose(Connection con) {
+		return (con != this.connection && con != this.target.get());
+	}
+
+	/**
+	 * Close the underlying {@link Connection}. The provider of this {@link ConnectionFactory} needs to care for proper
+	 * shutdown.
+	 * <p>
+	 * As this bean implements {@link DisposableBean}, a bean factory will automatically invoke this on destruction of its
+	 * cached singletons.
+	 */
+	@Override
+	public void destroy() {
+		resetConnection().block();
+	}
+
+	/**
+	 * Reset the underlying shared Connection, to be reinitialized on next access.
+	 */
+	public Mono<Void> resetConnection() {
+
+		Connection connection = this.target.get();
+
+		if (connection == null) {
+			return Mono.empty();
+		}
+
+		return Mono.defer(() -> {
+
+			if (this.target.compareAndSet(connection, null)) {
+
+				this.connection = null;
+
+				return Mono.from(connection.close());
+			}
+
+			return Mono.empty();
+		});
+	}
+
+	/**
+	 * Prepare the {@link Connection} before using it. Applies {@link #getAutoCommitValue() auto-commit} settings if
+	 * configured.
+	 *
+	 * @param connection the requested {@link Connection}.
+	 * @return the prepared {@link Connection}.
+	 */
+	protected Mono<Connection> prepareConnection(Connection connection) {
+
+		Boolean autoCommit = getAutoCommitValue();
+		if (autoCommit != null) {
+			return Mono.from(connection.setAutoCommit(autoCommit)).thenReturn(connection);
+		}
+
+		return Mono.just(connection);
+	}
+
+	/**
+	 * Wrap the given {@link Connection} with a proxy that delegates every method call to it but suppresses close calls.
+	 *
+	 * @param target the original {@link Connection} to wrap.
+	 * @return the wrapped Connection.
+	 */
+	protected Connection getCloseSuppressingConnectionProxy(Connection target) {
+		return (Connection) Proxy.newProxyInstance(ConnectionProxy.class.getClassLoader(),
+				new Class<?>[] { ConnectionProxy.class }, new CloseSuppressingInvocationHandler(target));
+	}
+
+	/**
+	 * Invocation handler that suppresses close calls on R2DBC Connections.
+	 *
+	 * @see io.r2dbc.spi.Connection#close()
+	 */
+	private static class CloseSuppressingInvocationHandler implements InvocationHandler {
+
+		private final io.r2dbc.spi.Connection target;
+
+		CloseSuppressingInvocationHandler(io.r2dbc.spi.Connection target) {
+			this.target = target;
+		}
+
+		@Override
+		@Nullable
+		public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
+			// Invocation on ConnectionProxy interface coming in...
+
+			if (method.getName().equals("equals")) {
+				// Only consider equal when proxies are identical.
+				return proxy == args[0];
+			} else if (method.getName().equals("hashCode")) {
+				// Use hashCode of PersistenceManager proxy.
+				return System.identityHashCode(proxy);
+			} else if (method.getName().equals("unwrap")) {
+				return target;
+			} else if (method.getName().equals("close")) {
+				// Handle close method: suppress, not valid.
+				return Mono.empty();
+			} else if (method.getName().equals("getTargetConnection")) {
+				// Handle getTargetConnection method: return underlying Connection.
+				return this.target;
+			}
+
+			// Invoke method on target Connection.
+			try {
+				Object retVal = method.invoke(this.target, args);
+
+				return retVal;
+			} catch (InvocationTargetException ex) {
+				throw ex.getTargetException();
+			}
+		}
+	}
+
+}

src/test/java/org/springframework/data/r2dbc/connectionfactory/SingleConnectionConnectionFactoryUnitTests.java

@@ -0,0 +1,138 @@
+/*
+ * 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.
+ * You may obtain a copy of the License at
+ *
+ *      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,
+ * 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.springframework.data.r2dbc.connectionfactory;
+
+import static org.assertj.core.api.Assertions.*;
+import static org.mockito.Mockito.*;
+
+import io.r2dbc.h2.H2Connection;
+import io.r2dbc.spi.Connection;
+import io.r2dbc.spi.ConnectionFactoryMetadata;
+import io.r2dbc.spi.IsolationLevel;
+import io.r2dbc.spi.R2dbcNonTransientResourceException;
+import io.r2dbc.spi.Wrapped;
+import reactor.core.publisher.Mono;
+import reactor.test.StepVerifier;
+
+import org.junit.Test;
+
+/**
+ * Unit tests for {@link SingleConnectionConnectionFactory}.
+ *
+ * @author Mark Paluch
+ */
+public class SingleConnectionConnectionFactoryUnitTests {
+
+	@Test // gh-204
+	public void shouldAllocateSameConnection() {
+
+		SingleConnectionConnectionFactory factory = new SingleConnectionConnectionFactory("r2dbc:h2:mem:///foo", false);
+
+		Mono<? extends Connection> cf1 = factory.create();
+		Mono<? extends Connection> cf2 = factory.create();
+
+		Connection c1 = cf1.block();
+		Connection c2 = cf2.block();
+
+		assertThat(c1).isSameAs(c2);
+		factory.destroy();
+	}
+
+	@Test // gh-204
+	public void shouldApplyAutoCommit() {
+
+		SingleConnectionConnectionFactory factory = new SingleConnectionConnectionFactory("r2dbc:h2:mem:///foo", false);
+		factory.setAutoCommit(false);
+
+		factory.create().as(StepVerifier::create).consumeNextWith(actual -> {
+			assertThat(actual.isAutoCommit()).isFalse();
+		}).verifyComplete();
+
+		factory.setAutoCommit(true);
+
+		factory.create().as(StepVerifier::create).consumeNextWith(actual -> {
+			assertThat(actual.isAutoCommit()).isTrue();
+		}).verifyComplete();
+
+		factory.destroy();
+	}
+
+	@Test // gh-204
+	public void shouldSuppressClose() {
+
+		SingleConnectionConnectionFactory factory = new SingleConnectionConnectionFactory("r2dbc:h2:mem:///foo", true);
+
+		Connection connection = factory.create().block();
+
+		StepVerifier.create(connection.close()).verifyComplete();
+		assertThat(connection).isInstanceOf(Wrapped.class);
+		assertThat(((Wrapped) connection).unwrap()).isInstanceOf(H2Connection.class);
+
+		StepVerifier.create(connection.setTransactionIsolationLevel(IsolationLevel.READ_COMMITTED)) //
+				.verifyComplete();
+		factory.destroy();
+	}
+
+	@Test // gh-204
+	public void shouldNotSuppressClose() {
+
+		SingleConnectionConnectionFactory factory = new SingleConnectionConnectionFactory("r2dbc:h2:mem:///foo", false);
+
+		Connection connection = factory.create().block();
+
+		StepVerifier.create(connection.close()).verifyComplete();
+
+		StepVerifier.create(connection.setTransactionIsolationLevel(IsolationLevel.READ_COMMITTED))
+				.verifyError(R2dbcNonTransientResourceException.class);
+		factory.destroy();
+	}
+
+	@Test // gh-204
+	public void releaseConnectionShouldNotCloseConnection() {
+
+		Connection connectionMock = mock(Connection.class);
+		ConnectionFactoryMetadata metadata = mock(ConnectionFactoryMetadata.class);
+
+		SingleConnectionConnectionFactory factory = new SingleConnectionConnectionFactory(connectionMock, metadata, false);
+
+		Connection connection = factory.create().block();
+
+		ConnectionFactoryUtils.releaseConnection(connection, factory) //
+				.as(StepVerifier::create) //
+				.verifyComplete();
+
+		verify(connectionMock, never()).close();
+	}
+
+	@Test // gh-204
+	public void releaseConnectionShouldCloseUnrelatedConnection() {
+
+		Connection connectionMock = mock(Connection.class);
+		Connection otherConnection = mock(Connection.class);
+		ConnectionFactoryMetadata metadata = mock(ConnectionFactoryMetadata.class);
+		when(otherConnection.close()).thenReturn(Mono.empty());
+
+		SingleConnectionConnectionFactory factory = new SingleConnectionConnectionFactory(connectionMock, metadata, false);
+
+		factory.create().as(StepVerifier::create).expectNextCount(1).verifyComplete();
+
+		ConnectionFactoryUtils.releaseConnection(otherConnection, factory) //
+				.as(StepVerifier::create) //
+				.verifyComplete();
+
+		verify(otherConnection).close();
+	}
+}