Links: upsert iterable observation targets

MongoDB cannot expand an upserting $in filter into one cache document per target. Iterable observation links could therefore be missing from the cache and escape deletion cleanup.

Use one exact-key UpdateOne upsert per distinct target in an unordered bulk write, and add real-MongoDB coverage for empty and partially populated caches.

Author: xsedla1o

dp3/core/link_manager.py

@@ -113,15 +113,24 @@ def add_observation_to_link_cache(
     def add_iterable_observation_to_link_cache(
         self, etype_to: str, post_validity: timedelta, eid: AnyEidT, dp: DataPointObservationsBase
     ):
-        linked_eids = [v.eid for v in dp.v]
-        self.cache.update_many(
-            {
-                "to": {"$in": [f"{etype_to}#{eid_}" for eid_ in linked_eids]},
-                "from": f"{dp.etype}#{eid}",
-                "using_attr": f"{dp.etype}#{dp.attr}",
-            },
-            {"$max": {"ttl": dp.t2 + post_validity}},
-            upsert=True,
+        targets = {f"{etype_to}#{value.eid}" for value in dp.v}
+        if not targets:
+            return
+
+        link_key = {
+            "from": f"{dp.etype}#{eid}",
+            "using_attr": f"{dp.etype}#{dp.attr}",
+        }
+        self.cache.bulk_write(
+            [
+                UpdateOne(
+                    link_key | {"to": target},
+                    {"$max": {"ttl": dp.t2 + post_validity}},
+                    upsert=True,
+                )
+                for target in targets
+            ],
+            ordered=False,
         )
 
     def remove_link_cache_of_deleted(self, etype: str, eid: AnyEidT):

tests/test_api/test_iterable_observation_links.py

@@ -0,0 +1,63 @@
+import unittest
+from datetime import UTC, datetime, timedelta
+from types import SimpleNamespace
+
+import common
+
+from dp3.core.link_manager import LinkManager
+from dp3.database.config import MongoConfig
+from dp3.database.database import EntityDatabase
+from dp3.database.encodings import get_codec_options
+
+
+class IterableObservationLinks(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        db_config = MongoConfig.model_validate(common.CONFIG.get("database", {}))
+        cls.client = EntityDatabase.connect(db_config)
+        cls.client.admin.command("ping")
+        cls.db = cls.client.get_database(db_config.db_name, codec_options=get_codec_options())
+        cls.cache = cls.db["#cache#IterableObservationLinkRegression"]
+
+        cls.link_manager = LinkManager.__new__(LinkManager)
+        cls.link_manager.cache = cls.cache
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.cache.drop()
+        cls.client.close()
+
+    def setUp(self):
+        self.cache.delete_many({})
+
+    @staticmethod
+    def _datapoint(*targets):
+        return SimpleNamespace(
+            etype="source",
+            attr="observed_targets",
+            v=[SimpleNamespace(eid=target) for target in targets],
+            t2=datetime(2025, 1, 1, tzinfo=UTC),
+        )
+
+    def _add_links(self, *targets):
+        self.link_manager.add_iterable_observation_to_link_cache(
+            "target", timedelta(hours=1), "source-1", self._datapoint(*targets)
+        )
+
+    def _cached_links(self):
+        return list(self.cache.find({}, {"_id": 0, "to": 1}))
+
+    def test_empty_cache_creates_one_link_per_target(self):
+        self._add_links("a", "a", "b")
+
+        self.assertCountEqual(self._cached_links(), [{"to": "target#a"}, {"to": "target#b"}])
+
+    def test_existing_links_do_not_prevent_new_target_upserts(self):
+        self._add_links("a")
+        self._add_links("a", "b")
+
+        self.assertCountEqual(self._cached_links(), [{"to": "target#a"}, {"to": "target#b"}])
+
+
+if __name__ == "__main__":
+    unittest.main()