django-redis cluster client

cluster_client.py

"""
A sketch of a django-redis integration with the Redis cluster features recently added to
https://github.com/redis/redis-py from https://github.com/Grokzen/redis-py-cluster

For non-clustered Redis, django-redis's ConnectionFactory does some management of
connection pools shared between client instances.

The cluster client in redis-py doesn't accept a connection pool from outside, they're
managed internally. To support that, we won't be caching connection pools and passing
them into clients, we will instead be caching client instances.
"""

import threading
from copy import deepcopy

from django.core.exceptions import ImproperlyConfigured

from django_redis.client.default import DefaultClient
from django_redis.pool import ConnectionFactory

from redis.cluster import RedisCluster


class DjangoRedisClusterClient(DefaultClient):
    """A django-redis client compatible with redis.cluster.RedisCluster

    We don't do much different here, except for using our own
    ClusterConnectionFactory (the base class would instead use the value of the
    DJANGO_REDIS_CONNECTION_FACTORY setting, but we don't care about that
    setting here.)
    """

    def __init__(self, server, params, backend) -> None:
        super().__init__(server, params, backend)
        self.connection_factory = ClusterConnectionFactory(options=self._options)


class ClusterConnectionFactory(ConnectionFactory):
    """A connection factory compatible with redis.cluster.RedisCluster

    The cluster client manages connection pools internally, so we don't want to
    do it at this level like the base ConnectionFactory does.
    """

    # A global cache of URL->client so that within a process, we will reuse a
    # single client, and therefore a single set of connection pools.
    _clients = {}
    _clients_lock = threading.Lock()

    def __init__(self, options):
        # set appropriate default, but allow overriding client class
        options.setdefault("REDIS_CLIENT_CLASS", "redis.cluster.RedisCluster")
        super().__init__(options)

    def connect(self, url: str) -> RedisCluster:
        """Given a connection url, return a client instance.

        Prefer to return from our cache but if we don't yet have one build it
        to populate the cache.
        """
        if url not in self._clients:
            with self._clients_lock:
                if url not in self._clients:
                    self._clients[url] = self._connect(url)
        return self._clients[url]

    def _connect(self, url: str) -> RedisCluster:
        """
        Given a connection url, return a new client instance.

        Basic django-redis ConnectionFactory manages a cache of connection
        pools and builds a fresh client each time. because the cluster client
        manages its own connection pools, we will instead merge the
        "connection" and "client" kwargs and throw them all at the client to
        sort out.

        If we find conflicting client and connection kwargs, we'll raise an
        error.
        """
        # Get connection and client kwargs...
        connection_params = self.make_connection_params(url)
        client_cls_kwargs = deepcopy(self.redis_client_cls_kwargs)

        # ... and smash 'em together (crashing if there's conflicts)...
        for key, value in connection_params.items():
            if key in client_cls_kwargs:
                raise ImproperlyConfigured(f"Found '{key}' in both the connection and the client kwargs")
            client_cls_kwargs[key] = value

        # ... and then build and return the client
        return self.redis_client_cls(**client_cls_kwargs)

    def disconnect(self, connection: RedisCluster):
        connection.disconnect_connection_pools()

example_settings.py

# Example Django CACHES setting to use the django-redis ClusterClient proposed here. By
# doing that, you will get the redis.cluster.RedisCluster actually talking to your
# Redis

CACHES = {
    'default': {
        'BACKEND': 'django_redis.cache.RedisCache',
        'LOCATION': '...',
        'OPTIONS': {
            'CLIENT_CLASS': 'cluster_client.DjangoRedisClusterClient',
        },
    },
}

# When we use that DjangoRedisClusterClient, it will impose a default value of
# 'redis.cluster.RedisCluster' for the ['OPTIONS']['REDIS_CLIENT_CLASS'] config,
# you won't have to specify that.
#
# It will also ignore the DJANGO_REDIS_CONNECTION_FACTORY setting and use the
# cluster-specific connection factory here.