From e4d7b2309a03326dfd018bc9fb92006e2e8775e9 Mon Sep 17 00:00:00 2001 From: PoAn Yang Date: Tue, 4 Aug 2026 16:33:44 +0900 Subject: [PATCH] Add Redis cluster mode support to RedisHook Signed-off-by: PoAn Yang --- .../testing/integration_tests.rst | 3 +- providers/redis/docs/connections.rst | 46 +++++-- providers/redis/provider.yaml | 14 +++ .../providers/redis/get_provider_info.py | 9 ++ .../airflow/providers/redis/hooks/redis.py | 74 ++++++++++-- .../integration/redis/hooks/test_redis.py | 66 +++++++++++ .../tests/unit/redis/hooks/test_redis.py | 112 ++++++++++++++++++ .../ci/docker-compose/integration-redis.yml | 31 +++++ 8 files changed, 338 insertions(+), 17 deletions(-) diff --git a/contributing-docs/testing/integration_tests.rst b/contributing-docs/testing/integration_tests.rst index 4b0b617d93765..93217348f3bc2 100644 --- a/contributing-docs/testing/integration_tests.rst +++ b/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. | +---------------+-------------------------------------------------------+ diff --git a/providers/redis/docs/connections.rst b/providers/redis/docs/connections.rst index 35fc33e3c317e..ec4286715302a 100644 --- a/providers/redis/docs/connections.rst +++ b/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. diff --git a/providers/redis/provider.yaml b/providers/redis/provider.yaml index 0f006ad75a516..bda0c95e57c64 100644 --- a/providers/redis/provider.yaml +++ b/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: diff --git a/providers/redis/src/airflow/providers/redis/get_provider_info.py b/providers/redis/src/airflow/providers/redis/get_provider_info.py index 6d764d41bc49b..e7dc6b2adf409 100644 --- a/providers/redis/src/airflow/providers/redis/get_provider_info.py +++ b/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", diff --git a/providers/redis/src/airflow/providers/redis/hooks/redis.py b/providers/redis/src/airflow/providers/redis/hooks/redis.py index 7cae107e6ab5f..3510142c21e5c 100644 --- a/providers/redis/src/airflow/providers/redis/hooks/redis.py +++ b/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"), diff --git a/providers/redis/tests/integration/redis/hooks/test_redis.py b/providers/redis/tests/integration/redis/hooks/test_redis.py index eac17ee676edc..5092d756b0c64 100644 --- a/providers/redis/tests/integration/redis/hooks/test_redis.py +++ b/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") diff --git a/providers/redis/tests/unit/redis/hooks/test_redis.py b/providers/redis/tests/unit/redis/hooks/test_redis.py index 1c779c1911ba9..10ce65b8b6876 100644 --- a/providers/redis/tests/unit/redis/hooks/test_redis.py +++ b/providers/redis/tests/unit/redis/hooks/test_redis.py @@ -129,3 +129,115 @@ def test_get_conn_password_stays_none(self): hook = RedisHook(redis_conn_id="redis_default") hook.get_conn() assert hook.password is None + + @mock.patch("airflow.providers.redis.hooks.redis.RedisCluster") + @mock.patch("airflow.providers.redis.hooks.redis.Redis") + @mock.patch("airflow.providers.redis.hooks.redis.RedisHook.get_connection") + def test_get_conn_defaults_to_single_node_client( + self, mock_get_connection, mock_redis, mock_redis_cluster + ): + mock_get_connection.return_value = Connection(host="remote_host", port=1234) + + RedisHook().get_conn() + + mock_redis.assert_called_once() + mock_redis_cluster.assert_not_called() + + @mock.patch("airflow.providers.redis.hooks.redis.Redis") + @mock.patch("airflow.providers.redis.hooks.redis.RedisCluster") + @mock.patch("airflow.providers.redis.hooks.redis.RedisHook.get_connection") + def test_get_conn_cluster_mode(self, mock_get_connection, mock_redis_cluster, mock_redis): + connection = Connection(login="user", password="password", host="node-1", port=6379) + connection.set_extra('{"cluster": true, "ssl": true, "ssl_cert_reqs": "required"}') + mock_get_connection.return_value = connection + + RedisHook().get_conn() + + mock_redis.assert_not_called() + mock_redis_cluster.assert_called_once() + call_kwargs = mock_redis_cluster.call_args[1] + assert call_kwargs["host"] == connection.host + assert call_kwargs["port"] == connection.port + assert call_kwargs["username"] == connection.login + assert call_kwargs["password"] == connection.password + assert call_kwargs["ssl"] is True + assert call_kwargs["ssl_cert_reqs"] == "required" + # RedisCluster raises RedisClusterException when handed a `db` kwarg at all, even `db=0`. + assert "db" not in call_kwargs + + @pytest.mark.parametrize("db", [1, 2]) + @mock.patch("airflow.providers.redis.hooks.redis.RedisCluster") + @mock.patch("airflow.providers.redis.hooks.redis.RedisHook.get_connection") + def test_get_conn_cluster_mode_rejects_non_zero_db(self, mock_get_connection, mock_redis_cluster, db): + connection = Connection(host="node-1", port=6379) + connection.set_extra(f'{{"cluster": true, "db": {db}}}') + mock_get_connection.return_value = connection + + with pytest.raises(ValueError, match="only supports database 0"): + RedisHook().get_conn() + + mock_redis_cluster.assert_not_called() + + @mock.patch("airflow.providers.redis.hooks.redis.RedisCluster") + @mock.patch("airflow.providers.redis.hooks.redis.RedisHook.get_connection") + def test_get_conn_cluster_mode_accepts_db_zero(self, mock_get_connection, mock_redis_cluster): + """`db` defaults to 0 in the connection form, so that value must not be treated as a conflict.""" + connection = Connection(host="node-1", port=6379) + connection.set_extra('{"cluster": true, "db": 0}') + mock_get_connection.return_value = connection + + RedisHook().get_conn() + + mock_redis_cluster.assert_called_once() + assert "db" not in mock_redis_cluster.call_args[1] + + @mock.patch("airflow.providers.redis.hooks.redis.RedisCluster") + @mock.patch("airflow.providers.redis.hooks.redis.RedisHook.get_connection") + def test_get_conn_cluster_mode_passes_startup_nodes(self, mock_get_connection, mock_redis_cluster): + connection = Connection(host="node-1", port=6379) + connection.set_extra('{"cluster": true, "startup_nodes": ["node-2:6379", "node-3:6380"]}') + mock_get_connection.return_value = connection + + RedisHook().get_conn() + + startup_nodes = mock_redis_cluster.call_args[1]["startup_nodes"] + assert [(node.host, node.port) for node in startup_nodes] == [("node-2", 6379), ("node-3", 6380)] + + @pytest.mark.parametrize( + ("raw_nodes", "expected"), + [ + pytest.param(["node-2:6379", "node-3:6380"], [("node-2", 6379), ("node-3", 6380)], id="list"), + pytest.param("node-2:6379,node-3:6380", [("node-2", 6379), ("node-3", 6380)], id="csv-string"), + pytest.param("node-2 , node-3:6380", [("node-2", 6379), ("node-3", 6380)], id="default-port"), + pytest.param(None, [], id="unset"), + pytest.param([], [], id="empty"), + ], + ) + def test_build_startup_nodes(self, raw_nodes, expected): + nodes = RedisHook._build_startup_nodes(raw_nodes) + + assert [(node.host, node.port) for node in nodes] == expected + + @pytest.mark.parametrize( + ("raw_nodes", "expected_error"), + [ + pytest.param("node-2:not-a-port", "Invalid port", id="non-numeric-port"), + pytest.param(":6379", "Missing host", id="missing-host"), + ], + ) + def test_build_startup_nodes_rejects_invalid_entries(self, raw_nodes, expected_error): + with pytest.raises(ValueError, match=expected_error): + RedisHook._build_startup_nodes(raw_nodes) + + @mock.patch("airflow.providers.redis.hooks.redis.RedisCluster") + @mock.patch("airflow.providers.redis.hooks.redis.RedisHook.get_connection") + @mock.patch("airflow.providers.redis.hooks.redis.DriverInfo", None) + @mock.patch("airflow.providers.redis.hooks.redis._SUPPORTS_LIB_NAME", True) + def test_cluster_mode_passes_client_identification(self, mock_get_connection, mock_redis_cluster): + connection = Connection(host="node-1", port=6379) + connection.set_extra('{"cluster": true}') + mock_get_connection.return_value = connection + + RedisHook().get_conn() + + assert "apache-airflow-providers-redis" in mock_redis_cluster.call_args[1]["lib_name"] diff --git a/scripts/ci/docker-compose/integration-redis.yml b/scripts/ci/docker-compose/integration-redis.yml index 488c2c0cb7c2a..cb07ad3b18c59 100644 --- a/scripts/ci/docker-compose/integration-redis.yml +++ b/scripts/ci/docker-compose/integration-redis.yml @@ -31,11 +31,42 @@ services: start_period: 30s retries: 50 restart: "on-failure" + redis-cluster: + image: redis:7-alpine + labels: + breeze.description: "Integration required for Redis cluster mode tests." + # All three masters share one network namespace so the cluster can be formed without + # inter-container discovery. The cluster is created against the container's own IP rather + # than 127.0.0.1: nodes record the address they are created with and hand it back in + # CLUSTER SLOTS, and a client in another container cannot follow a 127.0.0.1 redirect. + command: + - sh + - -c + - | + set -e + ip=$$(hostname -i | awk '{print $$1}') + for port in 7001 7002 7003; do + redis-server --port $$port --bind 0.0.0.0 --protected-mode no \ + --cluster-enabled yes --cluster-config-file nodes-$$port.conf \ + --cluster-node-timeout 5000 --appendonly no --daemonize yes + done + until redis-cli -p 7003 ping > /dev/null 2>&1; do sleep 1; done + redis-cli --cluster create $$ip:7001 $$ip:7002 $$ip:7003 --cluster-yes + sleep infinity + healthcheck: + test: ["CMD-SHELL", "redis-cli -p 7001 cluster info | grep -q cluster_state:ok"] + interval: 5s + timeout: 30s + start_period: 30s + retries: 50 + restart: "on-failure" airflow: environment: - INTEGRATION_REDIS=true depends_on: redis: condition: service_healthy + redis-cluster: + condition: service_healthy volumes: redis-db-volume: