rate_limiter: Extract KEY_PREFIX to redis_utils.

This commit is contained in:
Mateusz Mandera 2024-03-19 02:46:33 +01:00 committed by Tim Abbott
parent 706c380971
commit 630335142a
3 changed files with 13 additions and 10 deletions

View File

@ -1,5 +1,4 @@
import logging
import os
import time
from abc import ABC, abstractmethod
from typing import Dict, List, Optional, Set, Tuple, Type, cast
@ -11,6 +10,7 @@ from django.conf import settings
from django.http import HttpRequest
from typing_extensions import override
from zerver.lib import redis_utils
from zerver.lib.cache import cache_with_key
from zerver.lib.exceptions import RateLimitedError
from zerver.lib.redis_utils import get_redis_client
@ -22,8 +22,6 @@ from zerver.models import UserProfile
client = get_redis_client()
rules: Dict[str, List[Tuple[int, int]]] = settings.RATE_LIMITING_RULES
KEY_PREFIX = ""
logger = logging.getLogger(__name__)
@ -159,11 +157,6 @@ class RateLimitedIPAddr(RateLimitedObject):
return rules[self.domain]
def bounce_redis_key_prefix_for_testing(test_name: str) -> None:
global KEY_PREFIX
KEY_PREFIX = test_name + ":" + str(os.getpid()) + ":"
class RateLimiterBackend(ABC):
@classmethod
@abstractmethod
@ -323,7 +316,8 @@ class RedisRateLimiterBackend(RateLimiterBackend):
@classmethod
def get_keys(cls, entity_key: str) -> List[str]:
return [
f"{KEY_PREFIX}ratelimit:{entity_key}:{keytype}" for keytype in ["list", "zset", "block"]
f"{redis_utils.REDIS_KEY_PREFIX}ratelimit:{entity_key}:{keytype}"
for keytype in ["list", "zset", "block"]
]
@classmethod

View File

@ -1,3 +1,4 @@
import os
import re
import secrets
from typing import Any, Dict, Mapping, Optional
@ -84,3 +85,11 @@ def validate_key_fits_format(key: str, key_format: str) -> None:
if not re.fullmatch(regex, key):
raise ZulipRedisKeyOfWrongFormatError(f"{key} does not match format {key_format}")
REDIS_KEY_PREFIX = ""
def bounce_redis_key_prefix_for_testing(test_name: str) -> None:
global REDIS_KEY_PREFIX
REDIS_KEY_PREFIX = test_name + ":" + str(os.getpid()) + ":"

View File

@ -67,7 +67,7 @@ from zerver.lib.initial_password import initial_password
from zerver.lib.message import access_message
from zerver.lib.notification_data import UserMessageNotificationsData
from zerver.lib.per_request_cache import flush_per_request_caches
from zerver.lib.rate_limiter import bounce_redis_key_prefix_for_testing
from zerver.lib.redis_utils import bounce_redis_key_prefix_for_testing
from zerver.lib.sessions import get_session_dict_user
from zerver.lib.soft_deactivation import do_soft_deactivate_users
from zerver.lib.stream_subscription import get_subscribed_stream_ids_for_user