CallbackRegistrar: add correlation hook with master record

Adds `register_correlation_hook_with_master_record` hook - identical
to existing `register_correlation_hook`, but with an additional
master record parameter sent to the callback. Some modules need
access to the master record and since it's already loaded during
correlation process (snapshot creation), providing it doesn't incur
any additional overhead. One such example is the new TS last activity
module in the NERD2 system.

For backwards compatibility, the new hook type was created instead of
modifying the `register_correlation_hook`. Internally, they both
share the same implementation.

Author: dbnk0

dp3/common/callback_registrar.py

@@ -13,6 +13,7 @@
 from dp3.common.state import SharedFlag
 from dp3.common.task import DataPointTask
 from dp3.common.types import ParsedTimedelta
+from dp3.common.utils import get_func_name
 from dp3.core.updater import Updater
 from dp3.snapshots.snapshooter import SnapShooter
 from dp3.task_processing.task_executor import TaskExecutor
@@ -365,6 +366,41 @@ def register_correlation_hook(
             may_change: each item should specify an attribute that `hook` may change.
                 specification format is identical to `depends_on`.
 
+        Raises:
+            ValueError: On failure of specification validation.
+        """
+        # Ignore master record for this variant of the hook
+        hook_name = get_func_name(hook)
+        self._snap_shooter.register_correlation_hook(lambda e, s, _m: hook(e, s), entity_type, depends_on, may_change, hook_name=hook_name)
+
+    def register_correlation_hook_with_master_record(
+        self,
+        hook: Callable[[str, dict, dict], Union[None, list[DataPointTask]]],
+        entity_type: str,
+        depends_on: list[list[str]],
+        may_change: list[list[str]],
+    ):
+        """
+        Registers passed hook to be called during snapshot creation.
+
+        Identical to `register_correlation_hook`, but the hook also receives the master record.
+
+        Binds hook to specified entity_type (though same hook can be bound multiple times).
+
+        `entity_type` and attribute specifications are validated, `ValueError` is raised on failure.
+
+        Args:
+            hook: `hook` callable should expect entity type as str;
+                its current values, including linked entities, as dict;
+                and its master record as dict.
+                Can optionally return a list of DataPointTask objects to perform.
+            entity_type: specifies entity type
+            depends_on: each item should specify an attribute that is depended on
+                in the form of a path from the specified entity_type to individual attributes
+                (even on linked entities).
+            may_change: each item should specify an attribute that `hook` may change.
+                specification format is identical to `depends_on`.
+
         Raises:
             ValueError: On failure of specification validation.
         """

dp3/snapshots/snapshooter.py

@@ -185,33 +185,39 @@ def register_timeseries_hook(
 
     def register_correlation_hook(
         self,
-        hook: Callable[[str, dict], Union[None, list[DataPointTask]]],
+        hook: Callable[[str, dict, dict], Union[None, list[DataPointTask]]],
         entity_type: str,
         depends_on: list[list[str]],
         may_change: list[list[str]],
+        hook_name: Optional[str] = None,
     ):
         """
         Registers passed hook to be called during snapshot creation.
 
+        Common implementation for hooks with and without master record.
+
         Binds hook to specified entity_type (though same hook can be bound multiple times).
 
         `entity_type` and attribute specifications are validated, `ValueError` is raised on failure.
 
         Args:
             hook: `hook` callable should expect entity type as str
-                and its current values, including linked entities, as dict
+                and its current values, including linked entities, as dict;
+                and its master record as dict.
                 Can optionally return a list of DataPointTask objects to perform.
             entity_type: specifies entity type
             depends_on: each item should specify an attribute that is depended on
                 in the form of a path from the specified entity_type to individual attributes
                 (even on linked entities).
             may_change: each item should specify an attribute that `hook` may change.
                 specification format is identical to `depends_on`.
+            hook_name: Optional custom name for the hook, used for logging purposes. If not
+                provided, the function name of `hook` will be used.
 
         Raises:
             ValueError: On failure of specification validation.
         """
-        self._correlation_hooks.register(hook, entity_type, depends_on, may_change)
+        self._correlation_hooks.register(hook, entity_type, depends_on, may_change, hook_name=hook_name)
 
     def register_run_init_hook(self, hook: Callable[[], list[DataPointTask]]):
         """
@@ -457,9 +463,10 @@ def make_linkless_snapshot(self, entity_type: str, master_record: dict, time: da
         self.run_timeseries_processing(entity_type, master_record)
         values = self.get_values_at_time(entity_type, master_record, time)
         self.add_mirrored_links(entity_type, values)
-        entity_values = {(entity_type, master_record["_id"]): values}
+        entity_id = master_record["_id"]
+        entity_values = {(entity_type, entity_id): values}
 
-        tasks = self._correlation_hooks.run(entity_values)
+        tasks = self._correlation_hooks.run(entity_values, {(entity_type, entity_id): master_record})
         for task in tasks:
             self.task_queue_writer.put_task(task)
 
@@ -499,6 +506,7 @@ def make_snapshot(self, task: Snapshot):
         The resulting snapshots are saved into DB.
         """
         entity_values = {}
+        entity_master_records = {}
         for entity_type, entity_id in task.entities:
             record = self.db.get_master_record(entity_type, entity_id) or {"_id": entity_id}
             if not self.config.keep_empty and len(record) == 1:
@@ -508,9 +516,10 @@ def make_snapshot(self, task: Snapshot):
             values = self.get_values_at_time(entity_type, record, task.time)
             self.add_mirrored_links(entity_type, values)
             entity_values[entity_type, entity_id] = values
+            entity_master_records[entity_type, entity_id] = record
 
         self.link_loaded_entities(entity_values)
-        created_tasks = self._correlation_hooks.run(entity_values)
+        created_tasks = self._correlation_hooks.run(entity_values, entity_master_records)
         for created_task in created_tasks:
             self.task_queue_writer.put_task(created_task)
 

dp3/snapshots/snapshot_hooks.py

@@ -84,10 +84,11 @@ def __init__(self, log: logging.Logger, model_spec: ModelSpec, elog: EventGroupT
 
     def register(
         self,
-        hook: Callable[[str, dict], Union[None, list[DataPointTask]]],
+        hook: Callable[[str, dict, dict], Union[None, list[DataPointTask]]],
         entity_type: str,
         depends_on: list[list[str]],
         may_change: list[list[str]],
+        hook_name: Union[None, str] = None,
     ) -> str:
         """
         Registers passed hook to be called during snapshot creation.
@@ -97,15 +98,18 @@ def register(
         If entity_type and attribute specifications are validated
         and ValueError is raised on failure.
         Args:
-            hook: `hook` callable should expect entity type as str
-                and its current values, including linked entities, as dict.
+            hook: `hook` callable should expect entity type as str;
+                its current values, including linked entities, as dict;
+                and its master record as dict.
                 Can optionally return a list of DataPointTask objects to perform.
             entity_type: specifies entity type
             depends_on: each item should specify an attribute that is depended on
                 in the form of a path from the specified entity_type to individual attributes
                 (even on linked entities).
             may_change: each item should specify an attribute that `hook` may change.
                 specification format is identical to `depends_on`.
+            hook_name: Optional custom name for the hook, used for logging purposes. If not
+                provided, the function name of `hook` will be used.
         Returns:
             Generated hook id.
         """
@@ -120,7 +124,8 @@ def register(
         may_change = self._get_attr_path_destinations(entity_type, may_change)
 
         hook_args = f"({entity_type}, [{','.join(depends_on)}], [{','.join(may_change)}])"
-        hook_id = f"{get_func_name(hook)}{hook_args}"
+        hook_name = hook_name if hook_name is not None else get_func_name(hook)
+        hook_id = f"{hook_name}{hook_args}"
         self._short_hook_ids[hook_id] = hook_args
         self._dependency_graph.add_hook_dependency(hook_id, depends_on, may_change)
 
@@ -191,7 +196,7 @@ def _resolve_entities_in_path(self, base_entity: str, path: list[str]) -> list[t
                 position = entity_attributes[position.relation_to]
         return resolved_path
 
-    def run(self, entities: dict) -> list[DataPointTask]:
+    def run(self, entities: dict, entity_master_records: dict) -> list[DataPointTask]:
         """Runs registered hooks."""
         entity_types = {etype for etype, _ in entities}
         hook_subset = [
@@ -200,18 +205,22 @@ def run(self, entities: dict) -> list[DataPointTask]:
         topological_order = self._dependency_graph.topological_order
         hook_subset.sort(key=lambda x: topological_order.index(x[0]))
         entities_by_etype = defaultdict(dict)
+        entity_master_records_by_etype = defaultdict(dict)
         for (etype, eid), values in entities.items():
             entities_by_etype[etype][eid] = values
+        for (etype, eid), mr in entity_master_records.items():
+            entity_master_records_by_etype[etype][eid] = mr
 
         created_tasks = []
 
         with task_context(self.model_spec):
             for hook_id, hook, etype in hook_subset:
                 short_id = hook_id if len(hook_id) < 160 else self._short_hook_ids[hook_id]
                 for eid, entity_values in entities_by_etype[etype].items():
+                    entity_master_record = entity_master_records_by_etype[etype].get(eid, {})
                     self.log.debug("Running hook %s on entity %s", short_id, eid)
                     try:
-                        tasks = hook(etype, entity_values)
+                        tasks = hook(etype, entity_values, entity_master_record)
                         if tasks is not None and tasks:
                             created_tasks.extend(tasks)
                     except Exception as e:

tests/test_common/test_snapshots.py

@@ -14,7 +14,7 @@
 from dp3.snapshots.snapshot_hooks import SnapshotCorrelationHookContainer
 
 
-def modify_value(_: str, record: dict, attr: str, value):
+def modify_value(_: str, record: dict, _master_record: dict, attr: str, value):
     record[attr] = value
 
 
@@ -37,7 +37,7 @@ def test_basic_function(self):
             hook=dummy_hook_abc, entity_type="A", depends_on=[["data1"]], may_change=[["data2"]]
         )
         values = {}
-        self.container.run({("A", "a1"): values})
+        self.container.run({("A", "a1"): values}, {})
         self.assertEqual(values["data2"], "abc")
 
     def test_circular_dependency_error(self):