Add protocol_version to conn_config for Cassandrahook (#11036)

airflow/providers/apache/cassandra/hooks/cassandra.py

@@ -113,6 +113,10 @@ class CassandraHook(BaseHook, LoggingMixin):
         if ssl_options:
             conn_config['ssl_options'] = ssl_options
 
+        protocol_version = conn.extra_dejson.get('protocol_version', None)
+        if protocol_version:
+            conn_config['protocol_version'] = protocol_version
+
         self.cluster = Cluster(**conn_config)
         self.keyspace = conn.schema
         self.session = None

docs/howto/connection/cassandra.rst

@@ -53,6 +53,7 @@ Extra (optional)
       ``RoundRobinPolicy``, ``DCAwareRoundRobinPolicy``, ``WhiteListRoundRobinPolicy`` and ``TokenAwarePolicy``. ``RoundRobinPolicy`` is the default load balancing policy.
     * ``load_balancing_policy_args`` - This parameter specifies the arguments for the load balancing policy being used.
     * ``cql_version`` - This parameter specifies the CQL version of cassandra.
+    * ``protocol_version`` - This parameter specifies the maximum version of the native protocol to use.
     * ``ssl_options`` - This parameter specifies the details related to SSL, if it's enabled in Cassandra.
 
 

tests/providers/apache/cassandra/hooks/test_cassandra.py

@@ -43,7 +43,7 @@ class TestCassandraHook(unittest.TestCase):
                 host='host-1,host-2',
                 port='9042',
                 schema='test_keyspace',
-                extra='{"load_balancing_policy":"TokenAwarePolicy"}',
+                extra='{"load_balancing_policy":"TokenAwarePolicy","protocol_version":4}',
             )
         )
         db.merge_conn(
@@ -84,6 +84,7 @@ class TestCassandraHook(unittest.TestCase):
             cluster = hook.get_cluster()
             self.assertEqual(cluster.contact_points, ['host-1', 'host-2'])
             self.assertEqual(cluster.port, 9042)
+            self.assertEqual(cluster.protocol_version, 4)
             self.assertTrue(isinstance(cluster.load_balancing_policy, TokenAwarePolicy))
 
     def test_get_lb_policy_with_no_args(self):