Use SpringSessionDataSource-annotated DataSource when one is available

Fixes gh-24624

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/session/JdbcSessionConfiguration.java

 /*
- * Copyright 2012-2020 the original author or authors.
+ * Copyright 2012-2021 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.
@@ -20,6 +20,7 @@ import java.time.Duration;
 
 import javax.sql.DataSource;
 
+import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
@@ -32,6 +33,7 @@ import org.springframework.core.io.ResourceLoader;
 import org.springframework.jdbc.core.JdbcTemplate;
 import org.springframework.session.SessionRepository;
 import org.springframework.session.jdbc.JdbcIndexedSessionRepository;
+import org.springframework.session.jdbc.config.annotation.SpringSessionDataSource;
 import org.springframework.session.jdbc.config.annotation.web.http.JdbcHttpSessionConfiguration;
 
 /**
@@ -51,9 +53,11 @@ class JdbcSessionConfiguration {
 
 	@Bean
 	@ConditionalOnMissingBean
-	JdbcSessionDataSourceInitializer jdbcSessionDataSourceInitializer(DataSource dataSource,
-			ResourceLoader resourceLoader, JdbcSessionProperties properties) {
-		return new JdbcSessionDataSourceInitializer(dataSource, resourceLoader, properties);
+	JdbcSessionDataSourceInitializer jdbcSessionDataSourceInitializer(
+			@SpringSessionDataSource ObjectProvider<DataSource> sessionDataSource,
+			ObjectProvider<DataSource> dataSource, ResourceLoader resourceLoader, JdbcSessionProperties properties) {
+		return new JdbcSessionDataSourceInitializer(sessionDataSource.getIfAvailable(dataSource::getObject),
+				resourceLoader, properties);
 	}
 
 	@Configuration(proxyBeanMethods = false)

spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/session/SessionAutoConfigurationJdbcTests.java

 /*
- * Copyright 2012-2019 the original author or authors.
+ * Copyright 2012-2021 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.
@@ -16,6 +16,9 @@
 
 package org.springframework.boot.autoconfigure.session;
 
+import javax.sql.DataSource;
+
+import org.apache.commons.dbcp2.BasicDataSource;
 import org.junit.jupiter.api.Test;
 
 import org.springframework.boot.autoconfigure.AutoConfigurations;
@@ -28,6 +31,9 @@ import org.springframework.boot.test.context.FilteredClassLoader;
 import org.springframework.boot.test.context.assertj.AssertableWebApplicationContext;
 import org.springframework.boot.test.context.runner.WebApplicationContextRunner;
 import org.springframework.boot.web.servlet.FilterRegistrationBean;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.context.annotation.Primary;
 import org.springframework.jdbc.BadSqlGrammarException;
 import org.springframework.jdbc.core.JdbcOperations;
 import org.springframework.session.FlushMode;
@@ -36,6 +42,7 @@ import org.springframework.session.data.mongo.MongoIndexedSessionRepository;
 import org.springframework.session.data.redis.RedisIndexedSessionRepository;
 import org.springframework.session.hazelcast.HazelcastIndexedSessionRepository;
 import org.springframework.session.jdbc.JdbcIndexedSessionRepository;
+import org.springframework.session.jdbc.config.annotation.SpringSessionDataSource;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
@@ -104,6 +111,19 @@ class SessionAutoConfigurationJdbcTests extends AbstractSessionAutoConfiguration
 				});
 	}
 
+	@Test
+	void sessionDataSourceIsUsedWhenAvailable() {
+		this.contextRunner.withUserConfiguration(SessionDataSourceConfiguration.class)
+				.withPropertyValues("spring.session.store-type=jdbc").run((context) -> {
+					JdbcIndexedSessionRepository repository = validateSessionRepository(context,
+							JdbcIndexedSessionRepository.class);
+					assertThat(repository).extracting("jdbcOperations").extracting("dataSource")
+							.isEqualTo(context.getBean("sessionDataSource"));
+					assertThatExceptionOfType(BadSqlGrammarException.class).isThrownBy(
+							() -> context.getBean(JdbcOperations.class).queryForList("select * from SPRING_SESSION"));
+				});
+	}
+
 	@Test
 	void customTableName() {
 		this.contextRunner
@@ -157,4 +177,29 @@ class SessionAutoConfigurationJdbcTests extends AbstractSessionAutoConfiguration
 				});
 	}
 
+	@Configuration
+	static class SessionDataSourceConfiguration {
+
+		@Bean
+		@SpringSessionDataSource
+		DataSource sessionDataSource() {
+			BasicDataSource dataSource = new BasicDataSource();
+			dataSource.setDriverClassName("org.hsqldb.jdbcDriver");
+			dataSource.setUrl("jdbc:hsqldb:mem:sessiondb");
+			dataSource.setUsername("sa");
+			return dataSource;
+		}
+
+		@Bean
+		@Primary
+		DataSource mainDataSource() {
+			BasicDataSource dataSource = new BasicDataSource();
+			dataSource.setDriverClassName("org.hsqldb.jdbcDriver");
+			dataSource.setUrl("jdbc:hsqldb:mem:maindb");
+			dataSource.setUsername("sa");
+			return dataSource;
+		}
+
+	}
+
 }