Merge pull request #14 from openaleph/feat/tracer

Feat/tracer

Author: simonwoerpel

openaleph_procrastinate/settings.py

@@ -1,8 +1,13 @@
 import random
+from typing import TYPE_CHECKING
 
+from anystore.types import Uri
 from pydantic import AliasChoices, Field
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
+if TYPE_CHECKING:
+    from openaleph_procrastinate.tracer import Tracer
+
 MAX_PRIORITY = 100
 MIN_PRIORITY = 0
 DEFAULT_DB_URI = "memory://"
@@ -42,6 +47,12 @@ def get_priority(self, priority: int | None = None) -> int:
         max_priority = max(min_priority, self.max_priority)
         return random.randint(min_priority, max_priority)
 
+    def get_tracer(self, uri: Uri | None = None) -> "Tracer":
+        """Get task status tracer."""
+        from openaleph_procrastinate.tracer import get_tracer
+
+        return get_tracer(self.queue, self.task, uri)
+
 
 class DeferSettings(BaseSettings):
     """
@@ -100,6 +111,11 @@ class DeferSettings(BaseSettings):
     )
     """ftm-assets"""
 
+    vectorize: ServiceSettings = ServiceSettings(
+        queue="vectorize", task="ftm_vectorize.tasks.vectorize", defer=False
+    )
+    """ftm-vectorize"""
+
     # OpenAleph
 
     index: ServiceSettings = ServiceSettings(
@@ -224,6 +240,9 @@ class OpenAlephSettings(BaseSettings):
     )
     """FollowTheMoney Fragments store uri"""
 
+    redis_url: str | None = Field(default=None, validation_alias="redis_url")
+    """Redis instance uri"""
+
     procrastinate_dehydrate_entities: bool = True
     """Dehydrate entity in job payload, jobs need to re-fetch entity from store"""
 

openaleph_procrastinate/tasks.py

@@ -1,12 +1,14 @@
 import functools
 import random
-from typing import Any, Callable
+from typing import Any, Callable, cast
 
 from anystore.logging import get_logger
+from anystore.types import Uri
 from procrastinate.app import App
 
 from openaleph_procrastinate.exceptions import ErrorHandler
-from openaleph_procrastinate.model import AnyJob, DatasetJob, Job
+from openaleph_procrastinate.model import AnyJob, DatasetJob, Job, Status
+from openaleph_procrastinate.tracer import get_tracer
 
 log = get_logger(__name__)
 
@@ -19,13 +21,28 @@ def unpack_job(data: dict[str, Any]) -> AnyJob:
         return Job(**data)
 
 
+def handle_trace(job: AnyJob, status: Status, tracer_uri: Uri | None) -> None:
+    if tracer_uri is not None and isinstance(job, DatasetJob):
+        tracer = get_tracer(job.queue, job.task, tracer_uri)
+        for entity in job.get_entities():
+            tracer.mark(cast(str, entity.id), status)
+
+
 def task(app: App, **kwargs):
     # https://procrastinate.readthedocs.io/en/stable/howto/advanced/middleware.html
+    tracer_uri = kwargs.pop("tracer_uri", None)
+
     def wrap(func: Callable[..., None]):
         def _inner(*job_args, **job_kwargs):
             # turn the json data into the job model instance
             job = unpack_job(job_kwargs)
-            func(*job_args, job)
+            handle_trace(job, "doing", tracer_uri)
+            try:
+                func(*job_args, job)
+                handle_trace(job, "succeeded", tracer_uri)
+            except Exception as e:
+                handle_trace(job, "failed", tracer_uri)
+                raise e
 
         # need to call to not register tasks twice (procrastinate complains)
         wrapped_func = functools.update_wrapper(_inner, func, updated=())

openaleph_procrastinate/tracer.py

@@ -0,0 +1,81 @@
+"""
+Simple, quick & dirty tracer module to track task status for entities in a
+shared redis instance. This is useful to share task status across services for
+quick lookup (psql would be too expensive), e.g. to show a loading spinner in
+the UI if something is processing. This module currently is a quick shot to
+solve the status lookup for specific queues, it should be refactored at one
+point into a more general tracer module that e.g. can also store processing
+exceptions in postgres. For this, creating a better class-based queue system and
+merge the tracing into it should be considered.
+
+The tracer is used in the @task decorator (middleware) when tracer_uri=... is set
+in the kwargs. e.g.:
+
+```python
+@task(queue="my-queue", trace_uri="redis://localhost")
+def process(job) -> None:
+    pass
+```
+
+(See test suite for example)
+
+The tracer backend accepts any uri (sql, file-like, ...), but it is preferable
+to use redis for performance reasons.
+"""
+
+from functools import cache
+
+from anystore.store import get_store
+from anystore.types import Uri
+from anystore.util import join_relpaths
+
+from openaleph_procrastinate.model import Status
+from openaleph_procrastinate.settings import OpenAlephSettings
+
+
+class Tracer:
+    def __init__(self, queue: str, task: str, uri: Uri | None = None) -> None:
+        if uri is None:
+            settings = OpenAlephSettings()
+            uri = settings.redis_url
+        self._store = get_store(uri or "memory://")
+        self.queue = queue
+        self.task = task
+
+    def _make_key(self, entity_id: str) -> str:
+        return join_relpaths(
+            "openaleph-procrastinate", "tracer", self.queue, self.task, entity_id
+        )
+
+    def mark(self, entity_id: str, status: Status) -> None:
+        """Mark an entity status for the given queue and task. If status is
+        'succeeded' remove the data from the tracer."""
+        key = self._make_key(entity_id)
+        if status == "succeeded":
+            return self._store.delete(key, ignore_errors=True)
+        self._store.put(key, status)
+
+    def add(self, entity_id: str) -> None:
+        """Mark as todo"""
+        self.mark(entity_id, "todo")
+
+    def start(self, entity_id: str) -> None:
+        """Mark as doing"""
+        self.mark(entity_id, "doing")
+
+    def finish(self, entity_id: str) -> None:
+        """Mark done which is actually popping (deleting) from the tracer"""
+        self.mark(entity_id, "succeeded")
+
+    def is_processing(self, entity_id: str) -> bool:
+        """Check if a task for the entity_id is either pending or doing"""
+        key = self._make_key(entity_id)
+        if not self._store.exists(key):
+            return False
+        status = self._store.get(key)
+        return status in ("todo", "doing")
+
+
+@cache
+def get_tracer(queue: str, task: str, uri: Uri | None) -> Tracer:
+    return Tracer(queue, task, uri)