diff --git a/core/management/commands/delete_rate_limit_keys.py b/core/management/commands/delete_rate_limit_keys.py new file mode 100644 index 0000000000..075a33862e --- /dev/null +++ b/core/management/commands/delete_rate_limit_keys.py @@ -0,0 +1,28 @@ +from django.core.management.base import BaseCommand, CommandParser + +from services.redis_configuration import get_redis_connection + + +class Command(BaseCommand): + help = "This command is meant to delete all rate limit redis keys for either userId or ip." + + def add_arguments(self, parser: CommandParser) -> None: + # This argument switches the command to "anonymous mode" deleting all the ip based keys + parser.add_argument("--ip", type=bool) + + def handle(self, *args, **options): + redis = get_redis_connection() + + path = "rl-user:*" + if options["ip"]: + path = "rl-ip:*" + + try: + for key in redis.scan_iter(path): + # -1 means the key has no expiry + if redis.ttl(key) == -1: + print(f"Deleting key: {key.decode('utf-8')}") + redis.delete(key) + except Exception as e: + print("Error occurred when deleting redis keys") + print(e) diff --git a/core/tests/test_management_commands.py b/core/tests/test_management_commands.py index bca145041a..0741698556 100644 --- a/core/tests/test_management_commands.py +++ b/core/tests/test_management_commands.py @@ -6,6 +6,8 @@ from shared.config import ConfigHelper from shared.django_apps.core.tests.factories import OwnerFactory, RepositoryFactory +from services.redis_configuration import get_redis_connection + @pytest.mark.django_db def test_update_gitlab_webhook_command(mocker): @@ -77,3 +79,45 @@ def test_update_gitlab_webhook_command(mocker): secret=repo3.webhook_secret, ), ] + + +@pytest.mark.django_db +def test_delete_rate_limit_keys_user_id(): + redis = get_redis_connection() + redis.set("rl-user:1", 1) + redis.set("rl-user:2", 1, ex=5000) + redis.set("rl-ip:1", 1) + + call_command( + "delete_rate_limit_keys", + stdout=StringIO(), + stderr=StringIO(), + ) + + assert redis.get("rl-user:1") is None + assert redis.get("rl-user:2") is not None + assert redis.get("rl-ip:1") is not None + + # Get rid of lingering keys + redis.delete("rl-ip:1") + redis.delete("rl-user:2") + + +@pytest.mark.django_db +def test_delete_rate_limit_keys_ip_option(): + redis = get_redis_connection() + redis.set("rl-ip:1", 1) + redis.set("rl-ip:2", 1, ex=5000) + redis.set("rl-user:1", 1) + + call_command( + "delete_rate_limit_keys", stdout=StringIO(), stderr=StringIO(), ip=True + ) + + assert redis.get("rl-ip:1") is None + assert redis.get("rl-ip:2") is not None + assert redis.get("rl-user:1") is not None + + # Get rid of lingering keys + redis.delete("rl-user:1") + redis.delete("rl-ip:2")