Loosen cassandra-driver requirement to allow latest version (#15022)

For py 3.6-3.8 we don't need the CASS_ settings in the Dockerfiles but
as Datastax haven't yet published a Python 3.9 wheel, I have kept the
settings for now -- they are ignored when installing a wheel, but will
be used if the python version (i.e. 3.9) doesn't have a published wheel.

.pre-commit-config.yaml

@@ -321,7 +321,9 @@ repos:
           ^airflow/providers/apache/hive/.*README.md$|
           ^tests/providers/apache/cassandra/hooks/test_cassandra.py$|
           ^docs/apache-airflow-providers-apache-cassandra/connections/cassandra.rst$|
-          ^docs/apache-airflow-providers-apache-hive/commits.rst$|git|
+          ^docs/apache-airflow-providers-apache-hive/commits.rst$|
+          git|
+          ^pylintrc |
           ^CHANGELOG.txt$
       - id: consistent-pylint
         language: pygrep

pylintrc

@@ -54,6 +54,8 @@ suggestion-mode=yes
 # active Python interpreter and may run arbitrary code.
 unsafe-load-any-extension=no
 
+# Pylint 2.8 changes this to `extension-pkg-allow-list`
+extension-pkg-whitelist=cassandra
 
 [MESSAGES CONTROL]
 

pylintrc-tests

@@ -54,6 +54,8 @@ suggestion-mode=yes
 # active Python interpreter and may run arbitrary code.
 unsafe-load-any-extension=no
 
+# Pylint 2.8 changes this to `extension-pkg-allow-list`
+extension-pkg-whitelist=cassandra
 
 [MESSAGES CONTROL]
 

setup.py

@@ -224,7 +224,7 @@ def get_sphinx_theme_version() -> str:
     'azure-storage-file>=2.1.0',
 ]
 cassandra = [
-    'cassandra-driver>=3.13.0,<3.21.0',
+    'cassandra-driver>=3.13.0,<4',
 ]
 celery = [
     'celery~=4.4.2',

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

@@ -72,20 +72,17 @@ def setUp(self):
         hook.shutdown_cluster()
 
     def test_get_conn(self):
-        with mock.patch.object(Cluster, "connect") as mock_connect, mock.patch(
-            "socket.getaddrinfo", return_value=[]
-        ) as mock_getaddrinfo:
-            mock_connect.return_value = 'session'
-            hook = CassandraHook(cassandra_conn_id='cassandra_test')
-            hook.get_conn()
-            assert mock_getaddrinfo.called
-            mock_connect.assert_called_once_with('test_keyspace')
-
-            cluster = hook.get_cluster()
-            assert cluster.contact_points == ['host-1', 'host-2']
-            assert cluster.port == 9042
-            assert cluster.protocol_version == 4
-            assert isinstance(cluster.load_balancing_policy, TokenAwarePolicy)
+        with mock.patch.object(Cluster, "__init__") as mock_cluster_ctor:
+            mock_cluster_ctor.return_value = None
+            CassandraHook(cassandra_conn_id='cassandra_test')
+            mock_cluster_ctor.assert_called_once_with(
+                contact_points=['host-1', 'host-2'],
+                port=9042,
+                protocol_version=4,
+                load_balancing_policy=mock.ANY,
+            )
+
+            assert isinstance(mock_cluster_ctor.call_args[1]['load_balancing_policy'], TokenAwarePolicy)
 
     def test_get_lb_policy_with_no_args(self):
         # test LB policies with no args