⚡️ Optionally don't dehydrate entity payloads

Author: simonwoerpel

openaleph_procrastinate/model.py

@@ -3,14 +3,15 @@
 from pathlib import Path
 from typing import Any, ContextManager, Generator, Iterable, Literal, Self, TypeAlias
 
-from anystore.logging import BoundLogger, get_logger
+from anystore.logging import get_logger
 from anystore.store.virtual import VirtualIO
 from anystore.util import clean_dict
 from banal import ensure_dict
 from followthemoney import model
 from followthemoney.proxy import EntityProxy
 from ftmq.store.fragments.loader import BulkLoader
 from pydantic import BaseModel, ConfigDict, computed_field
+from structlog.stdlib import BoundLogger
 
 from openaleph_procrastinate import helpers
 from openaleph_procrastinate.app import App, run_sync_worker
@@ -139,8 +140,13 @@ def get_entities(self) -> Generator[EntityProxy, None, None]:
     def load_entities(self: Self) -> Generator[EntityProxy, None, None]:
         """Load the entities from the store to refresh it to the latest data"""
         assert "entities" in self.payload, "No entities in payload"
-        for data in self.payload["entities"]:
-            yield helpers.load_entity(self.dataset, data["id"])
+
+        # if not dehydrating, just get the payload data:
+        if not settings.procrastinate_dehydrate_entities:
+            yield from self.get_entities()
+        else:
+            for data in self.payload["entities"]:
+                yield helpers.load_entity(self.dataset, data["id"])
 
     # Helpers for file jobs that access the servicelayer archive
 
@@ -179,7 +185,7 @@ def from_entities(
         queue: str,
         task: str,
         entities: Iterable[EntityProxy],
-        dehydrate: bool | None = False,
+        dehydrate: bool = settings.procrastinate_dehydrate_entities,
         **context: Any,
     ) -> Self:
         """

openaleph_procrastinate/settings.py

@@ -224,6 +224,9 @@ class OpenAlephSettings(BaseSettings):
     )
     """FollowTheMoney Fragments store uri"""
 
+    procrastinate_dehydrate_entities: bool = True
+    """Dehydrate entity in job payload, jobs need to re-fetch entity from store"""
+
     @property
     def in_memory_db(self) -> bool:
         return self.procrastinate_db_uri.startswith("memory:")