feat: Create Django Command for Deleting Rate Limit Redis Keys

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)

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")