Allow to customize Casssandra's Cluster

This commit allows to configure special customizers to be used by
Cluster.

Closes gh-7312

spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/cassandra/CassandraAutoConfiguration.java

@@ -16,6 +16,8 @@
 
 package org.springframework.boot.autoconfigure.cassandra;
 
+import java.util.List;
+
 import com.datastax.driver.core.Cluster;
 import com.datastax.driver.core.QueryOptions;
 import com.datastax.driver.core.SocketOptions;
@@ -24,6 +26,7 @@ import com.datastax.driver.core.policies.ReconnectionPolicy;
 import com.datastax.driver.core.policies.RetryPolicy;
 
 import org.springframework.beans.BeanUtils;
+import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
@@ -46,8 +49,12 @@ public class CassandraAutoConfiguration {
 
 	private final CassandraProperties properties;
 
-	public CassandraAutoConfiguration(CassandraProperties properties) {
+	private final List<ClusterCustomizer> customizers;
+
+	public CassandraAutoConfiguration(CassandraProperties properties,
+			ObjectProvider<List<ClusterCustomizer>> customizers) {
 		this.properties = properties;
+		this.customizers = customizers.getIfAvailable();
 	}
 
 	@Bean
@@ -82,7 +89,18 @@ public class CassandraAutoConfiguration {
 		}
 		String points = properties.getContactPoints();
 		builder.addContactPoints(StringUtils.commaDelimitedListToStringArray(points));
-		return builder.build();
+
+		Cluster cluster = builder.build();
+		customize(cluster);
+		return cluster;
+	}
+
+	private void customize(Cluster cluster) {
+		if (this.customizers != null) {
+			for (ClusterCustomizer customizer : this.customizers) {
+				customizer.customize(cluster);
+			}
+		}
 	}
 
 	public static <T> T instantiate(Class<T> type) {

spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/cassandra/ClusterCustomizer.java

@@ -0,0 +1,36 @@
+/*
+ * Copyright 2012-2016 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.cassandra;
+
+import com.datastax.driver.core.Cluster;
+
+/**
+ * Callback interface that can be used to customize a {@link Cluster}.
+ *
+ * @author Eddú Meléndez
+ * @since 1.5.0
+ * @see CassandraAutoConfiguration
+ */
+public interface ClusterCustomizer {
+
+	/**
+	 * Customize the cluster.
+	 * @param cluster the cluster to customize
+	 */
+	void customize(Cluster cluster);
+
+}

spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/cassandra/CassandraAutoConfigurationTests.java

@@ -23,8 +23,11 @@ import org.junit.Test;
 import org.springframework.boot.autoconfigure.PropertyPlaceholderAutoConfiguration;
 import org.springframework.boot.test.util.EnvironmentTestUtils;
 import org.springframework.context.annotation.AnnotationConfigApplicationContext;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
 
 /**
  * Tests for {@link CassandraAutoConfiguration}
@@ -45,6 +48,7 @@ public class CassandraAutoConfigurationTests {
 	@Test
 	public void createClusterWithDefault() {
 		this.context = doLoad();
+		this.context.refresh();
 		assertThat(this.context.getBeanNamesForType(Cluster.class).length).isEqualTo(1);
 		Cluster cluster = this.context.getBean(Cluster.class);
 		assertThat(cluster.getClusterName()).startsWith("cluster");
@@ -53,18 +57,47 @@ public class CassandraAutoConfigurationTests {
 	@Test
 	public void createClusterWithOverrides() {
 		this.context = doLoad("spring.data.cassandra.cluster-name=testcluster");
+		this.context.refresh();
 		assertThat(this.context.getBeanNamesForType(Cluster.class).length).isEqualTo(1);
 		Cluster cluster = this.context.getBean(Cluster.class);
 		assertThat(cluster.getClusterName()).isEqualTo("testcluster");
 	}
 
-	private AnnotationConfigApplicationContext doLoad(String... environment) {
+	@Test
+	public void createCustomizeCluster() {
+		this.context = doLoad(ClusterConfig.class);
+		this.context.refresh();
+		assertThat(this.context.getBeanNamesForType(Cluster.class).length).isEqualTo(1);
+		assertThat(this.context.getBeanNamesForType(ClusterCustomizer.class).length).isEqualTo(1);
+	}
+
+	private AnnotationConfigApplicationContext doLoad() {
 		AnnotationConfigApplicationContext applicationContext = new AnnotationConfigApplicationContext();
-		EnvironmentTestUtils.addEnvironment(applicationContext, environment);
 		applicationContext.register(PropertyPlaceholderAutoConfiguration.class,
 				CassandraAutoConfiguration.class);
-		applicationContext.refresh();
 		return applicationContext;
 	}
 
+	private AnnotationConfigApplicationContext doLoad(Class<?> clazz) {
+		AnnotationConfigApplicationContext applicationContext = doLoad();
+		applicationContext.register(clazz);
+		return applicationContext;
+	}
+
+	private AnnotationConfigApplicationContext doLoad(String... environment) {
+		AnnotationConfigApplicationContext applicationContext = doLoad();
+		EnvironmentTestUtils.addEnvironment(applicationContext, environment);
+		return applicationContext;
+	}
+
+	@Configuration
+	static class ClusterConfig {
+
+		@Bean
+		public ClusterCustomizer customizer() {
+			return mock(ClusterCustomizer.class);
+		}
+
+	}
+
 }