Add Redis cluster mode support to RedisHook

contributing-docs/testing/integration_tests.rst

@@ -89,7 +89,8 @@ The following integrations are available.
 +---------------+-------------------------------------------------------+
 | qdrant        | Integration required for Qdrant tests.                |
 +---------------+-------------------------------------------------------+
-| redis         | Integration required for Redis tests.                 |
+| redis         | * Integration required for Redis tests.               |
+|               | * Integration required for Redis cluster mode tests.  |
 +---------------+-------------------------------------------------------+
 | statsd        | Integration required for Statsd hooks.                |
 +---------------+-------------------------------------------------------+

providers/redis/docs/connections.rst

@@ -20,7 +20,8 @@
 Redis Connection
 ================
 
-The Redis connection type enables connection to Redis cluster.
+The Redis connection type enables connection to a Redis deployment, either a standalone
+server or one running in cluster mode.
 
 Default Connection IDs
 ----------------------
@@ -31,22 +32,30 @@ parameter as ``redis_default`` by default.
 Configuring the Connection
 --------------------------
 Host
-    The host of the Redis cluster.
+    The host of the Redis server.
 
 Port
-    Specify the port to use for connecting the Redis cluster (Default is ``6379``).
+    Specify the port to use for connecting the Redis server (Default is ``6379``).
 
 Login
-    The user that will be used for authentication against the Redis cluster (only applicable in Redis 6.0 and above).
+    The user that will be used for authentication against the Redis server (only applicable in Redis 6.0 and above).
 
 Password
-    The password of the user that will be used for authentication against the Redis cluster.
+    The password of the user that will be used for authentication against the Redis server.
 
 DB
-    The DB number to use in the Redis cluster (Default is ``0``).
+    The DB number to use in the Redis server (Default is ``0``). Not supported in cluster mode.
+
+Enable cluster mode
+    Whether to connect with a cluster-aware client that follows ``MOVED``/``ASK`` redirects
+    (Default is ``False``). See :ref:`redis-cluster-mode` below.
+
+Cluster startup nodes
+    Extra bootstrap nodes as a comma-separated ``host:port`` list. The port may be omitted and
+    defaults to ``6379``. Used in cluster mode only (Default is ``None``).
 
 Enable SSL
-    Whether to enable SSL connection to the Redis cluster (Default is ``False``).
+    Whether to enable SSL connection to the Redis server (Default is ``False``).
 
 SSL verify mode
     Whether to try to verify other peers' certificates and how to behave if verification fails.
@@ -64,3 +73,26 @@ Certificate path
 
 Enable hostname check
     If set, match the hostname during the SSL handshake (Default is ``False``).
+
+.. _redis-cluster-mode:
+
+Cluster mode
+------------
+
+Redis Cluster spreads the keyspace over 16384 hash slots owned by different masters, and expects
+the client to route each command to the node owning that key's slot. A standalone client does not
+do this: when it asks a node for a key that node does not serve, the node answers ``MOVED`` and
+the standalone client fails.
+
+Enable cluster mode to use a cluster-aware client that follows those redirects:
+
+.. code-block:: json
+
+    {
+      "cluster": true,
+      "startup_nodes": ["node-2:6379", "node-3:6379"]
+    }
+
+The client discovers the full topology from the first node it reaches. ``startup_nodes`` matters for
+bootstrap resilience: every task builds its own connection, so with a single seed node one unreachable
+node breaks every task.

providers/redis/provider.yaml

@@ -116,6 +116,20 @@ connection-types:
             - integer
             - 'null'
           default: 0
+      cluster:
+        label: Enable cluster mode
+        schema:
+          type:
+            - boolean
+            - 'null'
+          default: false
+      startup_nodes:
+        label: Cluster startup nodes
+        description: "Comma-separated extra bootstrap nodes as host:port. Cluster mode only."
+        schema:
+          type:
+            - string
+            - 'null'
       ssl:
         label: Enable SSL
         schema:

providers/redis/src/airflow/providers/redis/get_provider_info.py

@@ -65,6 +65,15 @@ def get_provider_info():
                 "ui-field-behaviour": {"hidden-fields": ["schema", "extra"], "relabeling": {}},
                 "conn-fields": {
                     "db": {"label": "DB", "schema": {"type": ["integer", "null"], "default": 0}},
+                    "cluster": {
+                        "label": "Enable cluster mode",
+                        "schema": {"type": ["boolean", "null"], "default": False},
+                    },
+                    "startup_nodes": {
+                        "label": "Cluster startup nodes",
+                        "description": "Comma-separated extra bootstrap nodes as host:port. Cluster mode only.",
+                        "schema": {"type": ["string", "null"]},
+                    },
                     "ssl": {"label": "Enable SSL", "schema": {"type": ["boolean", "null"], "default": False}},
                     "ssl_cert_reqs": {
                         "label": "SSL verify mode",

providers/redis/src/airflow/providers/redis/hooks/redis.py

@@ -24,6 +24,7 @@
 
 import redis
 from redis import Redis
+from redis.cluster import ClusterNode, RedisCluster
 
 from airflow.providers.common.compat.sdk import BaseHook
 from airflow.providers.redis import __version__ as provider_version
@@ -32,6 +33,7 @@
 
 DEFAULT_SSL_CERT_REQS = "required"
 ALLOWED_SSL_CERT_REQS = [DEFAULT_SSL_CERT_REQS, "optional", "none"]
+DEFAULT_REDIS_PORT = 6379
 
 # Check at module import time what Redis client identification features are supported
 _REDIS_PARAMS = inspect.signature(Redis.__init__).parameters
@@ -45,6 +47,11 @@ class RedisHook(BaseHook):
     You can set your db in the extra field of your connection as ``{"db": 3}``.
     Also you can set ssl parameters as:
     ``{"ssl": true, "ssl_cert_reqs": "require", "ssl_certfile": "/path/to/cert.pem", etc}``.
+
+    To talk to a Redis deployment running in cluster mode, set ``{"cluster": true}``. Additional
+    bootstrap nodes may be listed as ``{"startup_nodes": ["node-2:6379", "node-3:6379"]}`` so that
+    a single unreachable node does not leave the whole connection unusable. Cluster mode only
+    supports database 0, so ``db`` must be left unset or 0.
     """
 
     conn_name_attr = "redis_conn_id"
@@ -67,6 +74,7 @@ def __init__(self, redis_conn_id: str = default_conn_name, **kwargs) -> None:
         self.username = kwargs.get("username", None)
         self.password = kwargs.get("password", None)
         self.db = kwargs.get("db", None)
+        self.cluster = kwargs.get("cluster", False)
 
     def get_conn(self):
         """Return a Redis connection."""
@@ -76,6 +84,14 @@ def get_conn(self):
         self.username = conn.login
         self.password = None if str(conn.password).lower() in ["none", "false", ""] else conn.password
         self.db = conn.extra_dejson.get("db")
+        self.cluster = conn.extra_dejson.get("cluster", False)
+
+        # https://redis.io/docs/latest/operate/oss_and_stack/reference/cluster-spec/#implemented-subset
+        if self.cluster and self.db not in (None, 0):
+            raise ValueError(
+                f"Redis connection {self.redis_conn_id!r} sets `db` to {self.db!r}, but Redis in cluster "
+                "mode only supports database 0. Remove `db` from the connection extra."
+            )
 
         # check for ssl parameters in conn.extra
         ssl_arg_names = [
@@ -111,18 +127,50 @@ def get_conn(self):
                     "lib_name": f"redis-py(apache-airflow-providers-redis_v{provider_version})",
                 }
 
-            self.redis = Redis(
-                host=self.host,
-                port=self.port,
-                username=self.username,
-                password=self.password,
-                db=self.db,
-                **ssl_args,
-                **driver_info_options,
-            )
+            if self.cluster:
+                self.redis = RedisCluster(
+                    host=self.host,
+                    port=self.port,
+                    startup_nodes=self._build_startup_nodes(conn.extra_dejson.get("startup_nodes")),
+                    username=self.username,
+                    password=self.password,
+                    **ssl_args,
+                    **driver_info_options,
+                )
+            else:
+                self.redis = Redis(
+                    host=self.host,
+                    port=self.port,
+                    username=self.username,
+                    password=self.password,
+                    db=self.db,
+                    **ssl_args,
+                    **driver_info_options,
+                )
 
         return self.redis
 
+    @staticmethod
+    def _build_startup_nodes(raw_nodes: Any) -> list[ClusterNode]:
+        """Build redis-py cluster nodes from the ``startup_nodes`` extra, given as ``host`` or ``host:port``."""
+        if not raw_nodes:
+            return []
+
+        entries = raw_nodes.split(",") if isinstance(raw_nodes, str) else raw_nodes
+        nodes = []
+        for entry in entries:
+            host, _, port = str(entry).strip().partition(":")
+            if not host:
+                raise ValueError(f"Missing host in `startup_nodes` entry {entry!r}; expected `host:port`.")
+            try:
+                parsed_port = int(port) if port else DEFAULT_REDIS_PORT
+            except ValueError:
+                raise ValueError(
+                    f"Invalid port in `startup_nodes` entry {entry!r}; expected `host:port`."
+                ) from None
+            nodes.append(ClusterNode(host, parsed_port))
+        return nodes
+
     @classmethod
     def get_ui_field_behaviour(cls) -> dict[str, Any]:
         """Return custom UI field behaviour for Redis connection."""
@@ -141,6 +189,14 @@ def get_connection_form_widgets(cls) -> dict[str, Any]:
 
         return {
             "db": IntegerField(lazy_gettext("DB"), widget=BS3TextFieldWidget(), default=0),
+            "cluster": BooleanField(lazy_gettext("Enable cluster mode"), default=False),
+            "startup_nodes": StringField(
+                lazy_gettext("Cluster startup nodes"),
+                widget=BS3TextFieldWidget(),
+                validators=[Optional()],
+                description="Comma-separated extra bootstrap nodes as host:port. Cluster mode only.",
+                default=None,
+            ),
             "ssl": BooleanField(lazy_gettext("Enable SSL"), default=False),
             "ssl_cert_reqs": StringField(
                 lazy_gettext("SSL verify mode"),

providers/redis/tests/integration/redis/hooks/test_redis.py

@@ -17,10 +17,19 @@
 
 from __future__ import annotations
 
+import json
+
 import pytest
+from redis.cluster import RedisCluster
+from redis.exceptions import MovedError
 
 from airflow.providers.redis.hooks.redis import RedisHook
 
+CLUSTER_HOST = "redis-cluster"
+CLUSTER_SEED_PORT = 7001
+# Nothing listens here; used to prove `startup_nodes` is what establishes the connection.
+CLUSTER_DEAD_PORT = 7009
+
 
 @pytest.mark.integration("redis")
 class TestRedisHook:
@@ -37,3 +46,60 @@ def test_real_get_and_set(self):
         assert redis.set("test_key", "test_value"), "Connection to Redis with SET works."
         assert redis.get("test_key") == b"test_value", "Connection to Redis with GET works."
         assert redis.delete("test_key") == 1, "Connection to Redis with DELETE works."
+
+
+@pytest.mark.integration("redis")
+class TestRedisHookClusterMode:
+    @pytest.fixture(autouse=True)
+    def cluster_connections(self, monkeypatch):
+        seed = {"conn_type": "redis", "host": CLUSTER_HOST, "port": CLUSTER_SEED_PORT}
+        monkeypatch.setenv("AIRFLOW_CONN_REDIS_STANDALONE_TEST", json.dumps(seed))
+        monkeypatch.setenv(
+            "AIRFLOW_CONN_REDIS_CLUSTER_TEST", json.dumps({**seed, "extra": {"cluster": True}})
+        )
+        monkeypatch.setenv(
+            "AIRFLOW_CONN_REDIS_CLUSTER_SEEDS_TEST",
+            json.dumps(
+                {
+                    **seed,
+                    "port": CLUSTER_DEAD_PORT,
+                    "extra": {
+                        "cluster": True,
+                        "startup_nodes": [f"{CLUSTER_HOST}:7002", f"{CLUSTER_HOST}:7003"],
+                    },
+                }
+            ),
+        )
+
+    def test_cluster_mode_follows_moved_redirect(self):
+        """Both connections point at the same seed node; only the cluster client can reach the key."""
+        cluster = RedisHook(redis_conn_id="redis_cluster_test").get_conn()
+        assert isinstance(cluster, RedisCluster)
+
+        remote_keys = [
+            key
+            for key in (f"cluster_key_{i}" for i in range(20))
+            if cluster.get_node_from_key(key).port != CLUSTER_SEED_PORT
+        ]
+        assert remote_keys, "expected at least one key owned by a node other than the seed node"
+        remote_key = remote_keys[0]
+
+        standalone = RedisHook(redis_conn_id="redis_standalone_test").get_conn()
+        with pytest.raises(MovedError):
+            standalone.set(remote_key, "value")
+
+        try:
+            assert cluster.set(remote_key, "value")
+            assert cluster.get(remote_key) == b"value"
+        finally:
+            cluster.delete(remote_key)
+
+    def test_startup_nodes_connect_when_the_seed_node_is_unreachable(self):
+        """The connection's own host/port is dead, so only `startup_nodes` can bootstrap it."""
+        conn = RedisHook(redis_conn_id="redis_cluster_seeds_test").get_conn()
+
+        try:
+            assert conn.set("cluster_startup_nodes_key", "value")
+            assert conn.get("cluster_startup_nodes_key") == b"value"
+        finally:
+            conn.delete("cluster_startup_nodes_key")