🐛 Ensure file properties are passed through

Author: simonwoerpel

openaleph_procrastinate/helpers.py

@@ -4,7 +4,7 @@
 
 from contextlib import contextmanager
 from pathlib import Path
-from typing import ContextManager, Generator
+from typing import ContextManager, Generator, Iterable
 
 from anystore.logic.virtual import VirtualIO
 from followthemoney import EntityProxy
@@ -60,6 +60,18 @@ def load_entity(dataset: str, entity_id: str) -> EntityProxy:
     return entity
 
 
+def load_entities(
+    dataset: str, entity_ids: Iterable[str]
+) -> Generator[EntityProxy, None, None]:
+    """
+    Batch retrieve entities from the fragment store.
+    """
+    store = get_fragments(
+        dataset, database_uri=settings.fragments_uri, **sqlalchemy_pool
+    )
+    yield from store.iterate(entity_ids)
+
+
 @contextmanager
 def entity_writer(
     dataset: str, origin: str = OPAL_ORIGIN

openaleph_procrastinate/model.py

@@ -20,7 +20,7 @@
     MIN_PRIORITY,
     OpenAlephSettings,
 )
-from openaleph_procrastinate.util import make_checksum_entity
+from openaleph_procrastinate.util import make_file_entity
 
 settings = OpenAlephSettings()
 
@@ -125,7 +125,9 @@ def log(self) -> BoundLogger:
             batch=self.batch,
         )
 
-    def get_writer(self: Self, origin: str = helpers.OPAL_ORIGIN) -> ContextManager[BulkLoader]:
+    def get_writer(
+        self: Self, origin: str = helpers.OPAL_ORIGIN
+    ) -> ContextManager[BulkLoader]:
         """Get the writer for the dataset of the current job"""
         return helpers.entity_writer(self.dataset, origin)
 
@@ -145,8 +147,8 @@ def load_entities(self: Self) -> Generator[EntityProxy, None, None]:
         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"])
+            entity_ids = [e["id"] for e in self.payload["entities"]]
+            yield from helpers.load_entities(self.dataset, entity_ids)
 
     # Helpers for file jobs that access the servicelayer archive
 
@@ -201,7 +203,7 @@ def from_entities(
             context: Job context
         """
         if dehydrate:
-            entities_ = (make_checksum_entity(e, quiet=True) for e in entities)
+            entities_ = (make_file_entity(e, quiet=True) for e in entities)
             entities = (e for e in entities_ if e is not None)
         return cls(
             dataset=dataset,

openaleph_procrastinate/util.py

@@ -22,7 +22,7 @@ def make_stub_entity(
     )
 
 
-def make_checksum_entity(
+def make_file_entity(
     e: E, entity_type: Type[E] | None = ValueEntity, quiet: bool | None = False
 ) -> E | None:
     """
@@ -32,4 +32,6 @@ def make_checksum_entity(
     stub = make_stub_entity(e, entity_type)
     if stub is not None:
         stub.add("contentHash", e.get("contentHash", quiet=q), quiet=q)
+        stub.add("fileName", e.get("fileName", quiet=q), quiet=q)
+        stub.add("parent", e.get("parent", quiet=q), quiet=q)
         return stub

tests/test_util.py

@@ -8,6 +8,7 @@ def test_util():
     e.id = "a"
     e.add("fileName", "test.txt")
     e.add("contentHash", "123")
+    e.add("parent", "1")
 
     assert util.make_stub_entity(e).to_dict() == {
         "caption": "test.txt",
@@ -18,16 +19,20 @@ def test_util():
         "referents": [],
     }
 
-    assert util.make_checksum_entity(e).to_dict() == {
+    assert util.make_file_entity(e).to_dict() == {
         "caption": "test.txt",
         "id": "a",
         "schema": "Document",
-        "properties": {"contentHash": ["123"]},
+        "properties": {
+            "contentHash": ["123"],
+            "fileName": ["test.txt"],
+            "parent": ["1"],
+        },
         "datasets": ["default"],
         "referents": [],
     }
 
     for i in ("", None):
         e.id = i
         assert util.make_stub_entity(e) is None
-        assert util.make_checksum_entity(e) is None
+        assert util.make_file_entity(e) is None