Merge pull request #362 from phenobarbital/dev

changes on BackgrounQueue

Author: phenobarbital

navigator/background/queue/__init__.py

@@ -181,6 +181,13 @@ async def _execute_taskwrapper(self, task: TaskWrapper):
         with ThreadPoolExecutor(max_workers=1) as executor:
             try:
                 result = await task()
+            except asyncio.CancelledError:
+                self.logger.warning(
+                    f"TaskWrapper {task!r} was cancelled."
+                )
+                result = {
+                    "status": "cancelled"
+                }
             except Exception as e:
                 self.logger.exception(
                     f"Error executing TaskWrapper {task!r}: {e}",
@@ -190,6 +197,16 @@ async def _execute_taskwrapper(self, task: TaskWrapper):
                     "status": "failed",
                     "error": e
                 }
+                # Handle the exception and return a failure result
+                # This could be customized based on your needs
+                try:
+                    if hasattr(task, "tracker"):
+                        # set status = "failed"
+                        await task.tracker.set_failed(task.task_uuid)
+                except Exception as e:
+                    self.logger.error(
+                        f"Error updating tracker for TaskWrapper {task!r}: {e}"
+                    )
         return result
 
     async def _execute_coroutine(self, coro: coroutine):
@@ -305,6 +322,7 @@ async def process_queue(self):
                     print('TASK LOG ERROR > ', e)
                 # Call your task completion callback (if any)
                 try:
+                    print('RESULT > ', result)
                     await self._callback(task, result=result)
                 except Exception as e:
                     self.logger.error(

navigator/background/tracker/redis.py

@@ -2,6 +2,7 @@
 from typing import Any, Callable, Dict, Optional, Sequence, List, Mapping
 import asyncio
 import redis.asyncio as redis
+from datamodel.exceptions import ParserError
 from ...libs.json import json_encoder, json_decoder  # pylint: disable=E0611 # noqa
 from .models import JobRecord, time_now
 from ...conf import CACHE_URL
@@ -84,10 +85,15 @@ async def _update(self, job_id: str, **patch) -> None:
             rec: JobRecord = self._decoder(payload)
             for k, v in patch.items():
                 setattr(rec, k, v)
-            await self._redis.set(
-                key, self._encoder(rec), keepttl=True
-            )  # keep the TTL
-
+            try:
+                await self._redis.set(
+                    key, self._encoder(rec), keepttl=True
+                )  # keep the TTL
+            except ParserError as exc:
+                # we cannot encode the result of JobRecord, so we raise an error
+                raise ValueError(
+                    f"Invalid job record data: {exc}, payload: {exc.payload}"
+                ) from exc
             # Update secondary index for attributes if they are part of the patch
             if 'attributes' in patch:
                 # Remove old attributes from the index

navigator/background/wrappers/__init__.py

@@ -66,7 +66,7 @@ def __init__(
         self.fn = fn
         self.tracker = tracker
         self._name: str = kwargs.pop('name', fn.__name__ if fn else 'unknown_task')
-        self._callback_: Union[Callable, Awaitable] = kwargs.pop('callback', None)
+        self._user_callback: Union[Callable, Awaitable] = kwargs.pop('callback', None)
         job_status = kwargs.pop('status', 'pending')
         if job_status not in ['pending', 'running', 'done', 'failed']:
             raise ValueError(
@@ -76,12 +76,11 @@ def __init__(
         self.jitter: float = jitter
         # Create the Job Record at status "pending"
         # generate a list of arguments accepted by JobRecord:
-        job_args = {}
+        job_args = {
+            k: v for k, v in kwargs.items()
+            if not k.startswith('_') and k in JobRecord.__fields__
+        }
         content = kwargs.pop('content', None)
-        for k, v in kwargs.items():
-            if not k.startswith('_'):
-                if k in JobRecord.__fields__:
-                    job_args[k] = v
         self.job_record: JobRecord = JobRecord(
             name=self._name,
             content=content,
@@ -113,7 +112,28 @@ def add_callback(self, callback: Union[Callable, Awaitable]):
         - callback (Union[Callable, Awaitable]):
             Callback function to be called after the task is executed.
         """
-        self._callback_ = callback
+        self._user_callback = callback
+
+    async def _wrapped_callback(self, result, exc, loop):
+        """
+        Internal wrapper callback that injects JobRecord information
+        before calling the user's callback.
+
+        Args:
+        - result: The result of the task execution.
+        - exc: Exception raised during task execution, if any.
+        - loop: The event loop in which the task was executed.
+        """
+        if self._user_callback:
+            # Call user callback with additional JobRecord info
+            # New signature: callback(result, exc, loop, job_record, task_id)
+            await self._user_callback(
+                result,
+                exc,
+                loop,
+                job_record=self.job_record,
+                task_id=self.job_record.task_id
+            )
 
     async def __call__(self):
         result = None
@@ -166,7 +186,9 @@ async def _finish(result: Any, exc: Exception):
                 return result
             with ThreadPoolExecutor(max_workers=1) as executor:
                 coro = self.fn(*self.args, **self.kwargs)
-                coroutine_in_thread(coro, self._callback_, on_complete=_finish)
+                # Use the wrapped callback instead of the user callback directly
+                callback_to_use = self._wrapped_callback if self._user_callback else None
+                coroutine_in_thread(coro, callback_to_use, on_complete=_finish)
                 return {"status": "running"}
         except asyncio.CancelledError:
             self.logger.warning(

navigator/version.py

@@ -4,7 +4,7 @@
 __description__ = (
     "Navigator Web Framework based on aiohttp, " "with batteries included."
 )
-__version__ = "2.13.2"
+__version__ = "2.13.3"
 __copyright__ = "Copyright (c) 2020-2024 Jesus Lara"
 __author__ = "Jesus Lara"
 __author_email__ = "jesuslarag@gmail.com"