Uses new Bootstrapper interface to allow config data customization. (#675)

spring-cloud-consul-config/src/main/java/org/springframework/cloud/consul/config/ConsulBootstrapper.java

@@ -0,0 +1,35 @@
+/*
+ * Copyright 2015-2020 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.cloud.consul.config;
+
+import java.util.function.Function;
+
+import com.ecwid.consul.v1.ConsulClient;
+
+import org.springframework.boot.Bootstrapper;
+import org.springframework.cloud.consul.ConsulProperties;
+
+public abstract class ConsulBootstrapper {
+
+	static Bootstrapper withConsulClient(Function<ConsulProperties, ConsulClient> factory) {
+		return registry -> registry.register(ConsulClient.class, context -> {
+			ConsulProperties properties = context.get(ConsulProperties.class);
+			return factory.apply(properties);
+		});
+	}
+
+}

spring-cloud-consul-config/src/main/java/org/springframework/cloud/consul/config/ConsulConfigDataLoader.java

@@ -28,7 +28,6 @@ import org.springframework.boot.context.config.ConfigData;
 import org.springframework.boot.context.config.ConfigDataLoader;
 import org.springframework.boot.context.config.ConfigDataLoaderContext;
 import org.springframework.boot.context.config.ConfigDataLocationNotFoundException;
-import org.springframework.boot.env.BootstrapRegistry.Registration;
 
 import static org.springframework.cloud.consul.config.ConsulConfigProperties.Format.FILES;
 
@@ -76,11 +75,10 @@ public class ConsulConfigDataLoader implements ConfigDataLoader<ConsulConfigData
 	}
 
 	protected <T> T getBean(ConfigDataLoaderContext context, Class<T> type) {
-		Registration<T> registration = context.getBootstrapRegistry().getRegistration(type);
-		if (registration == null) {
-			return null;
+		if (context.getBootstrapContext().isRegistered(type)) {
+			return context.getBootstrapContext().get(type);
 		}
-		return registration.get();
+		return null;
 	}
 
 	protected ConsulPropertySource create(ConfigDataLoaderContext context, ConsulConfigDataLocation location) {

spring-cloud-consul-config/src/main/java/org/springframework/cloud/consul/config/ConsulConfigDataLocationResolver.java

@@ -26,6 +26,8 @@ import java.util.stream.Collectors;
 
 import com.ecwid.consul.v1.ConsulClient;
 
+import org.springframework.boot.BootstrapRegistry.InstanceSupplier;
+import org.springframework.boot.ConfigurableBootstrapContext;
 import org.springframework.boot.context.config.ConfigDataLocationNotFoundException;
 import org.springframework.boot.context.config.ConfigDataLocationResolver;
 import org.springframework.boot.context.config.ConfigDataLocationResolverContext;
@@ -78,14 +80,16 @@ public class ConsulConfigDataLocationResolver implements ConfigDataLocationResol
 
 		UriComponents locationUri = parseLocation(context, location);
 
-		ConsulConfigProperties properties = loadConfigProperties(context.getBinder(), locationUri);
+		ConsulConfigProperties properties = loadConfigProperties(context.getBinder());
 
 		List<String> contexts = (locationUri == null || CollectionUtils.isEmpty(locationUri.getPathSegments()))
 				? getAutomaticContexts(profiles, properties) : getCustomContexts(locationUri, properties);
 
-		registerBean(context, ConsulClient.class, () -> createConsulClient(context, locationUri));
+		registerBean(context, ConsulProperties.class, loadProperties(context.getBinder(), locationUri));
 
-		registerBean(context, ConsulConfigIndexes.class, ConsulConfigDataIndexes::new);
+		registerAndPromoteBean(context, ConsulClient.class, () -> createConsulClient(context));
+
+		registerAndPromoteBean(context, ConsulConfigIndexes.class, ConsulConfigDataIndexes::new);
 
 		return contexts.stream()
 				.map(propertySourceContext -> new ConsulConfigDataLocation(properties, propertySourceContext, optional))
@@ -171,14 +175,28 @@ public class ConsulConfigDataLocationResolver implements ConfigDataLocationResol
 		return UriComponentsBuilder.fromUriString(uri).build();
 	}
 
-	protected <T> void registerBean(ConfigDataLocationResolverContext context, Class<T> type, Supplier<T> supplier) {
-		context.getBootstrapRegistry().register(type, supplier)
-				.onApplicationContextPrepared((ctxt, consulClient) -> ctxt.getBeanFactory()
-						.registerSingleton("configData" + type.getSimpleName(), consulClient));
+	public <T> void registerBean(ConfigDataLocationResolverContext context, Class<T> type, T instance) {
+		context.getBootstrapContext().registerIfAbsent(type, InstanceSupplier.of(instance));
 	}
 
-	protected ConsulClient createConsulClient(ConfigDataLocationResolverContext context, UriComponents location) {
-		ConsulProperties properties = loadProperties(context.getBinder(), location);
+	protected <T> void registerAndPromoteBean(ConfigDataLocationResolverContext context, Class<T> type,
+			Supplier<T> supplier) {
+		registerBean(context, type, supplier);
+		context.getBootstrapContext().addCloseListener(event -> {
+			T instance = event.getBootstrapContext().get(type);
+			event.getApplicationContext().getBeanFactory().registerSingleton("configData" + type.getSimpleName(),
+					instance);
+		});
+	}
+
+	protected <T> void registerBean(ConfigDataLocationResolverContext context, Class<T> type, Supplier<T> supplier) {
+		ConfigurableBootstrapContext bootstrapContext = context.getBootstrapContext();
+		bootstrapContext.registerIfAbsent(type, InstanceSupplier.from(supplier));
+	}
+
+	protected ConsulClient createConsulClient(ConfigDataLocationResolverContext context) {
+		ConsulProperties properties = context.getBootstrapContext().get(ConsulProperties.class);
+
 		return ConsulAutoConfiguration.createConsulClient(properties);
 	}
 
@@ -198,7 +216,7 @@ public class ConsulConfigDataLocationResolver implements ConfigDataLocationResol
 		return consulProperties;
 	}
 
-	protected ConsulConfigProperties loadConfigProperties(Binder binder, UriComponents location) {
+	protected ConsulConfigProperties loadConfigProperties(Binder binder) {
 		ConsulConfigProperties properties = binder
 				.bind(ConsulConfigProperties.PREFIX, Bindable.of(ConsulConfigProperties.class))
 				.orElse(new ConsulConfigProperties());

spring-cloud-consul-config/src/test/java/org/springframework/cloud/consul/config/ConsulConfigDataCustomizationIntegrationTests.java

@@ -0,0 +1,92 @@
+/*
+ * Copyright 2013-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.
+ */
+
+package org.springframework.cloud.consul.config;
+
+import java.util.UUID;
+
+import com.ecwid.consul.v1.ConsulClient;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+
+import org.springframework.boot.SpringApplication;
+import org.springframework.boot.WebApplicationType;
+import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
+import org.springframework.cloud.consul.ConsulProperties;
+import org.springframework.cloud.consul.test.ConsulTestcontainers;
+import org.springframework.context.ConfigurableApplicationContext;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.test.annotation.DirtiesContext;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * @author Spencer Gibb
+ */
+@DirtiesContext
+public class ConsulConfigDataCustomizationIntegrationTests {
+
+	private static final String APP_NAME = "testConsulConfigDataCustomization";
+
+	private static final String PREFIX = "_configDataIntegrationTests_config__";
+
+	private static final String ROOT = PREFIX + UUID.randomUUID();
+
+	private static ConfigurableApplicationContext context;
+
+	@BeforeAll
+	public static void setup() {
+		ConsulTestcontainers.start();
+
+		SpringApplication application = new SpringApplication(Config.class);
+		application.setWebApplicationType(WebApplicationType.NONE);
+		application.addBootstrapper(ConsulBootstrapper.withConsulClient(TestConsulClient::new));
+		context = application.run("--spring.application.name=" + APP_NAME,
+				"--spring.config.import=consul:" + ConsulTestcontainers.getHost() + ":"
+						+ ConsulTestcontainers.getPort(),
+				"--spring.cloud.consul.config.prefix=" + ROOT, "--spring.cloud.consul.config.watch.delay=10");
+
+	}
+
+	@AfterAll
+	public static void teardown() {
+		if (context != null) {
+			context.close();
+		}
+	}
+
+	@Test
+	public void consulClientIsCustom() {
+		ConsulClient client = context.getBean(ConsulClient.class);
+		assertThat(client).isInstanceOf(TestConsulClient.class);
+	}
+
+	static class TestConsulClient extends ConsulClient {
+
+		TestConsulClient(ConsulProperties properties) {
+			super(properties.getHost(), properties.getPort());
+		}
+
+	}
+
+	@Configuration
+	@EnableAutoConfiguration
+	static class Config {
+
+	}
+
+}

spring-cloud-consul-config/src/test/java/org/springframework/cloud/consul/config/ConsulConfigDataLocationResolverTests.java

@@ -91,7 +91,7 @@ public class ConsulConfigDataLocationResolverTests {
 	private ConsulConfigDataLocationResolver createResolver() {
 		ConsulConfigDataLocationResolver resolver = new ConsulConfigDataLocationResolver() {
 			@Override
-			protected <T> void registerBean(ConfigDataLocationResolverContext context, Class<T> type,
+			protected <T> void registerAndPromoteBean(ConfigDataLocationResolverContext context, Class<T> type,
 					Supplier<T> supplier) {
 				// do nothing
 			}