#98 - Add support for AbstractRoutingConnectionFactory.

We now provide an abstract base class for ConnectionFactory routing. Routing keys are typically obtained from a subscriber context. AbstractRoutingConnectionFactory is backed by either a map of string identifiers or connection factories. When using string identifiers, these can map agains e.g. Spring bean names that can be resolved using BeanFactoryConnectionFactoryLookup.

class MyRoutingConnectionFactory extends AbstractRoutingConnectionFactory {

		@Override
		protected Mono<Object> determineCurrentLookupKey() {
			return Mono.subscriberContext().filter(it -> it.hasKey(ROUTING_KEY)).map(it -> it.get(ROUTING_KEY));
		}
}

@Bean
public void routingConnectionFactory() {

		MyRoutingConnectionFactory router = new MyRoutingConnectionFactory();

		Map<String, ConnectionFactory> factories = new HashMap<>();
		ConnectionFactory myDefault = …;
		ConnectionFactory primary = …;
		ConnectionFactory secondary = …;

		factories.put("primary", primary);
		factories.put("secondary", secondary);

		router.setTargetConnectionFactories(factories);
		router.setDefaultTargetConnectionFactory(myDefault);

		return router;
}

Original pull request: #132.

src/test/java/org/springframework/data/r2dbc/connectionfactory/lookup/AbstractRoutingConnectionFactoryUnitTests.java

@@ -0,0 +1,191 @@
+/*
+ * 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.lookup;
+
+import static org.assertj.core.api.Assertions.*;
+
+import io.r2dbc.spi.ConnectionFactory;
+import reactor.core.publisher.Mono;
+import reactor.test.StepVerifier;
+import reactor.util.context.Context;
+
+import java.util.Collections;
+import java.util.Map;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnitRunner;
+
+/**
+ * Unit tests for {@link AbstractRoutingConnectionFactory}.
+ *
+ * @author Mark Paluch
+ */
+@RunWith(MockitoJUnitRunner.class)
+public class AbstractRoutingConnectionFactoryUnitTests {
+
+	private static final String ROUTING_KEY = "routingKey";
+
+	@Mock ConnectionFactory defaultConnectionFactory;
+	@Mock ConnectionFactory routedConnectionFactory;
+
+	DummyRoutingConnectionFactory sut;
+
+	@Before
+	public void before() {
+
+		sut = new DummyRoutingConnectionFactory();
+		sut.setDefaultTargetConnectionFactory(defaultConnectionFactory);
+	}
+
+	@Test // gh-98
+	public void shouldDetermineRoutedFactory() {
+
+		sut.setTargetConnectionFactories(Collections.singletonMap("key", routedConnectionFactory));
+		sut.setConnectionFactoryLookup(new MapConnectionFactoryLookup());
+		sut.afterPropertiesSet();
+
+		sut.determineTargetConnectionFactory() //
+				.subscriberContext(Context.of(ROUTING_KEY, "key")) //
+				.as(StepVerifier::create) //
+				.expectNext(routedConnectionFactory) //
+				.verifyComplete();
+	}
+
+	@Test // gh-98
+	public void shouldFallbackToDefaultConnectionFactory() {
+
+		sut.setTargetConnectionFactories(Collections.singletonMap("key", routedConnectionFactory));
+		sut.afterPropertiesSet();
+
+		sut.determineTargetConnectionFactory() //
+				.as(StepVerifier::create) //
+				.expectNext(defaultConnectionFactory) //
+				.verifyComplete();
+	}
+
+	@Test // gh-98
+	public void initializationShouldFailUnsupportedLookupKey() {
+
+		sut.setTargetConnectionFactories(Collections.singletonMap("key", new Object()));
+
+		assertThatThrownBy(() -> sut.afterPropertiesSet()).isInstanceOf(IllegalArgumentException.class);
+	}
+
+	@Test // gh-98
+	public void initializationShouldFailUnresolvableKey() {
+
+		sut.setTargetConnectionFactories(Collections.singletonMap("key", "value"));
+		sut.setConnectionFactoryLookup(new MapConnectionFactoryLookup());
+
+		assertThatThrownBy(() -> sut.afterPropertiesSet()).isInstanceOf(ConnectionFactoryLookupFailureException.class)
+				.hasMessageContaining("No ConnectionFactory with name 'value' registered");
+	}
+
+	@Test // gh-98
+	public void unresolvableConnectionFactoryRetrievalShouldFail() {
+
+		sut.setLenientFallback(false);
+		sut.setConnectionFactoryLookup(new MapConnectionFactoryLookup());
+		sut.setTargetConnectionFactories(Collections.singletonMap("key", routedConnectionFactory));
+		sut.afterPropertiesSet();
+
+		sut.determineTargetConnectionFactory() //
+				.subscriberContext(Context.of(ROUTING_KEY, "unknown")) //
+				.as(StepVerifier::create) //
+				.verifyError(IllegalStateException.class);
+	}
+
+	@Test // gh-98
+	public void connectionFactoryRetrievalWithUnknownLookupKeyShouldReturnDefaultConnectionFactory() {
+
+		sut.setTargetConnectionFactories(Collections.singletonMap("key", routedConnectionFactory));
+		sut.setDefaultTargetConnectionFactory(defaultConnectionFactory);
+		sut.afterPropertiesSet();
+
+		sut.determineTargetConnectionFactory() //
+				.subscriberContext(Context.of(ROUTING_KEY, "unknown")) //
+				.as(StepVerifier::create) //
+				.expectNext(defaultConnectionFactory) //
+				.verifyComplete();
+	}
+
+	@Test // gh-98
+	public void connectionFactoryRetrievalWithoutLookupKeyShouldReturnDefaultConnectionFactory() {
+
+		sut.setTargetConnectionFactories(Collections.singletonMap("key", routedConnectionFactory));
+		sut.setDefaultTargetConnectionFactory(defaultConnectionFactory);
+		sut.setLenientFallback(false);
+		sut.afterPropertiesSet();
+
+		sut.determineTargetConnectionFactory() //
+				.as(StepVerifier::create) //
+				.expectNext(defaultConnectionFactory) //
+				.verifyComplete();
+	}
+
+	@Test // gh-98
+	public void shouldLookupFromMap() {
+
+		MapConnectionFactoryLookup lookup = new MapConnectionFactoryLookup("lookup-key", routedConnectionFactory);
+
+		sut.setConnectionFactoryLookup(lookup);
+		sut.setTargetConnectionFactories(Collections.singletonMap("my-key", "lookup-key"));
+		sut.afterPropertiesSet();
+
+		sut.determineTargetConnectionFactory() //
+				.subscriberContext(Context.of(ROUTING_KEY, "my-key")) //
+				.as(StepVerifier::create) //
+				.expectNext(routedConnectionFactory) //
+				.verifyComplete();
+	}
+
+	@Test // gh-98
+	@SuppressWarnings("unchecked")
+	public void shouldAllowModificationsAfterInitialization() {
+
+		MapConnectionFactoryLookup lookup = new MapConnectionFactoryLookup();
+
+		sut.setConnectionFactoryLookup(lookup);
+		sut.setTargetConnectionFactories((Map) lookup.getConnectionFactories());
+		sut.afterPropertiesSet();
+
+		sut.determineTargetConnectionFactory() //
+				.subscriberContext(Context.of(ROUTING_KEY, "lookup-key")) //
+				.as(StepVerifier::create) //
+				.expectNext(defaultConnectionFactory) //
+				.verifyComplete();
+
+		lookup.addConnectionFactory("lookup-key", routedConnectionFactory);
+		sut.afterPropertiesSet();
+
+		sut.determineTargetConnectionFactory() //
+				.subscriberContext(Context.of(ROUTING_KEY, "lookup-key")) //
+				.as(StepVerifier::create) //
+				.expectNext(routedConnectionFactory) //
+				.verifyComplete();
+	}
+
+	static class DummyRoutingConnectionFactory extends AbstractRoutingConnectionFactory {
+
+		@Override
+		protected Mono<Object> determineCurrentLookupKey() {
+			return Mono.subscriberContext().filter(it -> it.hasKey(ROUTING_KEY)).map(it -> it.get(ROUTING_KEY));
+		}
+	}
+}

src/test/java/org/springframework/data/r2dbc/connectionfactory/lookup/BeanFactoryConnectionFactoryLookupUnitTests.java

@@ -0,0 +1,80 @@
+/*
+ * 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.lookup;
+
+import static org.assertj.core.api.Assertions.*;
+import static org.mockito.Mockito.*;
+
+import io.r2dbc.spi.ConnectionFactory;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnitRunner;
+
+import org.springframework.beans.factory.BeanFactory;
+import org.springframework.beans.factory.BeanNotOfRequiredTypeException;
+
+/**
+ * Unit tests for {@link BeanFactoryConnectionFactoryLookup}.
+ *
+ * @author Mark Paluch
+ */
+@RunWith(MockitoJUnitRunner.class)
+public class BeanFactoryConnectionFactoryLookupUnitTests {
+
+	private static final String CONNECTION_FACTORY_BEAN_NAME = "connectionFactory";
+
+	@Mock BeanFactory beanFactory;
+
+	@Test // gh-98
+	public void shouldLookupConnectionFactory() {
+
+		DummyConnectionFactory expectedConnectionFactory = new DummyConnectionFactory();
+		when(beanFactory.getBean(CONNECTION_FACTORY_BEAN_NAME, ConnectionFactory.class))
+				.thenReturn(expectedConnectionFactory);
+
+		BeanFactoryConnectionFactoryLookup lookup = new BeanFactoryConnectionFactoryLookup();
+		lookup.setBeanFactory(beanFactory);
+		ConnectionFactory connectionFactory = lookup.getConnectionFactory(CONNECTION_FACTORY_BEAN_NAME);
+
+		assertThat(connectionFactory).isNotNull();
+		assertThat(connectionFactory).isSameAs(expectedConnectionFactory);
+	}
+
+	@Test // gh-98
+	public void shouldLookupWhereBeanFactoryYieldsNonConnectionFactoryType() {
+
+		BeanFactory beanFactory = mock(BeanFactory.class);
+
+		when(beanFactory.getBean(CONNECTION_FACTORY_BEAN_NAME, ConnectionFactory.class)).thenThrow(
+				new BeanNotOfRequiredTypeException(CONNECTION_FACTORY_BEAN_NAME, ConnectionFactory.class, String.class));
+
+		BeanFactoryConnectionFactoryLookup lookup = new BeanFactoryConnectionFactoryLookup(beanFactory);
+
+		assertThatExceptionOfType(ConnectionFactoryLookupFailureException.class)
+				.isThrownBy(() -> lookup.getConnectionFactory(CONNECTION_FACTORY_BEAN_NAME));
+	}
+
+	@Test // gh-98
+	public void shouldLookupWhereBeanFactoryHasNotBeenSupplied() {
+
+		BeanFactoryConnectionFactoryLookup lookup = new BeanFactoryConnectionFactoryLookup();
+
+		assertThatThrownBy(() -> lookup.getConnectionFactory(CONNECTION_FACTORY_BEAN_NAME))
+				.isInstanceOf(IllegalStateException.class);
+	}
+}

src/test/java/org/springframework/data/r2dbc/connectionfactory/lookup/DummyConnectionFactory.java

@@ -0,0 +1,42 @@
+/*
+ * 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.lookup;
+
+import io.r2dbc.spi.Connection;
+import io.r2dbc.spi.ConnectionFactory;
+import io.r2dbc.spi.ConnectionFactoryMetadata;
+
+import org.reactivestreams.Publisher;
+
+/**
+ * Stub, do-nothing {@link ConnectionFactory} implementation.
+ * <p>
+ * All methods throw {@link UnsupportedOperationException}.
+ *
+ * @author Mark Paluch
+ */
+class DummyConnectionFactory implements ConnectionFactory {
+
+	@Override
+	public Publisher<? extends Connection> create() {
+		throw new UnsupportedOperationException();
+	}
+
+	@Override
+	public ConnectionFactoryMetadata getMetadata() {
+		throw new UnsupportedOperationException();
+	}
+}

src/test/java/org/springframework/data/r2dbc/connectionfactory/lookup/MapConnectionFactoryLookupUnitTests.java

@@ -0,0 +1,102 @@
+/*
+ * 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.lookup;
+
+import static org.assertj.core.api.Assertions.*;
+
+import io.r2dbc.spi.ConnectionFactory;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import org.junit.Test;
+
+/**
+ * Unit tests for {@link MapConnectionFactoryLookup}.
+ *
+ * @author Mark Paluch
+ */
+public class MapConnectionFactoryLookupUnitTests {
+
+	private static final String CONNECTION_FACTORY_NAME = "connectionFactory";
+
+	@Test // gh-98
+	public void getConnectionFactorysReturnsUnmodifiableMap() {
+
+		MapConnectionFactoryLookup lookup = new MapConnectionFactoryLookup();
+		Map<String, ConnectionFactory> connectionFactories = lookup.getConnectionFactories();
+
+		assertThatThrownBy(() -> connectionFactories.put("", new DummyConnectionFactory()))
+				.isInstanceOf(UnsupportedOperationException.class);
+	}
+
+	@Test // gh-98
+	public void shouldLookupConnectionFactory() {
+
+		Map<String, ConnectionFactory> connectionFactories = new HashMap<>();
+		DummyConnectionFactory expectedConnectionFactory = new DummyConnectionFactory();
+
+		connectionFactories.put(CONNECTION_FACTORY_NAME, expectedConnectionFactory);
+		MapConnectionFactoryLookup lookup = new MapConnectionFactoryLookup();
+
+		lookup.setConnectionFactories(connectionFactories);
+
+		ConnectionFactory connectionFactory = lookup.getConnectionFactory(CONNECTION_FACTORY_NAME);
+
+		assertThat(connectionFactory).isNotNull();
+		assertThat(connectionFactory).isSameAs(expectedConnectionFactory);
+	}
+
+	@Test // gh-98
+	public void addingConnectionFactoryPermitsOverride() {
+
+		Map<String, ConnectionFactory> connectionFactories = new HashMap<>();
+		DummyConnectionFactory overriddenConnectionFactory = new DummyConnectionFactory();
+		DummyConnectionFactory expectedConnectionFactory = new DummyConnectionFactory();
+		connectionFactories.put(CONNECTION_FACTORY_NAME, overriddenConnectionFactory);
+
+		MapConnectionFactoryLookup lookup = new MapConnectionFactoryLookup();
+
+		lookup.setConnectionFactories(connectionFactories);
+		lookup.addConnectionFactory(CONNECTION_FACTORY_NAME, expectedConnectionFactory);
+
+		ConnectionFactory connectionFactory = lookup.getConnectionFactory(CONNECTION_FACTORY_NAME);
+
+		assertThat(connectionFactory).isNotNull();
+		assertThat(connectionFactory).isSameAs(expectedConnectionFactory);
+	}
+
+	@Test // gh-98
+	@SuppressWarnings("unchecked")
+	public void getConnectionFactoryWhereSuppliedMapHasNonConnectionFactoryTypeUnderSpecifiedKey() {
+
+		Map connectionFactories = new HashMap<>();
+		connectionFactories.put(CONNECTION_FACTORY_NAME, new Object());
+		MapConnectionFactoryLookup lookup = new MapConnectionFactoryLookup(connectionFactories);
+
+		assertThatThrownBy(() -> lookup.getConnectionFactory(CONNECTION_FACTORY_NAME))
+				.isInstanceOf(ClassCastException.class);
+	}
+
+	@Test // gh-98
+	public void getConnectionFactoryWhereSuppliedMapHasNoEntryForSpecifiedKey() {
+
+		MapConnectionFactoryLookup lookup = new MapConnectionFactoryLookup();
+
+		assertThatThrownBy(() -> lookup.getConnectionFactory(CONNECTION_FACTORY_NAME))
+				.isInstanceOf(ConnectionFactoryLookupFailureException.class);
+	}
+}