Beta protocol support for cass plugins

Netflix · Nov 11, 2020 · 4d82dc3 · 4d82dc3
1 parent 91ba19a
commit 4d82dc3
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 5 deletions.
diff --git a/...nch-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CJavaDriverBasePlugin.java b/...nch-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/...nch-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CassJavaDriverManager.java b/...nch-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/...cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CassJavaDriverManagerImpl.java b/...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/...-plugins/src/main/java/com/netflix/ndbench/plugin/configs/CassandraConfigurationBase.java b/...-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();
 }