Merge pull request #37630 from OlgaMaciaszek

* gh-37630:
  Polish "Instrument user-created DataSource for checkpoint-restore"
  Instrument user-created DataSource for checkpoint-restore

Closes gh-37630

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

@@ -51,13 +51,14 @@ import org.springframework.util.StringUtils;
  * @author Phillip Webb
  * @author Stephane Nicoll
  * @author Kazuki Shimizu
+ * @author Olga Maciaszek-Sharma
  * @since 1.0.0
  */
 @AutoConfiguration(before = SqlInitializationAutoConfiguration.class)
 @ConditionalOnClass({ DataSource.class, EmbeddedDatabaseType.class })
 @ConditionalOnMissingBean(type = "io.r2dbc.spi.ConnectionFactory")
 @EnableConfigurationProperties(DataSourceProperties.class)
-@Import(DataSourcePoolMetadataProvidersConfiguration.class)
+@Import({ DataSourcePoolMetadataProvidersConfiguration.class, DataSourceCheckpointRestoreConfiguration.class })
 public class DataSourceAutoConfiguration {
 
 	@Configuration(proxyBeanMethods = false)

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

@@ -0,0 +1,53 @@
+/*
+ * Copyright 2012-2023 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.boot.autoconfigure.jdbc;
+
+import javax.sql.DataSource;
+
+import com.zaxxer.hikari.HikariDataSource;
+
+import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnCheckpointRestore;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
+import org.springframework.boot.jdbc.HikariCheckpointRestoreLifecycle;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+
+/**
+ * Checkpoint-restore specific configuration.
+ *
+ * @author Olga Maciaszek-Sharma
+ */
+@Configuration(proxyBeanMethods = false)
+@ConditionalOnCheckpointRestore
+@ConditionalOnBean(DataSource.class)
+class DataSourceCheckpointRestoreConfiguration {
+
+	@Configuration(proxyBeanMethods = false)
+	@ConditionalOnClass(HikariDataSource.class)
+	static class Hikari {
+
+		@Bean
+		@ConditionalOnMissingBean
+		HikariCheckpointRestoreLifecycle hikariCheckpointRestoreLifecycle(DataSource dataSource) {
+			return new HikariCheckpointRestoreLifecycle(dataSource);
+		}
+
+	}
+
+}

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

@@ -25,14 +25,12 @@ import oracle.jdbc.OracleConnection;
 import oracle.ucp.jdbc.PoolDataSourceImpl;
 
 import org.springframework.beans.factory.ObjectProvider;
-import org.springframework.boot.autoconfigure.condition.ConditionalOnCheckpointRestore;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
 import org.springframework.boot.context.properties.ConfigurationProperties;
 import org.springframework.boot.jdbc.DataSourceBuilder;
 import org.springframework.boot.jdbc.DatabaseDriver;
-import org.springframework.boot.jdbc.HikariCheckpointRestoreLifecycle;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.util.StringUtils;
@@ -124,12 +122,6 @@ abstract class DataSourceConfiguration {
 			return dataSource;
 		}
 
-		@Bean
-		@ConditionalOnCheckpointRestore
-		HikariCheckpointRestoreLifecycle hikariCheckpointRestoreLifecycle(DataSource hikariDataSource) {
-			return new HikariCheckpointRestoreLifecycle(hikariDataSource);
-		}
-
 	}
 
 	/**

spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/jdbc/HikariDataSourceConfigurationTests.java

@@ -25,6 +25,7 @@ import org.junit.jupiter.api.Test;
 import org.springframework.beans.BeansException;
 import org.springframework.beans.factory.config.BeanPostProcessor;
 import org.springframework.boot.autoconfigure.AutoConfigurations;
+import org.springframework.boot.jdbc.DataSourceBuilder;
 import org.springframework.boot.jdbc.HikariCheckpointRestoreLifecycle;
 import org.springframework.boot.test.context.runner.ApplicationContextRunner;
 import org.springframework.boot.testsupport.classpath.ClassPathOverrides;
@@ -42,6 +43,7 @@ import static org.assertj.core.api.Assertions.assertThat;
  * @author Moritz Halbritter
  * @author Andy Wilkinson
  * @author Phillip Webb
+ * @author Olga Maciaszek-Sharma
  */
 class HikariDataSourceConfigurationTests {
 
@@ -138,16 +140,22 @@ class HikariDataSourceConfigurationTests {
 	@ClassPathOverrides("org.crac:crac:1.3.0")
 	void whenCheckpointRestoreIsAvailableAndDataSourceHasBeenWrappedHikariAutoConfigRegistersLifecycleBean() {
 		this.contextRunner.withUserConfiguration(DataSourceWrapperConfiguration.class)
-			.withPropertyValues("spring.datasource.type=" + HikariDataSource.class.getName())
 			.run((context) -> assertThat(context).hasSingleBean(HikariCheckpointRestoreLifecycle.class));
 	}
 
 	@Test
 	void whenCheckpointRestoreIsNotAvailableHikariAutoConfigDoesNotRegisterLifecycleBean() {
-		this.contextRunner.withPropertyValues("spring.datasource.type=" + HikariDataSource.class.getName())
+		this.contextRunner
 			.run((context) -> assertThat(context).doesNotHaveBean(HikariCheckpointRestoreLifecycle.class));
 	}
 
+	@Test
+	@ClassPathOverrides("org.crac:crac:1.3.0")
+	void whenCheckpointRestoreIsAvailableAndDataSourceIsFromUserConfigurationHikariAutoConfigRegistersLifecycleBean() {
+		this.contextRunner.withUserConfiguration(UserDataSourceConfiguration.class)
+			.run((context) -> assertThat(context).hasSingleBean(HikariCheckpointRestoreLifecycle.class));
+	}
+
 	@Configuration(proxyBeanMethods = false)
 	static class ConnectionDetailsConfiguration {
 
@@ -178,4 +186,19 @@ class HikariDataSourceConfigurationTests {
 
 	}
 
+	@Configuration(proxyBeanMethods = false)
+	static class UserDataSourceConfiguration {
+
+		@Bean
+		DataSource dataSource() {
+			return DataSourceBuilder.create()
+				.driverClassName("org.postgresql.Driver")
+				.url("jdbc:postgresql://localhost:5432/database")
+				.username("user")
+				.password("password")
+				.build();
+		}
+
+	}
+
 }