feat(manual-mode): Add manual model

Author: gagantrivedi

settings/dev.py

@@ -68,5 +68,9 @@
 DOCGEN_MODE = env.bool("DOCGEN_MODE", default=False)
 TASK_RUN_METHOD = TaskRunMethod.TASK_PROCESSOR
 
+# To be used by test that depend on task processor to run tasks manually by
+# calling `run_tasks`
+TASK_PROCESSOR_MANUAL_MODE = env.bool("TASK_PROCESSOR_MANUAL_MODE", default=False)
+
 # Avoid models.W042 warnings
 DEFAULT_AUTO_FIELD = "django.db.models.AutoField"

src/task_processor/decorators.py

@@ -177,7 +177,7 @@ def register_recurring_task(
     first_run_time: time | None = None,
     timeout: timedelta | None = timedelta(minutes=30),
 ) -> typing.Callable[[TaskCallable[TaskParameters]], TaskCallable[TaskParameters]]:
-    if not settings.TASK_PROCESSOR_MODE:
+    if not settings.TASK_PROCESSOR_MODE or settings.TASK_PROCESSOR_MANUAL_MODE:
         # Do not register recurring tasks if not invoked by task processor
         return lambda f: f
 

src/task_processor/processor.py

@@ -1,3 +1,4 @@
+import functools
 import logging
 import traceback
 import typing
@@ -109,33 +110,65 @@ def run_recurring_tasks(database: str) -> list[RecurringTaskRun]:
     return []
 
 
+def task_metrics(
+    func: typing.Callable[[T], typing.Tuple[T, AnyTaskRun]],
+) -> typing.Callable[[T], typing.Tuple[T, AnyTaskRun]]:
+    @functools.wraps(func)
+    def wrapper(
+        task: T, *args: typing.Any, **kwargs: typing.Any
+    ) -> typing.Tuple[T, AnyTaskRun]:
+        # Only collect metrics when TASK_PROCESSOR_MODE is True
+        if not settings.TASK_PROCESSOR_MODE:
+            return func(task, *args, **kwargs)
+
+        ctx = ExitStack()
+        timer = metrics.flagsmith_task_processor_task_duration_seconds.time()
+        ctx.enter_context(timer)
+
+        task_obj, task_run = func(task, *args, **kwargs)
+        result = task_run.result
+
+        # Get task info for labels
+        task_identifier = task.task_identifier
+        registered_task = get_task(task_identifier)
+
+        labels = {
+            "task_identifier": task_identifier,
+            "task_type": registered_task.task_type.value.lower(),
+            "result": result.lower(),
+        }  # type: ignore[union-attr]
+
+        timer.labels(**labels)  # type: ignore[no-untyped-call]
+        ctx.close()
+        metrics.flagsmith_task_processor_finished_tasks_total.labels(**labels).inc()
+
+        return task_obj, task_run
+
+    return wrapper
+
+
+@task_metrics
 def _run_task(
     task: T,
 ) -> typing.Tuple[T, AnyTaskRun]:
-    assert settings.TASK_PROCESSOR_MODE, (
-        "Attempt to run tasks in a non-task-processor environment"
-    )
-
-    ctx = ExitStack()
-    timer = metrics.flagsmith_task_processor_task_duration_seconds.time()
-    ctx.enter_context(timer)
-
+    assert (
+        settings.TASK_PROCESSOR_MODE or settings.TASK_PROCESSOR_MANUAL_MODE
+    ), "Attempt to run tasks in a non-task-processor environment"
     task_identifier = task.task_identifier
     registered_task = get_task(task_identifier)
 
     logger.debug(
         f"Running task {task_identifier} id={task.pk} args={task.args} kwargs={task.kwargs}"
     )
     task_run: AnyTaskRun = task.task_runs.model(started_at=timezone.now(), task=task)  # type: ignore[attr-defined]
-    result: str
 
     try:
         with ThreadPoolExecutor(max_workers=1) as executor:
             future = executor.submit(task.run)
             timeout = task.timeout.total_seconds() if task.timeout else None
             future.result(timeout=timeout)  # Wait for completion or timeout
 
-        task_run.result = result = TaskResult.SUCCESS.value
+        task_run.result = TaskResult.SUCCESS.value
         task_run.finished_at = timezone.now()
         task.mark_success()
 
@@ -148,7 +181,7 @@ def _run_task(
 
         task.mark_failure()
 
-        task_run.result = result = TaskResult.FAILURE.value
+        task_run.result = TaskResult.FAILURE.value
         task_run.error_details = str(traceback.format_exc())
 
         logger.error(
@@ -160,9 +193,9 @@ def _run_task(
         )
 
         if isinstance(e, TaskBackoffError):
-            assert registered_task.task_type == TaskType.STANDARD, (
-                "Attempt to back off a recurring task (currently not supported)"
-            )
+            assert (
+                registered_task.task_type == TaskType.STANDARD
+            ), "Attempt to back off a recurring task (currently not supported)"
             if typing.TYPE_CHECKING:
                 assert isinstance(task, Task)
             if task.num_failures <= 3:
@@ -176,15 +209,4 @@ def _run_task(
                     delay_until,
                 )
 
-    labels = {
-        "task_identifier": task_identifier,
-        "task_type": registered_task.task_type.value.lower(),
-        "result": result.lower(),
-    }
-
-    timer.labels(**labels)  # type: ignore[no-untyped-call]
-    ctx.close()
-
-    metrics.flagsmith_task_processor_finished_tasks_total.labels(**labels).inc()
-
     return task, task_run