Browse Source

Add auth options in cassandra backend

David Harrigan 9 years ago
parent
commit
afbd2330ed
2 changed files with 12 additions and 3 deletions
  1. 11 2
      celery/backends/cassandra.py
  2. 1 1
      celery/tests/backends/test_cassandra.py

+ 11 - 2
celery/backends/cassandra.py

@@ -11,6 +11,7 @@ from __future__ import absolute_import
 import sys
 try:  # pragma: no cover
     import cassandra
+    import cassandra.auth
     import cassandra.cluster
 except ImportError:  # pragma: no cover
     cassandra = None   # noqa
@@ -121,6 +122,13 @@ class CassandraBackend(BaseBackend):
             cassandra.ConsistencyLevel.LOCAL_QUORUM,
         )
 
+        self.auth_provider = None
+        auth_provider = conf.get('cassandra_auth_provider', None)
+        auth_kwargs = conf.get('cassandra_auth_kwargs', None)
+        if auth_provider and auth_kwargs:
+            auth_provider_class = getattr(cassandra.auth, auth_provider)
+            self.auth_provider = auth_provider_class(**auth_kwargs)
+
         self._connection = None
         self._session = None
         self._write_stmt = None
@@ -142,8 +150,9 @@ class CassandraBackend(BaseBackend):
         """
         if self._connection is None:
             try:
-                self._connection = cassandra.cluster.Cluster(self.servers,
-                                                             port=self.port)
+                self._connection = cassandra.cluster.Cluster(
+                    self.servers, port=self.port,
+                    auth_provider=self.auth_provider)
                 self._session = self._connection.connect(self.keyspace)
 
                 # We are forced to do concatenation below, as formatting would

+ 1 - 1
celery/tests/backends/test_cassandra.py

@@ -9,7 +9,7 @@ from celery.tests.case import (
     AppCase, Mock, mock_module, depends_on_current_app
 )
 
-CASSANDRA_MODULES = ['cassandra', 'cassandra.cluster']
+CASSANDRA_MODULES = ['cassandra', 'cassandra.auth', 'cassandra.cluster']
 
 
 class Object(object):