From 0f14ade35a1e2187b977df2dd0d9dbb062bb17a9 Mon Sep 17 00:00:00 2001
From: Kirk Clemens <Kirk.Clemens@target.com>
Date: Thu, 13 Aug 2015 12:07:37 -0500
Subject: [PATCH] DATACASS-231 - Add ProtocolVersion to
 CassandraCqlClusterFactoryBean.

Original pull request: #41.
CLA: 134320150813044603 (Kirk Clemens)
---
 .../CassandraCqlClusterFactoryBean.java       | 10 +++++
 .../java/AbstractClusterConfiguration.java    |  6 +++
 .../CassandraCqlClusterFactoryBeanTest.java   | 43 +++++++++++++++++++
 3 files changed, 59 insertions(+)
 create mode 100644 spring-cql/src/test/java/org/springframework/cassandra/test/unit/config/CassandraCqlClusterFactoryBeanTest.java

diff --git a/spring-cql/src/main/java/org/springframework/cassandra/config/CassandraCqlClusterFactoryBean.java b/spring-cql/src/main/java/org/springframework/cassandra/config/CassandraCqlClusterFactoryBean.java
index b728de101..48f4bf675 100644
--- a/spring-cql/src/main/java/org/springframework/cassandra/config/CassandraCqlClusterFactoryBean.java
+++ b/spring-cql/src/main/java/org/springframework/cassandra/config/CassandraCqlClusterFactoryBean.java
@@ -43,6 +43,7 @@ import com.datastax.driver.core.Host;
 import com.datastax.driver.core.LatencyTracker;
 import com.datastax.driver.core.PoolingOptions;
 import com.datastax.driver.core.ProtocolOptions.Compression;
+import com.datastax.driver.core.ProtocolVersion;
 import com.datastax.driver.core.SSLOptions;
 import com.datastax.driver.core.Session;
 import com.datastax.driver.core.SocketOptions;
@@ -85,6 +86,7 @@ public class CassandraCqlClusterFactoryBean implements FactoryBean<Cluster>, Ini
 	private LoadBalancingPolicy loadBalancingPolicy;
 	private ReconnectionPolicy reconnectionPolicy;
 	private RetryPolicy retryPolicy;
+	private ProtocolVersion protocolVersion;
 	private boolean metricsEnabled = DEFAULT_METRICS_ENABLED;
 	private boolean jmxReportingEnabled = DEFAULT_JMX_REPORTING_ENABLED;
 	private boolean sslEnabled = DEFAULT_SSL_ENABLED;
@@ -178,6 +180,10 @@ public class CassandraCqlClusterFactoryBean implements FactoryBean<Cluster>, Ini
 			}
 		}
 
+		if(protocolVersion != null) {
+			builder.withProtocolVersion(protocolVersion);
+		}
+
 		cluster = builder.build();
 
 		if (hostStateListener != null) {
@@ -286,6 +292,10 @@ public class CassandraCqlClusterFactoryBean implements FactoryBean<Cluster>, Ini
 		this.poolingOptions = poolingOptions;
 	}
 
+	public void setProtocolVersion(ProtocolVersion protocolVersion){
+		this.protocolVersion = protocolVersion;
+	}
+
 	public void setSocketOptions(SocketOptions socketOptions) {
 		this.socketOptions = socketOptions;
 	}
diff --git a/spring-cql/src/main/java/org/springframework/cassandra/config/java/AbstractClusterConfiguration.java b/spring-cql/src/main/java/org/springframework/cassandra/config/java/AbstractClusterConfiguration.java
index ae99cd977..6662d32a7 100644
--- a/spring-cql/src/main/java/org/springframework/cassandra/config/java/AbstractClusterConfiguration.java
+++ b/spring-cql/src/main/java/org/springframework/cassandra/config/java/AbstractClusterConfiguration.java
@@ -27,6 +27,7 @@ import org.springframework.context.annotation.Configuration;
 
 import com.datastax.driver.core.AuthProvider;
 import com.datastax.driver.core.PoolingOptions;
+import com.datastax.driver.core.ProtocolVersion;
 import com.datastax.driver.core.SocketOptions;
 import com.datastax.driver.core.policies.LoadBalancingPolicy;
 import com.datastax.driver.core.policies.ReconnectionPolicy;
@@ -59,6 +60,7 @@ public abstract class AbstractClusterConfiguration {
 		bean.setShutdownScripts(getShutdownScripts());
 		bean.setSocketOptions(getSocketOptions());
 		bean.setStartupScripts(getStartupScripts());
+		bean.setProtocolVersion(getProtocolVersion());
 
 		return bean;
 	}
@@ -118,4 +120,8 @@ public abstract class AbstractClusterConfiguration {
 	protected AuthProvider getAuthProvider() {
 		return null;
 	}
+
+	protected ProtocolVersion getProtocolVersion(){
+		return null;
+    }
 }
diff --git a/spring-cql/src/test/java/org/springframework/cassandra/test/unit/config/CassandraCqlClusterFactoryBeanTest.java b/spring-cql/src/test/java/org/springframework/cassandra/test/unit/config/CassandraCqlClusterFactoryBeanTest.java
new file mode 100644
index 000000000..e43edf540
--- /dev/null
+++ b/spring-cql/src/test/java/org/springframework/cassandra/test/unit/config/CassandraCqlClusterFactoryBeanTest.java
@@ -0,0 +1,43 @@
+package org.springframework.cassandra.test.unit.config;
+
+import com.datastax.driver.core.ProtocolVersion;
+import org.junit.Assert;
+import org.junit.Test;
+import org.springframework.cassandra.config.CassandraCqlClusterFactoryBean;
+
+/**
+ * JUnit test case for CassandraCqlClusterFactoryBean
+ *
+ * @author Kirk Clemens
+ */
+public class CassandraCqlClusterFactoryBeanTest{
+
+	@Test
+	public void testProtocolVersion(){
+		final CassandraCqlClusterFactoryBean cassandraCqlClusterFactoryBean = new CassandraCqlClusterFactoryBean();
+		cassandraCqlClusterFactoryBean.setProtocolVersion(ProtocolVersion.V2);
+		try {
+			cassandraCqlClusterFactoryBean.afterPropertiesSet();
+			Assert.assertEquals(
+				ProtocolVersion.V2,
+				cassandraCqlClusterFactoryBean.getObject().getConfiguration().getProtocolOptions().getProtocolVersionEnum()
+			);
+		} catch (Exception e) {
+			Assert.fail("Unable to create CQL cluster bean" + e.getMessage());
+		}
+	}
+
+	@Test
+	public void testDefaultProtocolVersion(){
+		final CassandraCqlClusterFactoryBean cassandraCqlClusterFactoryBean = new CassandraCqlClusterFactoryBean();
+
+		try {
+			cassandraCqlClusterFactoryBean.afterPropertiesSet();
+			Assert.assertNull(
+				cassandraCqlClusterFactoryBean.getObject().getConfiguration().getProtocolOptions().getProtocolVersionEnum()
+			);
+		} catch (Exception e) {
+			Assert.fail("Unable to create CQL cluster bean" + e.getMessage());
+		}
+	}
+}