Fix stuck Celery tasks and add purge all tasks functionality

crawler/redis_repository.py

@@ -2,29 +2,31 @@ from datetime import timedelta
 from functools import lru_cache
 import json
 from string import Template
-from typing import Any
+from typing import Any, TypeVar
 from api.auth import User
 import redis
 from celery_app import app
 
+T = TypeVar("T")
+
 
 class RedisRepository:
-    redis_client: redis.Redis
-    tasks_key_template = Template("user:{user_id}/tasks")
+    redis_client: redis.Redis  # type: ignore[type-arg]
+    tasks_key_template: Template = Template("user:{user_id}/tasks")
 
-    def __init__(self):
-        redis_hostname = app.broker_connection().info()["hostname"]
-        redis_port = app.broker_connection().info()["port"]
+    def __init__(self) -> None:
+        redis_hostname: str = app.broker_connection().info()["hostname"]
+        redis_port: int = app.broker_connection().info()["port"]
         self.redis_client = redis.Redis(
             host=redis_hostname, port=redis_port, db=0, decode_responses=True
         )  # decode_responses=True returns str, not bytes
 
-    @lru_cache(maxsize=None)
     @staticmethod
-    def instance():
+    @lru_cache(maxsize=None)
+    def instance() -> "RedisRepository":
         return RedisRepository()
 
-    def set_key(self, key: str, value, ttl: timedelta | None = None) -> None:
+    def set_key(self, key: str, value: Any, ttl: timedelta | None = None) -> None:
         serialized_value = self.__serialize_value(value)
         self.redis_client.set(key, serialized_value)
 
@@ -37,9 +39,9 @@ class RedisRepository:
             return None
         return self.__deserialize_value(serialized_value)
 
-    def add_task_for_user(self, user: User, task_id: str):
+    def add_task_for_user(self, user: User, task_id: str) -> None:
         # Add the task ID to the Redis set for the user
-        current_tasks = (
+        current_tasks: list[str] = (
             self.get_key(self.tasks_key_template.substitute(user_id=user.email)) or []
         )
         self.set_key(
@@ -53,8 +55,27 @@ class RedisRepository:
             self.get_key(self.tasks_key_template.substitute(user_id=user.email)) or []
         )
 
+    def remove_task_for_user(self, user: User, task_id: str) -> bool:
+        """Remove a specific task from the user's task list."""
+        current_tasks: list[str] = self.get_tasks_for_user(user)
+        if task_id not in current_tasks:
+            return False
+        updated_tasks = [t for t in current_tasks if t != task_id]
+        self.set_key(
+            self.tasks_key_template.substitute(user_id=user.email),
+            updated_tasks,
+        )
+        return True
+
+    def clear_tasks_for_user(self, user: User) -> int:
+        """Clear all tasks for a user. Returns the number of tasks cleared."""
+        current_tasks: list[str] = self.get_tasks_for_user(user)
+        count = len(current_tasks)
+        self.redis_client.delete(self.tasks_key_template.substitute(user_id=user.email))
+        return count
+
     def __serialize_value(self, value: Any) -> str:
         return json.dumps(value)
 
-    def __deserialize_value(self, value_str) -> Any:
+    def __deserialize_value(self, value_str: str) -> Any:
         return json.loads(value_str)