Add a generic to DataSourceBuilder

Closes gh-7942

Author: wilkinsona

spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/jdbc/DataSourceAutoConfiguration.java

@@ -131,8 +131,8 @@ public ConditionOutcome getMatchOutcome(ConditionContext context,
 		 * @return the class loader
 		 */
 		private ClassLoader getDataSourceClassLoader(ConditionContext context) {
-			Class<?> dataSourceClass = new DataSourceBuilder(context.getClassLoader())
-					.findType();
+			Class<?> dataSourceClass = DataSourceBuilder
+					.findType(context.getClassLoader());
 			return (dataSourceClass == null ? null : dataSourceClass.getClassLoader());
 		}
 

spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/jdbc/DataSourceProperties.java

@@ -180,7 +180,7 @@ public void afterPropertiesSet() throws Exception {
 	 * @return a {@link DataSourceBuilder} initialized with the customizations defined on
 	 * this instance
 	 */
-	public DataSourceBuilder initializeDataSourceBuilder() {
+	public DataSourceBuilder<?> initializeDataSourceBuilder() {
 		return DataSourceBuilder.create(getClassLoader()).type(getType())
 				.driverClassName(determineDriverClassName()).url(determineUrl())
 				.username(determineUsername()).password(determinePassword());

spring-boot-docs/src/main/java/org/springframework/boot/jdbc/ConfigurableDataSourceExample.java

@@ -51,8 +51,8 @@ public DataSourceProperties dataSourceProperties() {
 		@Bean
 		@ConfigurationProperties("app.datasource")
 		public HikariDataSource dataSource(DataSourceProperties properties) {
-			return (HikariDataSource) properties.initializeDataSourceBuilder()
-					.type(HikariDataSource.class).build();
+			return properties.initializeDataSourceBuilder().type(HikariDataSource.class)
+					.build();
 		}
 		// end::configuration[]
 

spring-boot-docs/src/main/java/org/springframework/boot/jdbc/SimpleDataSourceExample.java

@@ -41,8 +41,7 @@ static class SimpleDataSourceConfiguration {
 		@Bean
 		@ConfigurationProperties("app.datasource")
 		public HikariDataSource dataSource() {
-			return (HikariDataSource) DataSourceBuilder.create()
-					.type(HikariDataSource.class).build();
+			return DataSourceBuilder.create().type(HikariDataSource.class).build();
 		}
 		// end::configuration[]
 

spring-boot-docs/src/main/java/org/springframework/boot/jdbc/SimpleTwoDataSourcesExample.java

@@ -58,8 +58,7 @@ public DataSource fooDataSource() {
 		@Bean
 		@ConfigurationProperties("app.datasource.bar")
 		public BasicDataSource barDataSource() {
-			return (BasicDataSource) DataSourceBuilder.create()
-					.type(BasicDataSource.class).build();
+			return DataSourceBuilder.create().type(BasicDataSource.class).build();
 		}
 		// end::configuration[]
 

spring-boot/src/main/java/org/springframework/boot/jdbc/DataSourceBuilder.java

@@ -39,11 +39,12 @@
  * inject additional properties into the result you can downcast it, or use
  * {@code @ConfigurationProperties}.
  *
+ * @param <T> type of DataSource produced by the builder
  * @author Dave Syer
  * @author Madhura Bhave
  * @since 2.0.0
  */
-public class DataSourceBuilder {
+public final class DataSourceBuilder<T extends DataSource> {
 
 	private static final String[] DATA_SOURCE_TYPE_NAMES = new String[] {
 			"com.zaxxer.hikari.HikariDataSource",
@@ -56,24 +57,25 @@ public class DataSourceBuilder {
 
 	private Map<String, String> properties = new HashMap<>();
 
-	public static DataSourceBuilder create() {
-		return new DataSourceBuilder(null);
+	public static DataSourceBuilder<?> create() {
+		return new DataSourceBuilder<DataSource>(null);
 	}
 
-	public static DataSourceBuilder create(ClassLoader classLoader) {
-		return new DataSourceBuilder(classLoader);
+	public static DataSourceBuilder<?> create(ClassLoader classLoader) {
+		return new DataSourceBuilder<DataSource>(classLoader);
 	}
 
-	public DataSourceBuilder(ClassLoader classLoader) {
+	private DataSourceBuilder(ClassLoader classLoader) {
 		this.classLoader = classLoader;
 	}
 
-	public DataSource build() {
+	@SuppressWarnings("unchecked")
+	public T build() {
 		Class<? extends DataSource> type = getType();
 		DataSource result = BeanUtils.instantiateClass(type);
 		maybeGetDriverClassName();
 		bind(result);
-		return result;
+		return (T) result;
 	}
 
 	private void maybeGetDriverClassName() {
@@ -95,40 +97,38 @@ private void bind(DataSource result) {
 		binder.bind(ConfigurationPropertyName.EMPTY, Bindable.ofInstance(result));
 	}
 
-	public DataSourceBuilder type(Class<? extends DataSource> type) {
+	@SuppressWarnings("unchecked")
+	public <D extends DataSource> DataSourceBuilder<D> type(Class<D> type) {
 		this.type = type;
-		return this;
+		return (DataSourceBuilder<D>) this;
 	}
 
-	public DataSourceBuilder url(String url) {
+	public DataSourceBuilder<T> url(String url) {
 		this.properties.put("url", url);
 		return this;
 	}
 
-	public DataSourceBuilder driverClassName(String driverClassName) {
+	public DataSourceBuilder<T> driverClassName(String driverClassName) {
 		this.properties.put("driverClassName", driverClassName);
 		return this;
 	}
 
-	public DataSourceBuilder username(String username) {
+	public DataSourceBuilder<T> username(String username) {
 		this.properties.put("username", username);
 		return this;
 	}
 
-	public DataSourceBuilder password(String password) {
+	public DataSourceBuilder<T> password(String password) {
 		this.properties.put("password", password);
 		return this;
 	}
 
 	@SuppressWarnings("unchecked")
-	public Class<? extends DataSource> findType() {
-		if (this.type != null) {
-			return this.type;
-		}
+	public static Class<? extends DataSource> findType(ClassLoader classLoader) {
 		for (String name : DATA_SOURCE_TYPE_NAMES) {
 			try {
 				return (Class<? extends DataSource>) ClassUtils.forName(name,
-						this.classLoader);
+						classLoader);
 			}
 			catch (Exception ex) {
 				// Swallow and continue
@@ -138,7 +138,8 @@ public Class<? extends DataSource> findType() {
 	}
 
 	private Class<? extends DataSource> getType() {
-		Class<? extends DataSource> type = findType();
+		Class<? extends DataSource> type = this.type != null ? this.type
+				: findType(this.classLoader);
 		if (type != null) {
 			return type;
 		}

spring-boot/src/test/java/org/springframework/boot/jdbc/DataSourceBuilderTests.java

@@ -71,6 +71,13 @@ public void defaultToCommonsDbcp2AsLastResort() {
 		assertThat(this.dataSource).isInstanceOf(BasicDataSource.class);
 	}
 
+	@Test
+	public void specificTypeOfDataSource() {
+		HikariDataSource hikariDataSource = DataSourceBuilder.create().type(HikariDataSource.class)
+				.build();
+		assertThat(hikariDataSource).isInstanceOf(HikariDataSource.class);
+	}
+
 	final class HidePackagesClassLoader extends URLClassLoader {
 
 		private final String[] hiddenPackages;

spring-boot/src/test/java/org/springframework/boot/jdbc/metadata/AbstractDataSourcePoolMetadataTests.java

@@ -86,7 +86,7 @@ public void getPoolSizeTwoConnections() {
 	@Test
 	public abstract void getValidationQuery();
 
-	protected DataSourceBuilder initializeBuilder() {
+	protected DataSourceBuilder<?> initializeBuilder() {
 		return DataSourceBuilder.create().driverClassName("org.hsqldb.jdbc.JDBCDriver")
 				.url("jdbc:hsqldb:mem:test").username("sa");
 	}

spring-boot/src/test/java/org/springframework/boot/jdbc/metadata/CommonsDbcp2DataSourcePoolMetadataTests.java

@@ -90,7 +90,7 @@ private CommonsDbcp2DataSourcePoolMetadata createDataSourceMetadata(int minSize,
 	}
 
 	private BasicDataSource createDataSource() {
-		return (BasicDataSource) initializeBuilder().type(BasicDataSource.class).build();
+		return initializeBuilder().type(BasicDataSource.class).build();
 	}
 
 }

spring-boot/src/test/java/org/springframework/boot/jdbc/metadata/HikariDataSourcePoolMetadataTests.java

@@ -51,8 +51,8 @@ public void getValidationQuery() {
 	}
 
 	private HikariDataSource createDataSource(int minSize, int maxSize) {
-		HikariDataSource dataSource = (HikariDataSource) initializeBuilder()
-				.type(HikariDataSource.class).build();
+		HikariDataSource dataSource = initializeBuilder().type(HikariDataSource.class)
+				.build();
 		dataSource.setMinimumIdle(minSize);
 		dataSource.setMaximumPoolSize(maxSize);
 		return dataSource;

spring-boot/src/test/java/org/springframework/boot/jdbc/metadata/TomcatDataSourcePoolMetadataTests.java

@@ -51,8 +51,7 @@ public void getValidationQuery() {
 	}
 
 	private DataSource createDataSource(int minSize, int maxSize) {
-		DataSource dataSource = (DataSource) initializeBuilder().type(DataSource.class)
-				.build();
+		DataSource dataSource = initializeBuilder().type(DataSource.class).build();
 		dataSource.setMinIdle(minSize);
 		dataSource.setMaxActive(maxSize);