Spring session in a multidatasource environment

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

@@ -60,9 +60,21 @@ public class MongoHttpSessionConfiguration extends SpringHttpSessionConfiguratio
 	private List<SessionRepositoryCustomizer<MongoIndexedSessionRepository>> sessionRepositoryCustomizers;
 	private ClassLoader classLoader;
 	private IndexResolver<MongoSession> indexResolver;
+    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 MongoIndexedSessionRepository mongoSessionRepository(MongoOperations mongoOperations) {
+	public MongoIndexedSessionRepository mongoSessionRepository() {
 
 		MongoIndexedSessionRepository repository = new MongoIndexedSessionRepository(mongoOperations);
 		repository.setMaxInactiveIntervalInSeconds(this.maxInactiveIntervalInSeconds);
@@ -95,35 +107,35 @@ public MongoIndexedSessionRepository mongoSessionRepository(MongoOperations mong
 		return repository;
 	}
 
-	public void setCollectionName(String collectionName) {
-		this.collectionName = collectionName;
-	}
+    public void setCollectionName(String collectionName) {
+        this.collectionName = collectionName;
+    }
 
-	public void setMaxInactiveIntervalInSeconds(Integer maxInactiveIntervalInSeconds) {
-		this.maxInactiveIntervalInSeconds = maxInactiveIntervalInSeconds;
-	}
+    public void setMaxInactiveIntervalInSeconds(Integer maxInactiveIntervalInSeconds) {
+        this.maxInactiveIntervalInSeconds = maxInactiveIntervalInSeconds;
+    }
 
-	public void setImportMetadata(AnnotationMetadata importMetadata) {
+    public void setImportMetadata(AnnotationMetadata importMetadata) {
 
-		AnnotationAttributes attributes = AnnotationAttributes
-				.fromMap(importMetadata.getAnnotationAttributes(EnableMongoHttpSession.class.getName()));
+        AnnotationAttributes attributes = AnnotationAttributes
+                .fromMap(importMetadata.getAnnotationAttributes(EnableMongoHttpSession.class.getName()));
 
 		if (attributes != null) {
 			this.maxInactiveIntervalInSeconds = attributes.getNumber("maxInactiveIntervalInSeconds");
 		} else {
 			this.maxInactiveIntervalInSeconds = MongoIndexedSessionRepository.DEFAULT_INACTIVE_INTERVAL;
 		}
 
-		String collectionNameValue = attributes != null ? attributes.getString("collectionName") : "";
-		if (StringUtils.hasText(collectionNameValue)) {
-			this.collectionName = this.embeddedValueResolver.resolveStringValue(collectionNameValue);
-		}
-	}
+        String collectionNameValue = attributes != null ? attributes.getString("collectionName") : "";
+        if (StringUtils.hasText(collectionNameValue)) {
+            this.collectionName = this.embeddedValueResolver.resolveStringValue(collectionNameValue);
+        }
+    }
 
-	@Autowired(required = false)
-	public void setMongoSessionConverter(AbstractMongoSessionConverter mongoSessionConverter) {
-		this.mongoSessionConverter = mongoSessionConverter;
-	}
+    @Autowired(required = false)
+    public void setMongoSessionConverter(AbstractMongoSessionConverter mongoSessionConverter) {
+        this.mongoSessionConverter = mongoSessionConverter;
+    }
 
 	@Autowired(required = false)
 	public void setSessionRepositoryCustomizers(
@@ -136,10 +148,10 @@ public void setBeanClassLoader(ClassLoader classLoader) {
 		this.classLoader = classLoader;
 	}
 
-	@Override
-	public void setEmbeddedValueResolver(StringValueResolver resolver) {
-		this.embeddedValueResolver = resolver;
-	}
+    @Override
+    public void setEmbeddedValueResolver(StringValueResolver resolver) {
+        this.embeddedValueResolver = resolver;
+    }
 
 	@Autowired(required = false)
 	public void setIndexResolver(IndexResolver<MongoSession> indexResolver) {

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

@@ -0,0 +1,38 @@
+package org.springframework.session.data.mongo.config.annotation.web.http;
+
+import org.springframework.beans.factory.annotation.Qualifier;
+
+import java.lang.annotation.*;
+
+/*
+ * 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.
+ */
+/**
+ * Qualifier annotation for a {@link org.springframework.data.mongodb.core.MongoOperations} to be injected in
+ * {@link org.springframework.session.data.mongo.MongoIndexedSessionRepository}.
+ *
+ * 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

@@ -24,11 +24,9 @@
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Order;
 import org.junit.jupiter.api.Test;
+import org.springframework.beans.factory.BeanCreationException;
 import org.springframework.beans.factory.UnsatisfiedDependencyException;
-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.*;
 import org.springframework.context.support.PropertySourcesPlaceholderConfigurer;
 import org.springframework.data.mongodb.core.MongoOperations;
 import org.springframework.data.mongodb.core.index.IndexOperations;
@@ -66,9 +64,9 @@ public void after() {
 	@Test
 	public void noMongoOperationsConfiguration() {
 
-		assertThatExceptionOfType(UnsatisfiedDependencyException.class).isThrownBy(() -> {
+		assertThatExceptionOfType(BeanCreationException.class).isThrownBy(() -> {
 			registerAndRefresh(EmptyConfiguration.class);
-		}).withMessageContaining("mongoSessionRepository");
+		}).withMessageContaining("expected at least 1 bean which qualifies as autowire candidate");
 	}
 
 	@Test
@@ -166,7 +164,7 @@ void customIndexResolverConfigurationWithDefaultMongoSessionConverter() {
 
 		MongoIndexedSessionRepository repository = this.context.getBean(MongoIndexedSessionRepository.class);
 		IndexResolver<MongoSession> indexResolver = this.context.getBean(IndexResolver.class);
-		
+
 		assertThat(repository).isNotNull();
 		assertThat(indexResolver).isNotNull();
 		assertThat(repository).extracting("mongoSessionConverter").hasFieldOrPropertyWithValue("indexResolver", indexResolver);
@@ -191,13 +189,96 @@ private void registerAndRefresh(Class<?>... annotatedClasses) {
 		this.context.refresh();
 	}
 
-	@Configuration
-	@EnableMongoHttpSession
-	static class EmptyConfiguration {
+    @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);
+
+
+		MongoIndexedSessionRepository repository = this.context
+                .getBean(MongoIndexedSessionRepository.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);
+
+		MongoIndexedSessionRepository repository = this.context
+                .getBean(MongoIndexedSessionRepository.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);
+
+		MongoIndexedSessionRepository repository = this.context
+                .getBean(MongoIndexedSessionRepository.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 {
 
 	}
 
-	static class BaseConfiguration {
+    @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 {
@@ -277,6 +358,85 @@ 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;
+        }
+
+
+    }
+
 	@EnableMongoHttpSession
 	static class SessionRepositoryCustomizerConfiguration {