feat: add typings for utils/redis.py (#65454)

Adds typings to `utils/redis.py`

fixtures/stubs-for-mypy/rediscluster/client.pyi

@@ -1,8 +1,43 @@
-from typing import TypeVar
+from typing import TypeVar, Any, Literal, overload
 
 from redis import Redis
 
 T = TypeVar("T", str, bytes)
 
 class RedisCluster(Redis[T]):
-    ...
+    @overload
+    def __init__(
+        self: RedisCluster[str],
+        *,
+        startup_nodes: list[dict[str, Any]],
+        decode_responses: Literal[True],
+        skip_full_coverage_check: bool,
+        max_connections: int,
+        max_connections_per_node: bool,
+        readonly_mode: bool,
+        **client_args: object,
+    ) -> None: ...
+    @overload
+    def __init__(
+        self: RedisCluster[bytes],
+        *,
+        startup_nodes: list[dict[str, Any]],
+        decode_responses: Literal[False],
+        skip_full_coverage_check: bool,
+        max_connections: int,
+        max_connections_per_node: bool,
+        readonly_mode: bool,
+        **client_args: object,
+    ) -> None: ...
+    @overload
+    def __init__(
+        self,
+        *,
+        startup_nodes: list[dict[str, Any]],
+        decode_responses: bool,
+        skip_full_coverage_check: bool,
+        max_connections: int,
+        max_connections_per_node: bool,
+        readonly_mode: bool,
+        **client_args: object,
+    ) -> None: ...

pyproject.toml

@@ -517,7 +517,6 @@ module = [
     "sentry.utils.performance_issues.detectors.render_blocking_asset_span_detector",
     "sentry.utils.performance_issues.detectors.slow_db_query_detector",
     "sentry.utils.performance_issues.performance_detection",
-    "sentry.utils.redis",
     "sentry.utils.sentry_apps.webhooks",
     "sentry.utils.services",
     "sentry.utils.snowflake",

src/sentry/buffer/redis.py

@@ -6,8 +6,10 @@ import threading
 from datetime import date, datetime, timezone
 from time import time
 
+import rb
 from django.db import models
 from django.utils.encoding import force_bytes, force_str
+from rediscluster import RedisCluster
 
 from sentry.buffer.base import Buffer
 from sentry.tasks.process_buffer import process_incr, process_pending
@@ -15,7 +17,12 @@ from sentry.utils import json, metrics
 from sentry.utils.compat import crc32
 from sentry.utils.hashlib import md5_text
 from sentry.utils.imports import import_string
-from sentry.utils.redis import get_dynamic_cluster_from_options, validate_dynamic_cluster
+from sentry.utils.redis import (
+    get_dynamic_cluster_from_options,
+    is_instance_rb_cluster,
+    is_instance_redis_cluster,
+    validate_dynamic_cluster,
+)
 
 _local_buffers = None
 _local_buffers_lock = threading.Lock()
@@ -82,11 +89,13 @@ class RedisBuffer(Buffer):
         assert self.pending_partitions > 0
         assert self.incr_batch_size > 0
 
-    def get_routing_client(self):
-        if self.is_redis_cluster:
+    def get_routing_client(self) -> RedisCluster | rb.RoutingClient:
+        if is_instance_redis_cluster(self.cluster, self.is_redis_cluster):
             return self.cluster
-        else:
+        elif is_instance_rb_cluster(self.cluster, self.is_redis_cluster):
             return self.cluster.get_routing_client()
+        else:
+            raise AssertionError("unreachable")
 
     def validate(self):
         validate_dynamic_cluster(self.is_redis_cluster, self.cluster)
@@ -182,11 +191,13 @@ class RedisBuffer(Buffer):
         Fetches buffered values for a model/filter. Passed columns must be integer columns.
         """
         key = self._make_key(model, filters)
-        if self.is_redis_cluster:
+        if is_instance_redis_cluster(self.cluster, self.is_redis_cluster):
             pipe = self.cluster.pipeline(transaction=False)
-        else:
+        elif is_instance_rb_cluster(self.cluster, self.is_redis_cluster):
             conn = self.cluster.get_local_client_for_key(key)
             pipe = conn.pipeline()
+        else:
+            raise AssertionError("unreachable")
 
         for col in columns:
             pipe.hget(key, f"i+{col}")
@@ -211,16 +222,18 @@ class RedisBuffer(Buffer):
         pending_key = self._make_pending_key_from_key(key)
         # We can't use conn.map() due to wanting to support multiple pending
         # keys (one per Redis partition)
-        if self.is_redis_cluster:
+        if is_instance_redis_cluster(self.cluster, self.is_redis_cluster):
             conn = self.cluster
-        else:
+        elif is_instance_rb_cluster(self.cluster, self.is_redis_cluster):
             conn = self.cluster.get_local_client_for_key(key)
+        else:
+            raise AssertionError("unreachable")
 
         pipe = conn.pipeline()
         pipe.hsetnx(key, "m", f"{model.__module__}.{model.__name__}")
         _validate_json_roundtrip(filters, model)
 
-        if self.is_redis_cluster:
+        if is_instance_redis_cluster(self.cluster, self.is_redis_cluster):
             pipe.hsetnx(key, "f", json.dumps(self._dump_values(filters)))
         else:
             pipe.hsetnx(key, "f", pickle.dumps(filters))
@@ -234,7 +247,7 @@ class RedisBuffer(Buffer):
             # e.g. "update score if last_seen or times_seen is changed"
             _validate_json_roundtrip(extra, model)
             for column, value in extra.items():
-                if self.is_redis_cluster:
+                if is_instance_redis_cluster(self.cluster, self.is_redis_cluster):
                     pipe.hset(key, "e+" + column, json.dumps(self._dump_value(value)))
                 else:
                     pipe.hset(key, "e+" + column, pickle.dumps(value))
@@ -263,10 +276,7 @@ class RedisBuffer(Buffer):
             # super fast and is fine to do redundantly.
 
         pending_key = self._make_pending_key(partition)
-        if self.is_redis_cluster:
-            client = self.cluster
-        else:
-            client = self.cluster.get_routing_client()
+        client = self.get_routing_client()
         lock_key = self._make_lock_key(pending_key)
         # prevent a stampede due to celerybeat + periodic task
         if not client.set(lock_key, "1", nx=True, ex=60):
@@ -276,7 +286,7 @@ class RedisBuffer(Buffer):
 
         try:
             keycount = 0
-            if self.is_redis_cluster:
+            if is_instance_redis_cluster(self.cluster, self.is_redis_cluster):
                 keys = self.cluster.zrange(pending_key, 0, -1)
                 keycount += len(keys)
 
@@ -286,22 +296,24 @@ class RedisBuffer(Buffer):
                         process_incr.apply_async(kwargs={"batch_keys": pending_buffer.flush()})
 
                 self.cluster.zrem(pending_key, *keys)
-            else:
+            elif is_instance_rb_cluster(self.cluster, self.is_redis_cluster):
                 with self.cluster.all() as conn:
                     results = conn.zrange(pending_key, 0, -1)
 
                 with self.cluster.all() as conn:
-                    for host_id, keys in results.value.items():
-                        if not keys:
+                    for host_id, keysb in results.value.items():
+                        if not keysb:
                             continue
-                        keycount += len(keys)
-                        for key in keys:
-                            pending_buffer.append(key.decode("utf-8"))
+                        keycount += len(keysb)
+                        for keyb in keysb:
+                            pending_buffer.append(keyb.decode("utf-8"))
                             if pending_buffer.full():
                                 process_incr.apply_async(
                                     kwargs={"batch_keys": pending_buffer.flush()}
                                 )
-                        conn.target([host_id]).zrem(pending_key, *keys)
+                        conn.target([host_id]).zrem(pending_key, *keysb)
+            else:
+                raise AssertionError("unreachable")
 
             # queue up remainder of pending keys
             if not pending_buffer.empty():
@@ -325,11 +337,7 @@ class RedisBuffer(Buffer):
         return super().process(model, columns, filters, extra, signal_only)
 
     def _process_single_incr(self, key):
-        if self.is_redis_cluster:
-            client = self.cluster
-        else:
-            client = self.cluster.get_routing_client()
-
+        client = self.get_routing_client()
         lock_key = self._make_lock_key(key)
         # prevent a stampede due to the way we use celery etas + duplicate
         # tasks
@@ -341,12 +349,13 @@ class RedisBuffer(Buffer):
         pending_key = self._make_pending_key_from_key(key)
 
         try:
-            if self.is_redis_cluster:
+            if is_instance_redis_cluster(self.cluster, self.is_redis_cluster):
                 pipe = self.cluster.pipeline(transaction=False)
-            else:
+            elif is_instance_rb_cluster(self.cluster, self.is_redis_cluster):
                 conn = self.cluster.get_local_client_for_key(key)
                 pipe = conn.pipeline()
-
+            else:
+                raise AssertionError("unreachable")
             pipe.hgetall(key)
             pipe.zrem(pending_key, key)
             pipe.delete(key)

src/sentry/cache/redis.py

@@ -1,5 +1,5 @@
 from sentry.utils import json
-from sentry.utils.redis import get_cluster_from_options, redis_clusters
+from sentry.utils.redis import get_cluster_from_options, is_instance_rb_cluster, redis_clusters
 
 from .base import BaseCache
 
@@ -55,7 +55,10 @@ class CommonRedisCache(BaseCache):
 class RbCache(CommonRedisCache):
     def __init__(self, **options: object) -> None:
         cluster, options = get_cluster_from_options("SENTRY_CACHE_OPTIONS", options)
-        client = cluster.get_routing_client()
+        if is_instance_rb_cluster(cluster, False):
+            client = cluster.get_routing_client()
+        else:
+            raise AssertionError("unreachable")
         # XXX: rb does not have a "raw" client -- use the default client
         super().__init__(client=client, raw_client=client, **options)
 

src/sentry/digests/backends/redis.py

@@ -4,6 +4,7 @@ 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
 
@@ -71,7 +72,9 @@ class RedisBackend(Backend):
     """
 
     def __init__(self, **options: Any) -> None:
-        self.cluster, options = get_cluster_from_options("SENTRY_DIGESTS_OPTIONS", options)
+        cluster, options = get_cluster_from_options("SENTRY_DIGESTS_OPTIONS", options)
+        assert isinstance(cluster, rb.Cluster)
+        self.cluster = cluster
         self.locks = LockManager(RedisLockBackend(self.cluster))
 
         self.namespace = options.pop("namespace", "d")

src/sentry/incidents/subscription_processor.py

@@ -10,6 +10,7 @@ from typing import TypeVar, cast
 from django.conf import settings
 from django.db import router, transaction
 from django.utils import timezone
+from sentry_redis_tools.retrying_cluster import RetryingRedisCluster
 from snuba_sdk import Column, Condition, Limit, Op
 
 from sentry import features
@@ -47,7 +48,6 @@ from sentry.snuba.models import QuerySubscription
 from sentry.snuba.tasks import build_query_builder
 from sentry.utils import metrics, redis
 from sentry.utils.dates import to_datetime, to_timestamp
-from sentry.utils.redis import RetryingRedisCluster
 
 logger = logging.getLogger(__name__)
 REDIS_TTL = int(timedelta(days=7).total_seconds())

src/sentry/models/featureadoption.py

@@ -1,9 +1,11 @@
 import logging
 from typing import ClassVar, cast
 
+import rb
 from django.conf import settings
 from django.db import IntegrityError, models, router, transaction
 from django.utils import timezone
+from rediscluster import RedisCluster
 
 from sentry.adoption import manager
 from sentry.adoption.manager import UnknownFeature
@@ -16,7 +18,11 @@ from sentry.db.models import (
     region_silo_only_model,
     sane_repr,
 )
-from sentry.utils.redis import get_dynamic_cluster_from_options
+from sentry.utils.redis import (
+    get_dynamic_cluster_from_options,
+    is_instance_rb_cluster,
+    is_instance_redis_cluster,
+)
 from sentry.utils.services import build_instance_from_options
 
 logger = logging.getLogger(__name__)
@@ -125,12 +131,14 @@ class FeatureAdoptionRedisBackend:
             "SENTRY_FEATURE_ADOPTION_CACHE_OPTIONS", options
         )
 
-    def get_client(self, key):
+    def get_client(self, key: str) -> rb.Cluster | RedisCluster:
         # WARN: Carefully as this works only for single key operations.
-        if self.is_redis_cluster:
+        if is_instance_redis_cluster(self.cluster, self.is_redis_cluster):
             return self.cluster
-        else:
+        elif is_instance_rb_cluster(self.cluster, self.is_redis_cluster):
             return self.cluster.get_local_client_for_key(key)
+        else:
+            raise AssertionError("unreachable")
 
     def in_cache(self, organization_id, feature_id):
         org_key = self.key_tpl.format(organization_id)

src/sentry/options/__init__.py

@@ -51,6 +51,7 @@ __all__ = (
     "register",
     "unregister",
     "set",
+    "OptionsManager",
 )
 
 # See notes in ``runner.initializer`` regarding lazy cache configuration.

src/sentry/relay/projectconfig_debounce_cache/redis.py

@@ -1,6 +1,14 @@
+import rb
+from rediscluster import RedisCluster
+
 from sentry.relay.projectconfig_debounce_cache.base import ProjectConfigDebounceCache
 from sentry.utils import metrics
-from sentry.utils.redis import get_dynamic_cluster_from_options, validate_dynamic_cluster
+from sentry.utils.redis import (
+    get_dynamic_cluster_from_options,
+    is_instance_rb_cluster,
+    is_instance_redis_cluster,
+    validate_dynamic_cluster,
+)
 
 REDIS_CACHE_TIMEOUT = 3600  # 1 hr
 
@@ -28,11 +36,13 @@ class RedisProjectConfigDebounceCache(ProjectConfigDebounceCache):
     def validate(self):
         validate_dynamic_cluster(self.is_redis_cluster, self.cluster)
 
-    def _get_redis_client(self, routing_key):
-        if self.is_redis_cluster:
+    def _get_redis_client(self, routing_key: str) -> rb.Cluster | RedisCluster:
+        if is_instance_redis_cluster(self.cluster, self.is_redis_cluster):
             return self.cluster
-        else:
+        elif is_instance_rb_cluster(self.cluster, self.is_redis_cluster):
             return self.cluster.get_local_client_for_key(routing_key)
+        else:
+            raise AssertionError("unreachable")
 
     def is_debounced(self, *, public_key, project_id, organization_id):
         if organization_id:

src/sentry/utils/imports.py

@@ -17,7 +17,7 @@ class ModuleProxyCache(dict):
 _cache = ModuleProxyCache()
 
 
-def import_string(path):
+def import_string(path: str):
     """
     Path must be module.path.ClassName