From 4d82dc3e127618e22bcdc749a4437db22cff2bb7 Mon Sep 17 00:00:00 2001
From: Sumanth Pasupuleti <spasupuleti@netflix.com>
Date: Wed, 11 Nov 2020 00:50:28 -0800
Subject: [PATCH] Beta protocol support for cass plugins

---
 .../ndbench/plugin/cass/CJavaDriverBasePlugin.java    |  7 ++++++-
 .../ndbench/plugin/cass/CassJavaDriverManager.java    |  2 +-
 .../plugin/cass/CassJavaDriverManagerImpl.java        | 11 ++++++++---
 .../plugin/configs/CassandraConfigurationBase.java    |  3 +++
 4 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CJavaDriverBasePlugin.java b/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CJavaDriverBasePlugin.java
index 3f4c0d14..fbaf054d 100644
--- a/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CJavaDriverBasePlugin.java
+++ b/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CJavaDriverBasePlugin.java
@@ -46,6 +46,7 @@ public abstract class CJavaDriverBasePlugin<C extends CassandraConfigurationBase
     protected volatile Session session;
     protected volatile PreparedStatement readPstmt;
     protected volatile PreparedStatement writePstmt;
+    protected volatile boolean allowBetaProtocol;
 
     /**
      * Creates an instance of the abstract CJavaDriverBasePlugin class. Subclasses calling this method should use
@@ -70,6 +71,7 @@ public void init(DataGenerator dataGenerator) throws Exception {
         this.connections = config.getConnections();
         this.username = config.getUsername();
         this.password = config.getPassword();
+        this.allowBetaProtocol = config.getAllowBetaProtocol();
 
         // we do not set ReadConsistencyLevel and WriteConsistencyLevel and MaxColCount here because the
         // enum classes corresponding to the consistency levels differ among the concrete subclasses and because
@@ -99,9 +101,12 @@ public String runWorkFlow() {
     private void initDriver() {
         logger.info("Cassandra  Cluster: " + clusterName);
 
-        this.cluster = cassJavaDriverManager.registerCluster(clusterName, clusterContactPoint, connections, port,
+        this.cluster = cassJavaDriverManager.registerCluster(clusterName, clusterContactPoint, connections, port, allowBetaProtocol,
                                                              username, password);
         this.session = cassJavaDriverManager.getSession(cluster);
+
+        logger.info("Protocol version in use: {}", session.getCluster().getConfiguration().getProtocolOptions().getProtocolVersion());
+
         if(config.getCreateSchema())
         {
             logger.info("Trying to upsert schema");
diff --git a/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CassJavaDriverManager.java b/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CassJavaDriverManager.java
index cd2ba11b..3f5a5dd5 100644
--- a/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CassJavaDriverManager.java
+++ b/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CassJavaDriverManager.java
@@ -12,7 +12,7 @@
  */
 @ImplementedBy(CassJavaDriverManagerImpl.class)
 public interface CassJavaDriverManager {
-    Cluster registerCluster(String clName, String contactPoint, int connections, int port, String username, String password);
+    Cluster registerCluster(String clName, String contactPoint, int connections, int port, boolean allowBetaProtocol, String username, String password);
     Cluster registerCluster(String clName, String contactPoint, int connections, int port);
     Session getSession(Cluster cluster);
 
diff --git a/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CassJavaDriverManagerImpl.java b/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CassJavaDriverManagerImpl.java
index 967496c1..934d9c39 100644
--- a/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CassJavaDriverManagerImpl.java
+++ b/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CassJavaDriverManagerImpl.java
@@ -20,10 +20,11 @@ public class CassJavaDriverManagerImpl implements CassJavaDriverManager {
 
     @Override
     public Cluster registerCluster(String clName, String contactPoint, int connections, int port) {
-        return registerCluster(clName,contactPoint,connections,port,null,null);
+        return registerCluster(clName,contactPoint,connections,port, false, null,null);
     }
-        @Override
-    public Cluster registerCluster(String clName, String contactPoint, int connections, int port, String username, String password) {
+
+    @Override
+    public Cluster registerCluster(String clName, String contactPoint, int connections, int port, boolean allowBetaProtocol, String username, String password) {
     
         PoolingOptions poolingOpts = new PoolingOptions()
                                      .setConnectionsPerHost(HostDistance.LOCAL, connections, connections)
@@ -40,6 +41,10 @@ public Cluster registerCluster(String clName, String contactPoint, int connectio
             clusterBuilder = clusterBuilder.withCredentials(username, password);
         }
 
+        if (allowBetaProtocol) {
+            clusterBuilder.allowBetaProtocolVersion();
+        }
+
         cluster = clusterBuilder.build();
         return cluster;
     }
diff --git a/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/configs/CassandraConfigurationBase.java b/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/configs/CassandraConfigurationBase.java
index 909ba465..dc99e137 100644
--- a/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/configs/CassandraConfigurationBase.java
+++ b/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/configs/CassandraConfigurationBase.java
@@ -75,4 +75,7 @@ public interface CassandraConfigurationBase {
 
     @DefaultValue("true")
     Boolean getCreateSchema();
+
+    @DefaultValue("false")
+    Boolean getAllowBetaProtocol();
 }