Spring session in a multidatasource environment

src/main/java/org/springframework/session/data/mongo/config/annotation/web/http/MongoHttpSessionConfiguration.java

@@ -18,6 +18,7 @@
 import java.time.Duration;
 
 import org.springframework.beans.factory.BeanClassLoaderAware;
+import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.context.EmbeddedValueResolverAware;
 import org.springframework.context.annotation.Bean;
@@ -52,11 +53,23 @@ public class MongoHttpSessionConfiguration extends SpringHttpSessionConfiguratio
 	private String collectionName;
 	private StringValueResolver embeddedValueResolver;
 	private ClassLoader classLoader;
+	private MongoOperations mongoOperations;
+
+	@Autowired
+	public void setMongoOperations(
+			@SpringSessionMongoOperations ObjectProvider<MongoOperations> springSessionMongoOperations,
+			ObjectProvider<MongoOperations> mongoOperations) {
+		MongoOperations mongoOperationsToUse = springSessionMongoOperations.getIfAvailable();
+		if (mongoOperationsToUse == null) {
+			mongoOperationsToUse = mongoOperations.getObject();
+		}
+		this.mongoOperations = mongoOperationsToUse;
+	}
 
 	@Bean
-	public MongoOperationsSessionRepository mongoSessionRepository(MongoOperations mongoOperations) {
+	public MongoOperationsSessionRepository mongoSessionRepository() {
 
-		MongoOperationsSessionRepository repository = new MongoOperationsSessionRepository(mongoOperations);
+		MongoOperationsSessionRepository repository = new MongoOperationsSessionRepository(this.mongoOperations);
 		repository.setMaxInactiveIntervalInSeconds(this.maxInactiveIntervalInSeconds);
 
 		if (this.mongoSessionConverter != null) {

src/main/java/org/springframework/session/data/mongo/config/annotation/web/http/SpringSessionMongoOperations.java

@@ -0,0 +1,43 @@
+/*
+ * Copyright 2014-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.session.data.mongo.config.annotation.web.http;
+
+import java.lang.annotation.Documented;
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+import org.springframework.beans.factory.annotation.Qualifier;
+
+
+/**
+ * Qualifier annotation for a {@link org.springframework.data.mongodb.core.MongoOperations} to be injected in
+ * {@link org.springframework.session.data.mongo.MongoOperationsSessionRepository}.
+ * <p>
+ * This will enable us to have multiple MongoOperations in the application.
+ *
+ * @author Visweshwar Ganesh
+ * @since 2.2.0
+ */
+@Target({ ElementType.FIELD, ElementType.METHOD, ElementType.PARAMETER, ElementType.TYPE, ElementType.ANNOTATION_TYPE })
+@Retention(RetentionPolicy.RUNTIME)
+@Documented
+@Qualifier
+public @interface SpringSessionMongoOperations {
+
+}

src/test/java/org/springframework/session/data/mongo/config/annotation/web/http/MongoHttpSessionConfigurationTest.java

@@ -19,17 +19,17 @@
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.BDDMockito.*;
 
-import java.net.UnknownHostException;
 
 import org.junit.After;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
-import org.springframework.beans.factory.UnsatisfiedDependencyException;
+import org.springframework.beans.factory.BeanCreationException;
 import org.springframework.context.annotation.AnnotationConfigApplicationContext;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.context.annotation.Import;
+import org.springframework.context.annotation.Primary;
 import org.springframework.context.support.PropertySourcesPlaceholderConfigurer;
 import org.springframework.data.mongodb.core.MongoOperations;
 import org.springframework.data.mongodb.core.index.IndexOperations;
@@ -65,10 +65,9 @@ public void after() {
 	@Test
 	public void noMongoOperationsConfiguration() {
 
-		this.thrown.expect(UnsatisfiedDependencyException.class);
-		this.thrown.expectMessage("mongoSessionRepository");
-
-		registerAndRefresh(EmptyConfiguration.class);
+		assertThatExceptionOfType(BeanCreationException.class).isThrownBy(
+				() -> registerAndRefresh(NoMongoOperationsConfiguration.class)).withMessageContaining(
+				"expected at least 1 bean which qualifies as autowire candidate");
 	}
 
 	@Test
@@ -155,16 +154,82 @@ private void registerAndRefresh(Class<?>... annotatedClasses) {
 		this.context.refresh();
 	}
 
+	@Test
+	public void multipleDataSourceConfiguration() {
+		assertThatExceptionOfType(BeanCreationException.class).isThrownBy(
+				() -> registerAndRefresh(MongoOperationConfiguration.class,
+						MultipleMongoOperationsConfiguration.class)).withMessageContaining(
+				"expected single matching bean but found 2");
+	}
+
+	@Test
+	public void primaryMongoOperationConfiguration() {
+
+		registerAndRefresh(MongoOperationConfiguration.class, PrimaryMongoOperationsConfiguration.class);
+
+		MongoOperationsSessionRepository repository = this.context.getBean(MongoOperationsSessionRepository.class);
+		MongoOperations mongoOperations = this.context.getBean("primaryMongoOperations", MongoOperations.class);
+		assertThat(repository).isNotNull();
+		assertThat(mongoOperations).isNotNull();
+		MongoOperations mongoOperationsReflection = (MongoOperations) ReflectionTestUtils.getField(repository,
+				"mongoOperations");
+		assertThat(mongoOperationsReflection).isNotNull();
+		assertThat((mongoOperationsReflection)).isEqualTo(mongoOperations);
+	}
+
+	@Test
+	public void qualifiedDataSourceConfiguration() {
+		registerAndRefresh(MongoOperationConfiguration.class, QualifiedMongoOperationsConfiguration.class);
+
+		MongoOperationsSessionRepository repository = this.context.getBean(MongoOperationsSessionRepository.class);
+		MongoOperations mongoOperations = this.context.getBean("qualifiedMongoOperations", MongoOperations.class);
+		assertThat(repository).isNotNull();
+		assertThat(mongoOperations).isNotNull();
+		MongoOperations mongoOperationsReflection = (MongoOperations) ReflectionTestUtils.getField(repository,
+				"mongoOperations");
+		assertThat(mongoOperationsReflection).isNotNull();
+		assertThat(mongoOperationsReflection).isEqualTo(mongoOperations);
+	}
+
+	@Test
+	public void qualifiedAndPrimaryDataSourceConfiguration() {
+		registerAndRefresh(MongoOperationConfiguration.class, QualifiedAndPrimaryMongoConfiguration.class);
+
+		MongoOperationsSessionRepository repository = this.context.getBean(MongoOperationsSessionRepository.class);
+		MongoOperations mongoOperations = this.context.getBean("qualifiedMongoOperations", MongoOperations.class);
+		assertThat(repository).isNotNull();
+		assertThat(mongoOperations).isNotNull();
+		MongoOperations mongoOperationsReflection = (MongoOperations) ReflectionTestUtils.getField(repository,
+				"mongoOperations");
+		assertThat(mongoOperations).isNotNull();
+		assertThat(mongoOperationsReflection).isEqualTo(mongoOperations);
+	}
+
 	@Configuration
 	@EnableMongoHttpSession
 	static class EmptyConfiguration {
 
 	}
 
+	@Configuration
+	static class MongoOperationConfiguration {
+
+		@Bean
+		public MongoOperations defaultMongoOperations() {
+			MongoOperations mongoOperations = mock(MongoOperations.class);
+			IndexOperations indexOperations = mock(IndexOperations.class);
+
+			given(mongoOperations.indexOps(anyString())).willReturn(indexOperations);
+
+			return mongoOperations;
+		}
+
+	}
+
 	static class BaseConfiguration {
 
 		@Bean
-		public MongoOperations mongoOperations() throws UnknownHostException {
+		public MongoOperations mongoOperations() {
 
 			MongoOperations mongoOperations = mock(MongoOperations.class);
 			IndexOperations indexOperations = mock(IndexOperations.class);
@@ -241,4 +306,78 @@ public PropertySourcesPlaceholderConfigurer propertySourcesPlaceholderConfigurer
 
 	}
 
+	@EnableMongoHttpSession
+	static class NoMongoOperationsConfiguration {
+
+	}
+
+	@EnableMongoHttpSession
+	static class MultipleMongoOperationsConfiguration {
+
+		@Bean
+		public MongoOperations secondaryDataSource() {
+			return mock(MongoOperations.class);
+		}
+
+	}
+
+	@EnableMongoHttpSession
+	static class PrimaryMongoOperationsConfiguration {
+
+		@Bean
+		@Primary
+		public MongoOperations primaryMongoOperations() {
+			MongoOperations mongoOperations = mock(MongoOperations.class);
+			IndexOperations indexOperations = mock(IndexOperations.class);
+
+			given(mongoOperations.indexOps(anyString())).willReturn(indexOperations);
+
+			return mongoOperations;
+		}
+
+	}
+
+	@EnableMongoHttpSession
+	static class QualifiedMongoOperationsConfiguration {
+
+		@Bean
+		@SpringSessionMongoOperations
+		public MongoOperations qualifiedMongoOperations() {
+			MongoOperations mongoOperations = mock(MongoOperations.class);
+			IndexOperations indexOperations = mock(IndexOperations.class);
+
+			given(mongoOperations.indexOps(anyString())).willReturn(indexOperations);
+
+			return mongoOperations;
+		}
+
+	}
+
+	@EnableMongoHttpSession
+	static class QualifiedAndPrimaryMongoConfiguration {
+
+		@Bean
+		@SpringSessionMongoOperations
+		public MongoOperations qualifiedMongoOperations() {
+			MongoOperations mongoOperations = mock(MongoOperations.class);
+			IndexOperations indexOperations = mock(IndexOperations.class);
+
+			given(mongoOperations.indexOps(anyString())).willReturn(indexOperations);
+
+			return mongoOperations;
+		}
+
+		@Bean
+		@Primary
+		public MongoOperations primaryMongoOperations() {
+			MongoOperations mongoOperations = mock(MongoOperations.class);
+			IndexOperations indexOperations = mock(IndexOperations.class);
+
+			given(mongoOperations.indexOps(anyString())).willReturn(indexOperations);
+
+			return mongoOperations;
+		}
+
+	}
+
 }