Merge pull request #298 from alaminopu/feat/add-support-redis-credential-provider

feat: add support for redis credential provider

Author: sibson

.gitignore

@@ -55,6 +55,10 @@ tox-*.xml
 # Django stuff:
 *.log
 
+# ide/editor
+.idea/
+.vscode/
+
 # Sphinx documentation
 docs/_build/
 

redbeat/schedulers.py

@@ -147,15 +147,19 @@ def get_redis(app=None):
             ssl_options = {'ssl_cert_reqs': ssl.CERT_REQUIRED}
             if isinstance(conf.redis_use_ssl, dict):
                 ssl_options.update(conf.redis_use_ssl)
-            connection = StrictRedis.from_url(conf.redis_url, decode_responses=True, **ssl_options)
+            extras = {"decode_responses": True, **ssl_options, **redis_options}
+            connection = StrictRedis.from_url(conf.redis_url, **extras)
         elif conf.redis_url.startswith('redis-cluster'):
             from rediscluster import RedisCluster
 
             if not redis_options.get('startup_nodes'):
-                redis_options = {'startup_nodes': [{"host": "localhost", "port": "30001"}]}
-            connection = RedisCluster(decode_responses=True, **redis_options)
+                startup_nodes_options = {'startup_nodes': [{"host": "localhost", "port": "30001"}]}
+                redis_options.update(startup_nodes_options)
+            redis_options.update({"decode_responses": True})
+            connection = RedisCluster(**redis_options)
         else:
-            connection = StrictRedis.from_url(conf.redis_url, decode_responses=True)
+            redis_options.update({"decode_responses": True})
+            connection = StrictRedis.from_url(conf.redis_url, **redis_options)
 
         if connection:
             _set_redbeat_connect(app, REDBEAT_REDIS_KEY, connection, retry_period)

tests/test_scheduler.py

@@ -9,6 +9,7 @@
 from celery.beat import DEFAULT_MAX_INTERVAL
 from celery.schedules import schedstate, schedule
 from celery.utils.time import maybe_timedelta
+from redis import CredentialProvider
 from redis.exceptions import ConnectionError
 
 from redbeat import RedBeatScheduler
@@ -389,6 +390,42 @@ def test_ssl_connection_scheduler(self):
         assert 'ssl_ca_certs' not in redis_client.connection_pool.connection_kwargs
 
 
+class RedisWithCredentialProvider(AppCase):
+    class UserCredProvider(CredentialProvider):
+        def __init__(self, username, password):
+            self.username = username
+            self.password = password
+
+        def get_credential(self):
+            return self.username, self.password
+
+    config_dict = {
+        'REDBEAT_KEY_PREFIX': 'rb-tests:',
+        'REDBEAT_REDIS_URL': 'rediss://redishost:26379/0',
+        'REDBEAT_REDIS_OPTIONS': {
+            'credential_provider': UserCredProvider("test_user", "test_pass")
+        },
+        'REDBEAT_REDIS_USE_SSL': True,
+    }
+
+    def setup(self):  # celery3
+        self.app.conf.add_defaults(deepcopy(self.config_dict))
+
+    def test_redis_with_credential_provider_scheduler(self):
+        redis_client = get_redis(app=self.app)
+
+        # existing ssl checks
+        assert 'SSLConnection' in str(redis_client.connection_pool)
+        assert redis_client.connection_pool.connection_kwargs['ssl_cert_reqs'] == ssl.CERT_REQUIRED
+
+        # check for credential provider
+        assert 'username' not in redis_client.connection_pool.connection_kwargs
+        assert 'password' not in redis_client.connection_pool.connection_kwargs
+        assert 'credential_provider' in redis_client.connection_pool.connection_kwargs
+        cred_provider = redis_client.connection_pool.connection_kwargs['credential_provider']
+        assert isinstance(cred_provider, CredentialProvider)
+
+
 class RedBeatLockTimeoutDefaultValues(RedBeatCase):
     def test_no_values(self):
         scheduler = RedBeatScheduler(app=self.app)