NIFI-4358 This closes #3363. cassandra connection enable compression at resquest and response

Signed-off-by: Joe Witt <joewitt@apache.org>

nifi-nar-bundles/nifi-cassandra-bundle/nifi-cassandra-processors/src/main/java/org/apache/nifi/processors/cassandra/AbstractCassandraProcessor.java

@@ -22,6 +22,7 @@ import com.datastax.driver.core.ConsistencyLevel;
 import com.datastax.driver.core.DataType;
 import com.datastax.driver.core.JdkSSLOptions;
 import com.datastax.driver.core.Metadata;
+import com.datastax.driver.core.ProtocolOptions;
 import com.datastax.driver.core.Row;
 import com.datastax.driver.core.Session;
 import com.datastax.driver.core.TypeCodec;
@@ -137,6 +138,15 @@ public abstract class AbstractCassandraProcessor extends AbstractProcessor {
             .defaultValue("ONE")
             .build();
 
+    static final PropertyDescriptor COMPRESSION_TYPE = new PropertyDescriptor.Builder()
+            .name("Compression Type")
+            .description("Enable compression at transport-level requests and responses")
+            .required(false)
+            .allowableValues(ProtocolOptions.Compression.values())
+            .addValidator(StandardValidators.NON_EMPTY_VALIDATOR)
+            .defaultValue("NONE")
+            .build();
+
     static final PropertyDescriptor CHARSET = new PropertyDescriptor.Builder()
             .name("Character Set")
             .description("Specifies the character set of the record data.")
@@ -172,6 +182,7 @@ public abstract class AbstractCassandraProcessor extends AbstractProcessor {
         descriptors.add(USERNAME);
         descriptors.add(PASSWORD);
         descriptors.add(CONSISTENCY_LEVEL);
+        descriptors.add(COMPRESSION_TYPE);
         descriptors.add(CHARSET);
     }
 
@@ -238,6 +249,7 @@ public abstract class AbstractCassandraProcessor extends AbstractProcessor {
             ComponentLog log = getLogger();
             final String contactPointList = context.getProperty(CONTACT_POINTS).evaluateAttributeExpressions().getValue();
             final String consistencyLevel = context.getProperty(CONSISTENCY_LEVEL).getValue();
+            final String compressionType = context.getProperty(COMPRESSION_TYPE).getValue();
             List<InetSocketAddress> contactPoints = getContactPoints(contactPointList);
 
             // Set up the client for secure (SSL/TLS communications) if configured to do so
@@ -277,7 +289,7 @@ public abstract class AbstractCassandraProcessor extends AbstractProcessor {
             }
 
             // Create the cluster and connect to it
-            Cluster newCluster = createCluster(contactPoints, sslContext, username, password);
+            Cluster newCluster = createCluster(contactPoints, sslContext, username, password, compressionType);
             PropertyValue keyspaceProperty = context.getProperty(KEYSPACE).evaluateAttributeExpressions();
 
             final Session newSession;
@@ -304,16 +316,22 @@ public abstract class AbstractCassandraProcessor extends AbstractProcessor {
      * @param sslContext    The SSL context (used for secure connections)
      * @param username      The username for connection authentication
      * @param password      The password for connection authentication
+     * @param compressionType Enable compression at transport-level requests and responses.
      * @return A reference to the Cluster object associated with the given Cassandra configuration
      */
     protected Cluster createCluster(List<InetSocketAddress> contactPoints, SSLContext sslContext,
-                                    String username, String password) {
+                                    String username, String password, String compressionType) {
         Cluster.Builder builder = Cluster.builder().addContactPointsWithPorts(contactPoints);
         if (sslContext != null) {
             JdkSSLOptions sslOptions = JdkSSLOptions.builder()
                     .withSSLContext(sslContext)
                     .build();
             builder = builder.withSSL(sslOptions);
+            if(ProtocolOptions.Compression.SNAPPY.equals(compressionType)) {
+                builder = builder.withCompression(ProtocolOptions.Compression.SNAPPY);
+            } else if(ProtocolOptions.Compression.LZ4.equals(compressionType)) {
+                builder = builder.withCompression(ProtocolOptions.Compression.LZ4);
+            }
         }
         if (username != null && password != null) {
             builder = builder.withCredentials(username, password);

nifi-nar-bundles/nifi-cassandra-bundle/nifi-cassandra-processors/src/test/java/org/apache/nifi/processors/cassandra/AbstractCassandraProcessorTest.java

@@ -303,7 +303,7 @@ public class AbstractCassandraProcessorTest {
 
         @Override
         protected Cluster createCluster(List<InetSocketAddress> contactPoints, SSLContext sslContext,
-                                        String username, String password) {
+                                        String username, String password, String compressionType) {
             Cluster mockCluster = mock(Cluster.class);
             Metadata mockMetadata = mock(Metadata.class);
             when(mockMetadata.getClusterName()).thenReturn("cluster1");

nifi-nar-bundles/nifi-cassandra-bundle/nifi-cassandra-processors/src/test/java/org/apache/nifi/processors/cassandra/PutCassandraQLTest.java

@@ -314,7 +314,7 @@ public class PutCassandraQLTest {
 
         @Override
         protected Cluster createCluster(List<InetSocketAddress> contactPoints, SSLContext sslContext,
-                                        String username, String password) {
+                                        String username, String password, String compressionType) {
             Cluster mockCluster = mock(Cluster.class);
             try {
                 Metadata mockMetadata = mock(Metadata.class);

nifi-nar-bundles/nifi-cassandra-bundle/nifi-cassandra-processors/src/test/java/org/apache/nifi/processors/cassandra/PutCassandraRecordTest.java

@@ -158,7 +158,7 @@ public class PutCassandraRecordTest {
 
         @Override
         protected Cluster createCluster(List<InetSocketAddress> contactPoints, SSLContext sslContext,
-                                        String username, String password) {
+                                        String username, String password, String compressionType) {
             Cluster mockCluster = mock(Cluster.class);
             try {
                 Metadata mockMetadata = mock(Metadata.class);

nifi-nar-bundles/nifi-cassandra-bundle/nifi-cassandra-processors/src/test/java/org/apache/nifi/processors/cassandra/QueryCassandraTest.java

@@ -385,7 +385,7 @@ public class QueryCassandraTest {
 
         @Override
         protected Cluster createCluster(List<InetSocketAddress> contactPoints, SSLContext sslContext,
-                                        String username, String password) {
+                                        String username, String password, String compressionType) {
             Cluster mockCluster = mock(Cluster.class);
             try {
                 Metadata mockMetadata = mock(Metadata.class);

nifi-nar-bundles/nifi-cassandra-bundle/nifi-cassandra-services/src/main/java/org/apache/nifi/service/CassandraSessionProvider.java

@@ -20,6 +20,7 @@ import com.datastax.driver.core.Cluster;
 import com.datastax.driver.core.ConsistencyLevel;
 import com.datastax.driver.core.JdkSSLOptions;
 import com.datastax.driver.core.Metadata;
+import com.datastax.driver.core.ProtocolOptions;
 import com.datastax.driver.core.Session;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.nifi.annotation.documentation.CapabilityDescription;
@@ -118,6 +119,15 @@ public class CassandraSessionProvider extends AbstractControllerService implemen
             .defaultValue("ONE")
             .build();
 
+    static final PropertyDescriptor COMPRESSION_TYPE = new PropertyDescriptor.Builder()
+            .name("Compression Type")
+            .description("Enable compression at transport-level requests and responses")
+            .required(false)
+            .allowableValues(ProtocolOptions.Compression.values())
+            .addValidator(StandardValidators.NON_EMPTY_VALIDATOR)
+            .defaultValue("NONE")
+            .build();
+
     private List<PropertyDescriptor> properties;
     private Cluster cluster;
     private Session cassandraSession;
@@ -129,6 +139,7 @@ public class CassandraSessionProvider extends AbstractControllerService implemen
         props.add(CONTACT_POINTS);
         props.add(CLIENT_AUTH);
         props.add(CONSISTENCY_LEVEL);
+        props.add(COMPRESSION_TYPE);
         props.add(KEYSPACE);
         props.add(USERNAME);
         props.add(PASSWORD);
@@ -190,6 +201,8 @@ public class CassandraSessionProvider extends AbstractControllerService implemen
             ComponentLog log = getLogger();
             final String contactPointList = context.getProperty(CONTACT_POINTS).evaluateAttributeExpressions().getValue();
             final String consistencyLevel = context.getProperty(CONSISTENCY_LEVEL).getValue();
+            final String compressionType = context.getProperty(COMPRESSION_TYPE).getValue();
+
             List<InetSocketAddress> contactPoints = getContactPoints(contactPointList);
 
             // Set up the client for secure (SSL/TLS communications) if configured to do so
@@ -228,7 +241,7 @@ public class CassandraSessionProvider extends AbstractControllerService implemen
             }
 
             // Create the cluster and connect to it
-            Cluster newCluster = createCluster(contactPoints, sslContext, username, password);
+            Cluster newCluster = createCluster(contactPoints, sslContext, username, password, compressionType);
             PropertyValue keyspaceProperty = context.getProperty(KEYSPACE).evaluateAttributeExpressions();
             final Session newSession;
             if (keyspaceProperty != null) {
@@ -266,7 +279,7 @@ public class CassandraSessionProvider extends AbstractControllerService implemen
     }
 
     private Cluster createCluster(List<InetSocketAddress> contactPoints, SSLContext sslContext,
-                                  String username, String password) {
+                                  String username, String password, String compressionType) {
         Cluster.Builder builder = Cluster.builder().addContactPointsWithPorts(contactPoints);
 
         if (sslContext != null) {
@@ -280,6 +293,12 @@ public class CassandraSessionProvider extends AbstractControllerService implemen
             builder = builder.withCredentials(username, password);
         }
 
+        if(ProtocolOptions.Compression.SNAPPY.equals(compressionType)) {
+            builder = builder.withCompression(ProtocolOptions.Compression.SNAPPY);
+        } else if(ProtocolOptions.Compression.LZ4.equals(compressionType)) {
+            builder = builder.withCompression(ProtocolOptions.Compression.LZ4);
+        }
+
         return builder.build();
     }
 }

nifi-nar-bundles/nifi-cassandra-bundle/nifi-cassandra-services/src/test/java/org/apache/nifi/service/TestCassandraSessionProvider.java

@@ -46,7 +46,7 @@ public class TestCassandraSessionProvider {
     public void testGetPropertyDescriptors() {
         List<PropertyDescriptor> properties = sessionProvider.getPropertyDescriptors();
 
-        assertEquals(7, properties.size());
+        assertEquals(8, properties.size());
         assertTrue(properties.contains(CassandraSessionProvider.CLIENT_AUTH));
         assertTrue(properties.contains(CassandraSessionProvider.CONSISTENCY_LEVEL));
         assertTrue(properties.contains(CassandraSessionProvider.CONTACT_POINTS));