replace RateLimit.Run with @contextmanager

posthog/clickhouse/client/limit.py

@@ -1,5 +1,6 @@
 import dataclasses
 import time
+from contextlib import contextmanager
 from functools import wraps
 from typing import Optional
 from collections.abc import Callable
@@ -64,29 +65,19 @@ class RateLimit:
     get_time: Callable[[], int] = lambda: int(time.time())
     applicable: Optional[Callable] = None  # allows to put a constraint on when rate limiting is used
     ttl: int = 60
+    redis_client = redis.get_client()
 
-    @property
-    def redis_client(self):
-        return redis.get_client()
-
-    class Run:
-        def __init__(self, limit, running_task_key, task_id):
-            self.limit = limit
-            self.running_task_key = running_task_key
-            self.task_id = task_id
-
-        def __enter__(self):
-            pass
-
-        def __exit__(self, applicable, running_task_key, task_id):
-            if self.limit:
-                self.limit.release(self.running_task_key, self.task_id)
-
+    @contextmanager
     def run(self, *args, **kwargs):
-        if self.applicable and not self.applicable(*args, **kwargs):
-            return self.Run(None, None, None)
-        running_tasks_key, task_id = self.use(*args, **kwargs)
-        return self.Run(self, running_tasks_key, task_id)
+        cleanup = None
+        if not self.applicable or self.applicable(*args, **kwargs):
+            running_task_key, task_id = self.use(*args, **kwargs)
+            cleanup = True
+        try:
+            yield
+        finally:
+            if cleanup:
+                self.release(running_task_key, task_id)
 
     def use(self, *args, **kwargs):
         """

posthog/clickhouse/client/test/test_limit.py

@@ -95,3 +95,17 @@ def test_run_applicable(self):
             result += 4
 
         assert result == 7
+
+    def test_exception(self):
+        result = 0
+        with self.assertRaises(Exception):
+            result += 1
+            with self.limit.run(is_api=True, team_id=9, task_id=17):
+                result += 2
+                raise Exception()
+            result += 4
+
+        with self.limit.run(is_api=True, team_id=9, task_id=17):
+            result += 8
+
+        assert result == 11

posthog/tasks/tasks.py

@@ -67,11 +67,11 @@ def process_query_task(
     is_api: bool,
     limit_context: Optional[LimitContext] = None,
 ) -> None:
+    """
+    Kick off query
+    Once complete save results to redis
+    """
     with get_api_personal_rate_limiter().run(is_api=is_api, team_id=team_id, task_id=query_id):
-        """
-        Kick off query
-        Once complete save results to redis
-        """
         from posthog.clickhouse.client import execute_process_query
 
         execute_process_query(