ref: break the generic / inheritance of redis ClusterManager (#68018)

this allow us to properly type redis with respect to `bytes` / `str`
(currently it is defaulting to `Any`!)

<!-- Describe your PR here. -->

src/sentry/digests/backends/redis.py

@@ -4,7 +4,6 @@ from collections.abc import Iterable
 from contextlib import contextmanager
 from typing import Any
 
-import rb
 from rb.clients import LocalClient
 from redis.exceptions import ResponseError
 
@@ -73,7 +72,6 @@ class RedisBackend(Backend):
 
     def __init__(self, **options: Any) -> None:
         cluster, options = get_cluster_from_options("SENTRY_DIGESTS_OPTIONS", options)
-        assert isinstance(cluster, rb.Cluster)
         self.cluster = cluster
         self.locks = LockManager(RedisLockBackend(self.cluster))
 

src/sentry/tsdb/redis.py

@@ -119,10 +119,7 @@ class RedisTSDB(BaseTSDB):
 
     def __init__(self, prefix: str = "ts:", vnodes: int = 64, **options: Any):
         cluster, options = get_cluster_from_options("SENTRY_TSDB_OPTIONS", options)
-        if is_instance_rb_cluster(cluster, False):
-            self.cluster = cluster
-        else:
-            raise AssertionError("unreachable")
+        self.cluster = cluster
         self.prefix = prefix
         self.vnodes = vnodes
         self.enable_frequency_sketches = options.pop("enable_frequency_sketches", False)

src/sentry/utils/redis.py

@@ -1,11 +1,10 @@
 from __future__ import annotations
 
-import functools
 import importlib.resources
 import logging
 from copy import deepcopy
 from threading import Lock
-from typing import Any, Generic, TypeGuard, TypeVar, overload
+from typing import Any, TypeGuard
 
 import rb
 from django.utils.functional import SimpleLazyObject
@@ -52,21 +51,17 @@ def _shared_pool(**opts: Any) -> ConnectionPool:
         return pool
 
 
-_make_rb_cluster = functools.partial(rb.Cluster, pool_cls=_shared_pool)
+class RBClusterManager:
+    def __init__(self, options_manager: OptionsManager) -> None:
+        self._clusters: dict[str, rb.Cluster] = {}
+        self._options_manager = options_manager
 
-
-class _RBCluster:
-    def supports(self, config: dict[str, Any]) -> bool:
-        return not config.get("is_redis_cluster", False)
-
-    def factory(
+    def _factory(
         self,
-        decode_responses: bool | None = None,
+        *,
         hosts: list[dict[int, Any]] | dict[int, Any] | None = None,
         **config: Any,
     ) -> rb.Cluster:
-        if not decode_responses:
-            raise NotImplementedError("decode_responses=False mode is not implemented for `rb`")
         if not hosts:
             hosts = []
         # rb expects a dict of { host, port } dicts where the key is the host
@@ -78,10 +73,26 @@ class _RBCluster:
         pool_options = {**_REDIS_DEFAULT_CLIENT_ARGS, **pool_options}
         config["pool_options"] = pool_options
 
-        return _make_rb_cluster(**config)
+        return rb.Cluster(**config, pool_cls=_shared_pool)
 
-    def __str__(self) -> str:
-        return "Redis Blaster Cluster"
+    def get(self, key: str, *, decode_responses: bool = True) -> rb.Cluster:
+        if not decode_responses:
+            raise NotImplementedError("rb does not support decode_responses")
+
+        try:
+            return self._clusters[key]
+        except KeyError:
+            pass
+
+        cfg = self._options_manager.get("redis.clusters", {}).get(key)
+        if cfg is None:
+            raise KeyError(f"Invalid cluster name: {key}")
+
+        if cfg.get("is_redis_cluster", False):
+            raise KeyError("Invalid cluster type, expected rb cluster")
+
+        ret = self._clusters[key] = self._factory(**cfg)
+        return ret
 
 
 class _RedisCluster:
@@ -101,7 +112,7 @@ class _RedisCluster:
         hosts: list[dict[Any, Any]] | dict[Any, Any] | None = None,
         client_args: dict[str, Any] | None = None,
         **config: Any,
-    ) -> SimpleLazyObject:
+    ) -> RetryingRedisCluster | FailoverRedis:
         # StrictRedisCluster expects a list of { host, port } dicts. Coerce the
         # configuration into the correct format if necessary.
         if not hosts:
@@ -142,72 +153,57 @@ class _RedisCluster:
                 host["decode_responses"] = decode_responses
                 return FailoverRedis(**host, **client_args)
 
-        return SimpleLazyObject(cluster_factory)
+        # losing some type safety: SimpleLazyObject acts like the underlying type
+        return SimpleLazyObject(cluster_factory)  # type: ignore[return-value]
 
     def __str__(self) -> str:
         return "Redis Cluster"
 
 
-TCluster = TypeVar("TCluster", rb.Cluster, RedisCluster | StrictRedis)
-
-
-class ClusterManager(Generic[TCluster]):
-    @overload
-    def __init__(self: ClusterManager[rb.Cluster], options_manager: OptionsManager) -> None:
-        ...
-
-    @overload
-    def __init__(
-        self: ClusterManager[RedisCluster | StrictRedis],
-        options_manager: OptionsManager,
-        cluster_type: type[Any],
-    ) -> None:
-        ...
-
-    def __init__(self, options_manager: OptionsManager, cluster_type: type[Any] = _RBCluster):
-        self.__clusters: dict[tuple[str, bool], TCluster] = {}
+class RedisClusterManager:
+    def __init__(self, options_manager: OptionsManager) -> None:
+        self.__clusters: dict[tuple[str, bool], RedisCluster | StrictRedis] = {}
         self.__options_manager = options_manager
-        self.__cluster_type = cluster_type()
+        self.__cluster_type = _RedisCluster()
 
-    def get(self, key: str, *, decode_responses: bool = True) -> TCluster:
+    def get(self, key: str, *, decode_responses: bool = True) -> RedisCluster | StrictRedis:
         cache_key = (key, decode_responses)
         try:
             return self.__clusters[cache_key]
         except KeyError:
-            # Do not access attributes of the `cluster` object to prevent
-            # setup/init of lazy objects. The _RedisCluster type will try to
-            # connect to the cluster during initialization.
-
-            # TODO: This would probably be safer with a lock, but I'm not sure
-            # that it's necessary.
-            configuration = self.__options_manager.get("redis.clusters", {}).get(key)
-            if configuration is None:
-                raise KeyError(f"Invalid cluster name: {key}")
-
-            if not self.__cluster_type.supports(configuration):
-                raise KeyError(f"Invalid cluster type, expected: {self.__cluster_type}")
-
-            ret = self.__clusters[cache_key] = self.__cluster_type.factory(
-                **configuration,
-                decode_responses=decode_responses,
-            )
-            return ret
+            pass
+
+        # Do not access attributes of the `cluster` object to prevent
+        # setup/init of lazy objects. The _RedisCluster type will try to
+        # connect to the cluster during initialization.
+
+        # TODO: This would probably be safer with a lock, but I'm not sure
+        # that it's necessary.
+        cfg = self.__options_manager.get("redis.clusters", {}).get(key)
+        if cfg is None:
+            raise KeyError(f"Invalid cluster name: {key}")
+
+        if not self.__cluster_type.supports(cfg):
+            raise KeyError(f"Invalid cluster type, expected: {self.__cluster_type}")
+
+        ret = self.__clusters[cache_key] = self.__cluster_type.factory(
+            **cfg, decode_responses=decode_responses
+        )
+        return ret
 
 
 # TODO(epurkhiser): When migration of all rb cluster to true redis clusters has
 # completed, remove the rb ``clusters`` module variable and rename
 # redis_clusters to clusters.
-clusters: ClusterManager[rb.Cluster] = ClusterManager(options.default_manager)
-redis_clusters: ClusterManager[RedisCluster | StrictRedis] = ClusterManager(
-    options.default_manager, _RedisCluster
-)
+clusters = RBClusterManager(options.default_manager)
+redis_clusters = RedisClusterManager(options.default_manager)
 
 
 def get_cluster_from_options(
     setting: str,
     options: dict[str, Any],
-    cluster_manager: ClusterManager = clusters,
-) -> tuple[rb.Cluster | RedisCluster | StrictRedis, dict[str, Any]]:
+    cluster_manager: RBClusterManager = clusters,
+) -> tuple[rb.Cluster, dict[str, Any]]:
     cluster_option_name = "cluster"
     default_cluster_name = "default"
     cluster_constructor_option_names = frozenset(("hosts",))

tests/sentry/processing/realtime_metrics/test_redis.py

@@ -23,10 +23,7 @@ def config() -> dict[str, Any]:
 
 @pytest.fixture
 def redis_cluster(config: dict[str, Any]) -> StrictRedis:
-    cluster, options = redis.get_cluster_from_options(
-        "TEST_CLUSTER", config, cluster_manager=redis.redis_clusters
-    )
-    return cluster
+    return redis.redis_clusters.get("default")
 
 
 @pytest.fixture

tests/sentry/tsdb/test_redis.py

@@ -2,7 +2,6 @@ from contextlib import contextmanager
 from datetime import datetime, timedelta, timezone
 
 import pytest
-import rb
 
 from sentry.testutils.cases import TestCase
 from sentry.testutils.helpers.options import override_options
@@ -46,12 +45,10 @@ class RedisTSDBTest(TestCase):
             cluster="tsdb",
         )
 
-        assert isinstance(self.db.cluster, rb.Cluster)
         # the point of this test is to demonstrate behaviour with a multi-host cluster
         assert len(self.db.cluster.hosts) == 3
 
     def tearDown(self):
-        assert isinstance(self.db.cluster, rb.Cluster)
         with self.db.cluster.all() as client:
             client.flushdb()
 
@@ -594,7 +591,6 @@ class RedisTSDBTest(TestCase):
         ) == {"organization:1": [], "organization:2": []}
 
     def test_frequency_table_import_export_no_estimators(self):
-        assert isinstance(self.db.cluster, rb.Cluster)
         client = self.db.cluster.get_local_client_for_key("key")
 
         parameters = [64, 5, 10]
@@ -645,7 +641,6 @@ class RedisTSDBTest(TestCase):
         assert client.exists("1:e")
 
     def test_frequency_table_import_export_both_estimators(self):
-        assert isinstance(self.db.cluster, rb.Cluster)
         client = self.db.cluster.get_local_client_for_key("key")
 
         parameters = [64, 5, 5]
@@ -708,7 +703,6 @@ class RedisTSDBTest(TestCase):
         ]
 
     def test_frequency_table_import_export_source_estimators(self):
-        assert isinstance(self.db.cluster, rb.Cluster)
         client = self.db.cluster.get_local_client_for_key("key")
 
         parameters = [64, 5, 5]
@@ -767,7 +761,6 @@ class RedisTSDBTest(TestCase):
         ]
 
     def test_frequency_table_import_export_destination_estimators(self):
-        assert isinstance(self.db.cluster, rb.Cluster)
         client = self.db.cluster.get_local_client_for_key("key")
 
         parameters = [64, 5, 5]

tests/sentry/utils/test_redis.py

@@ -1,27 +1,21 @@
 from __future__ import annotations
 
-import logging
 from unittest import TestCase, mock
 
 import pytest
 import rb
-from django.utils.functional import SimpleLazyObject
 from sentry_redis_tools.failover_redis import FailoverRedis
 
 from sentry.exceptions import InvalidConfiguration
 from sentry.utils import imports
 from sentry.utils.redis import (
-    ClusterManager,
-    _RedisCluster,
+    RBClusterManager,
+    RedisClusterManager,
     _shared_pool,
     get_cluster_from_options,
-    logger,
 )
 from sentry.utils.warnings import DeprecatedSettingWarning
 
-# Silence connection warnings
-logger.setLevel(logging.ERROR)
-
 
 def _options_manager():
     return {
@@ -38,7 +32,7 @@ class ClusterManagerTestCase(TestCase):
         imports._cache.clear()
 
     def test_get(self):
-        manager = ClusterManager(_options_manager())
+        manager = RBClusterManager(_options_manager())
         assert manager.get("foo") is manager.get("foo")
         assert manager.get("foo") is not manager.get("bar")
         assert manager.get("foo").pool_cls is _shared_pool
@@ -47,7 +41,7 @@ class ClusterManagerTestCase(TestCase):
 
     @mock.patch("sentry.utils.redis.RetryingRedisCluster")
     def test_specific_cluster(self, RetryingRedisCluster):
-        manager = ClusterManager(_options_manager(), cluster_type=_RedisCluster)
+        manager = RedisClusterManager(_options_manager())
 
         # We wrap the cluster in a Simple Lazy Object, force creation of the
         # object to verify it's correct.
@@ -62,26 +56,19 @@ class ClusterManagerTestCase(TestCase):
         with pytest.raises(KeyError):
             manager.get("bar")
 
-    def test_multiple_retrieval_do_not_setup_lazy_object(self):
-        class TestClusterType:
-            def supports(self, config):
-                return True
-
-            def factory(self, **config):
-                def setupfunc():
-                    assert False, "setupfunc should not be called"
-
-                return SimpleLazyObject(setupfunc)
+    @mock.patch("sentry.utils.redis.RetryingRedisCluster")
+    def test_multiple_retrieval_do_not_setup_lazy_object(self, RetryingRedisCluster):
+        RetryingRedisCluster.side_effect = AssertionError("should not be called")
 
-        manager = ClusterManager(_options_manager(), cluster_type=TestClusterType)
-        manager.get("foo")
+        manager = RedisClusterManager(_options_manager())
+        manager.get("baz")
         # repeated retrieval should not trigger call to setupfunc
-        manager.get("foo")
+        manager.get("baz")
 
 
 def test_get_cluster_from_options_cluster_provided():
     backend = mock.sentinel.backend
-    manager = ClusterManager(_options_manager())
+    manager = RBClusterManager(_options_manager())
 
     cluster, options = get_cluster_from_options(
         backend, {"cluster": "foo", "foo": "bar"}, cluster_manager=manager
@@ -95,7 +82,7 @@ def test_get_cluster_from_options_cluster_provided():
 
 def test_get_cluster_from_options_legacy_hosts_option():
     backend = mock.sentinel.backend
-    manager = ClusterManager(_options_manager())
+    manager = RBClusterManager(_options_manager())
 
     with pytest.warns(DeprecatedSettingWarning) as warninfo:
         cluster, options = get_cluster_from_options(
@@ -116,7 +103,7 @@ def test_get_cluster_from_options_legacy_hosts_option():
 
 def test_get_cluster_from_options_both_options_invalid():
     backend = mock.sentinel.backend
-    manager = ClusterManager(_options_manager())
+    manager = RBClusterManager(_options_manager())
 
     with pytest.raises(InvalidConfiguration):
         cluster, options = get_cluster_from_options(