Keep Liquibase-specific DataSource open for use by LiquibaseEndpoint

Closes gh-13832

Author: wilkinsona

spring-boot-project/spring-boot-actuator-autoconfigure/src/main/java/org/springframework/boot/actuate/autoconfigure/liquibase/LiquibaseEndpointAutoConfiguration.java

@@ -18,13 +18,16 @@
 
 import liquibase.integration.spring.SpringLiquibase;
 
+import org.springframework.beans.BeansException;
+import org.springframework.beans.factory.config.BeanPostProcessor;
 import org.springframework.boot.actuate.autoconfigure.endpoint.condition.ConditionalOnEnabledEndpoint;
 import org.springframework.boot.actuate.liquibase.LiquibaseEndpoint;
 import org.springframework.boot.autoconfigure.AutoConfigureAfter;
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
+import org.springframework.boot.autoconfigure.liquibase.DataSourceClosingSpringLiquibase;
 import org.springframework.boot.autoconfigure.liquibase.LiquibaseAutoConfiguration;
 import org.springframework.context.ApplicationContext;
 import org.springframework.context.annotation.Bean;
@@ -49,4 +52,29 @@ public LiquibaseEndpoint liquibaseEndpoint(ApplicationContext context) {
 		return new LiquibaseEndpoint(context);
 	}
 
+	@Bean
+	@ConditionalOnBean(SpringLiquibase.class)
+	@ConditionalOnEnabledEndpoint(endpoint = LiquibaseEndpoint.class)
+	public static BeanPostProcessor preventDataSourceCloseBeanPostProcessor() {
+		return new BeanPostProcessor() {
+
+			@Override
+			public Object postProcessBeforeInitialization(Object bean, String beanName)
+					throws BeansException {
+				if (bean instanceof DataSourceClosingSpringLiquibase) {
+					((DataSourceClosingSpringLiquibase) bean)
+							.setCloseDataSourceOnceMigrated(false);
+				}
+				return bean;
+			}
+
+			@Override
+			public Object postProcessAfterInitialization(Object bean, String beanName)
+					throws BeansException {
+				return bean;
+			}
+
+		};
+	}
+
 }

spring-boot-project/spring-boot-actuator-autoconfigure/src/test/java/org/springframework/boot/actuate/autoconfigure/liquibase/LiquibaseEndpointAutoConfigurationTests.java

@@ -16,11 +16,13 @@
 
 package org.springframework.boot.actuate.autoconfigure.liquibase;
 
+import liquibase.exception.LiquibaseException;
 import liquibase.integration.spring.SpringLiquibase;
 import org.junit.Test;
 
 import org.springframework.boot.actuate.liquibase.LiquibaseEndpoint;
 import org.springframework.boot.autoconfigure.AutoConfigurations;
+import org.springframework.boot.autoconfigure.liquibase.DataSourceClosingSpringLiquibase;
 import org.springframework.boot.test.context.runner.ApplicationContextRunner;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
@@ -37,23 +39,48 @@ public class LiquibaseEndpointAutoConfigurationTests {
 
 	private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
 			.withConfiguration(
-					AutoConfigurations.of(LiquibaseEndpointAutoConfiguration.class))
-			.withUserConfiguration(LiquibaseConfiguration.class);
+					AutoConfigurations.of(LiquibaseEndpointAutoConfiguration.class));
 
 	@Test
 	public void runShouldHaveEndpointBean() {
-		this.contextRunner.run(
+		this.contextRunner.withUserConfiguration(LiquibaseConfiguration.class).run(
 				(context) -> assertThat(context).hasSingleBean(LiquibaseEndpoint.class));
 	}
 
 	@Test
 	public void runWhenEnabledPropertyIsFalseShouldNotHaveEndpointBean() {
-		this.contextRunner
+		this.contextRunner.withUserConfiguration(LiquibaseConfiguration.class)
 				.withPropertyValues("management.endpoint.liquibase.enabled:false")
 				.run((context) -> assertThat(context)
 						.doesNotHaveBean(LiquibaseEndpoint.class));
 	}
 
+	@Test
+	public void disablesCloseOfDataSourceWhenEndpointIsEnabled() {
+		this.contextRunner
+				.withUserConfiguration(DataSourceClosingLiquibaseConfiguration.class)
+				.run((context) -> {
+					assertThat(context).hasSingleBean(LiquibaseEndpoint.class);
+					assertThat(context.getBean(DataSourceClosingSpringLiquibase.class))
+							.hasFieldOrPropertyWithValue("closeDataSourceOnceMigrated",
+									false);
+				});
+	}
+
+	@Test
+	public void doesNotDisableCloseOfDataSourceWhenEndpointIsDisabled() {
+		this.contextRunner
+				.withUserConfiguration(DataSourceClosingLiquibaseConfiguration.class)
+				.withPropertyValues("management.endpoint.liquibase.enabled:false")
+				.run((context) -> {
+					assertThat(context).doesNotHaveBean(LiquibaseEndpoint.class);
+					DataSourceClosingSpringLiquibase bean = context
+							.getBean(DataSourceClosingSpringLiquibase.class);
+					assertThat(bean).hasFieldOrPropertyWithValue(
+							"closeDataSourceOnceMigrated", true);
+				});
+	}
+
 	@Configuration
 	static class LiquibaseConfiguration {
 
@@ -64,4 +91,33 @@ public SpringLiquibase liquibase() {
 
 	}
 
+	@Configuration
+	static class DataSourceClosingLiquibaseConfiguration {
+
+		@Bean
+		public SpringLiquibase liquibase() {
+			return new DataSourceClosingSpringLiquibase() {
+
+				private boolean propertiesSet = false;
+
+				@Override
+				public void setCloseDataSourceOnceMigrated(
+						boolean closeDataSourceOnceMigrated) {
+					if (this.propertiesSet) {
+						throw new IllegalStateException("setCloseDataSourceOnceMigrated "
+								+ "invoked after afterPropertiesSet");
+					}
+					super.setCloseDataSourceOnceMigrated(closeDataSourceOnceMigrated);
+				}
+
+				@Override
+				public void afterPropertiesSet() throws LiquibaseException {
+					this.propertiesSet = true;
+				}
+
+			};
+		}
+
+	}
+
 }

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/liquibase/DataSourceClosingSpringLiquibase.java

@@ -0,0 +1,68 @@
+/*
+ * Copyright 2012-2018 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
+ *
+ *      http://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.liquibase;
+
+import java.lang.reflect.Method;
+
+import javax.sql.DataSource;
+
+import liquibase.exception.LiquibaseException;
+import liquibase.integration.spring.SpringLiquibase;
+
+import org.springframework.beans.factory.DisposableBean;
+import org.springframework.util.ReflectionUtils;
+
+/**
+ * A custom {@link SpringLiquibase} extension that closes the underlying
+ * {@link DataSource} once the database has been migrated.
+ *
+ * @author Andy Wilkinson
+ * @since 2.0.6
+ */
+public class DataSourceClosingSpringLiquibase extends SpringLiquibase
+		implements DisposableBean {
+
+	private volatile boolean closeDataSourceOnceMigrated = true;
+
+	public void setCloseDataSourceOnceMigrated(boolean closeDataSourceOnceMigrated) {
+		this.closeDataSourceOnceMigrated = closeDataSourceOnceMigrated;
+	}
+
+	@Override
+	public void afterPropertiesSet() throws LiquibaseException {
+		super.afterPropertiesSet();
+		if (this.closeDataSourceOnceMigrated) {
+			closeDataSource();
+		}
+	}
+
+	private void closeDataSource() {
+		Class<?> dataSourceClass = getDataSource().getClass();
+		Method closeMethod = ReflectionUtils.findMethod(dataSourceClass, "close");
+		if (closeMethod != null) {
+			ReflectionUtils.invokeMethod(closeMethod, getDataSource());
+		}
+	}
+
+	@Override
+	public void destroy() throws Exception {
+		if (!this.closeDataSourceOnceMigrated) {
+			closeDataSource();
+		}
+	}
+
+}

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/liquibase/LiquibaseAutoConfiguration.java

@@ -16,7 +16,6 @@
 
 package org.springframework.boot.autoconfigure.liquibase;
 
-import java.lang.reflect.Method;
 import java.util.Collections;
 import java.util.List;
 import java.util.function.Supplier;
@@ -26,7 +25,6 @@
 import javax.sql.DataSource;
 
 import liquibase.change.DatabaseChange;
-import liquibase.exception.LiquibaseException;
 import liquibase.integration.spring.SpringLiquibase;
 
 import org.springframework.beans.factory.ObjectProvider;
@@ -52,7 +50,6 @@
 import org.springframework.orm.jpa.AbstractEntityManagerFactoryBean;
 import org.springframework.orm.jpa.LocalContainerEntityManagerFactoryBean;
 import org.springframework.util.Assert;
-import org.springframework.util.ReflectionUtils;
 
 /**
  * {@link EnableAutoConfiguration Auto-configuration} for Liquibase.
@@ -207,26 +204,4 @@ public LiquibaseJdbcOperationsDependencyConfiguration() {
 
 	}
 
-	/**
-	 * A custom {@link SpringLiquibase} extension that closes the underlying
-	 * {@link DataSource} once the database has been migrated.
-	 */
-	private static final class DataSourceClosingSpringLiquibase extends SpringLiquibase {
-
-		@Override
-		public void afterPropertiesSet() throws LiquibaseException {
-			super.afterPropertiesSet();
-			closeDataSource();
-		}
-
-		private void closeDataSource() {
-			Class<?> dataSourceClass = getDataSource().getClass();
-			Method closeMethod = ReflectionUtils.findMethod(dataSourceClass, "close");
-			if (closeMethod != null) {
-				ReflectionUtils.invokeMethod(closeMethod, getDataSource());
-			}
-		}
-
-	}
-
 }