From 9ec7dbd5c79de9b96352096588d5610dc17a67ce Mon Sep 17 00:00:00 2001
From: Jonathan Shook <jshook@gmail.com>
Date: Thu, 11 Mar 2021 11:09:07 -0600
Subject: [PATCH] add protocol_version to cql driver

---
 .../cql/statements/core/CQLSessionCache.java  | 54 +++++++++----------
 driver-cql-shaded/src/main/resources/cql.md   | 20 ++++---
 2 files changed, 39 insertions(+), 35 deletions(-)

diff --git a/driver-cql-shaded/src/main/java/io/nosqlbench/activitytype/cql/statements/core/CQLSessionCache.java b/driver-cql-shaded/src/main/java/io/nosqlbench/activitytype/cql/statements/core/CQLSessionCache.java
index 7db1c24bd..f46bede03 100644
--- a/driver-cql-shaded/src/main/java/io/nosqlbench/activitytype/cql/statements/core/CQLSessionCache.java
+++ b/driver-cql-shaded/src/main/java/io/nosqlbench/activitytype/cql/statements/core/CQLSessionCache.java
@@ -1,32 +1,7 @@
 package io.nosqlbench.activitytype.cql.statements.core;
 
-import java.io.File;
-import java.io.IOException;
-import java.net.Inet6Address;
-import java.net.InetAddress;
-import java.net.InetSocketAddress;
-import java.net.UnknownHostException;
-import java.nio.file.Files;
-import java.nio.file.Path;
-import java.nio.file.Paths;
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.Optional;
-import javax.net.ssl.SSLContext;
-
-import com.datastax.driver.core.Cluster;
-import com.datastax.driver.core.NettyOptions;
-import com.datastax.driver.core.ProtocolOptions;
-import com.datastax.driver.core.RemoteEndpointAwareJdkSSLOptions;
-import com.datastax.driver.core.Session;
-import com.datastax.driver.core.policies.DefaultRetryPolicy;
-import com.datastax.driver.core.policies.LoadBalancingPolicy;
-import com.datastax.driver.core.policies.LoggingRetryPolicy;
-import com.datastax.driver.core.policies.RetryPolicy;
-import com.datastax.driver.core.policies.RoundRobinPolicy;
-import com.datastax.driver.core.policies.WhiteListPolicy;
+import com.datastax.driver.core.*;
+import com.datastax.driver.core.policies.*;
 import com.datastax.driver.dse.DseCluster;
 import io.netty.channel.socket.SocketChannel;
 import io.netty.handler.codec.haproxy.HAProxyCommand;
@@ -45,6 +20,18 @@ import io.nosqlbench.nb.api.errors.BasicError;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
+import javax.net.ssl.SSLContext;
+import java.io.File;
+import java.io.IOException;
+import java.net.Inet6Address;
+import java.net.InetAddress;
+import java.net.InetSocketAddress;
+import java.net.UnknownHostException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.*;
+
 public class CQLSessionCache implements Shutdownable {
 
     private final static Logger logger = LogManager.getLogger(CQLSessionCache.class);
@@ -81,8 +68,8 @@ public class CQLSessionCache implements Shutdownable {
         String driverType = activityDef.getParams().getOptionalString("cqldriver").orElse("dse");
 
         Cluster.Builder builder =
-            driverType.toLowerCase().equals("dse") ? DseCluster.builder() :
-                driverType.toLowerCase().equals("oss") ? Cluster.builder() : null;
+            driverType.equalsIgnoreCase("dse") ? DseCluster.builder() :
+                driverType.equalsIgnoreCase("oss") ? Cluster.builder() : null;
 
         if (builder == null) {
             throw new RuntimeException("The driver type '" + driverType + "' is not recognized");
@@ -185,6 +172,15 @@ public class CQLSessionCache implements Shutdownable {
                 .ifPresent(builder::withSpeculativeExecutionPolicy);
         }
 
+        activityDef.getParams().getOptionalString("protocol_version")
+            .map(String::toUpperCase)
+            .map(ProtocolVersion::valueOf)
+            .map(pv -> {
+                logger.info("protocol_version=>" + pv);
+                return pv;
+            })
+            .ifPresent(builder::withProtocolVersion);
+
         activityDef.getParams().getOptionalString("socketoptions")
             .map(sockopts -> {
                 logger.info("socketoptions=>" + sockopts);
diff --git a/driver-cql-shaded/src/main/resources/cql.md b/driver-cql-shaded/src/main/resources/cql.md
index 3f2c0e461..64500def6 100644
--- a/driver-cql-shaded/src/main/resources/cql.md
+++ b/driver-cql-shaded/src/main/resources/cql.md
@@ -62,14 +62,22 @@ activity types.
     The only option supported for this version is `retrypolicy=logging`,
     which uses the default retry policy, but with logging added.
 
-- **reconnectpolicy** default: none - Applies a reconnection policy in the driver
-    Supports either `reconnectpolicy=exponential(minDelayInMs,maxDelayInMs)` or `reconnectpolicy=constant(delayInMs)`.
-    The driver reconnects using this policy when the entire cluster becomes unavailable.
+- **reconnectpolicy** default: none - Applies a reconnection policy in the
+  driver Supports
+  either `reconnectpolicy=exponential(minDelayInMs,maxDelayInMs)`
+  or `reconnectpolicy=constant(delayInMs)`. The driver reconnects using
+  this policy when the entire cluster becomes unavailable.
 
+- **protocol_version** default: unset, defaults to driver default behavior
+  - Set the CQL protocol version. Valid values are V1, V2, V3, V4, V5,
+  DSE_V1, DSE_V2. Protocol is usually auto-negotiated, however, the
+  initial connection may use a lower protocol to ensure connectivity to
+  older server versions. If you know you are running on a newer server
+  version, you can set this to higher than the default for the initial
+  control connection, which is V2.
 
-- **pooling** default: none - Applies the connection pooling options
-    to the policy.
-    Examples:
+- **pooling** default: none - Applies the connection pooling options to
+  the policy. Examples:
     - `pooling=4:10`
       keep between 4 and 10 connections to LOCAL hosts
     - `pooling=4:10,2:5`