Batch-hydrate all nodes and CREs (N+1, supersedes #850)

application/database/db.py

@@ -1502,17 +1502,10 @@ def __get_unlinked_cres(self) -> List[CRE]:
     def __get_all_nodes_and_cres(
         self, cres_only: bool = False
     ) -> List[cre_defs.Document]:
-        result = []
-        nodes = []
-        cres = []
+        result: List[cre_defs.Document] = []
         if not cres_only:
-            node_ids = self.session.query(Node.id).all()
-            for nid in node_ids:
-                result.extend(self.get_nodes(db_id=nid[0]))
-
-        cre_ids = self.session.query(CRE.id).all()
-        for cid in cre_ids:
-            result.append(self.get_cre_by_db_id(cid[0]))
+            result.extend(self._hydrate_nodes_batch(self.session.query(Node).all()))
+        result.extend(self._hydrate_cres_batch(self.session.query(CRE).all()))
         return result
 
     @classmethod
@@ -1945,6 +1938,53 @@ def get_CREs(
 
         return self._hydrate_cres_batch(dbcres, include_only_nodes=include_only)
 
+    def _hydrate_nodes_batch(
+        self,
+        dbnodes: List[Node],
+        include_only: Optional[List[str]] = None,
+    ) -> List[cre_defs.Node]:
+        """Hydrate Standard/Tool/Code nodes and their CRE links without per-row queries."""
+        if not dbnodes:
+            return []
+
+        node_ids = [node.id for node in dbnodes]
+        links_by_node: Dict[str, List[Links]] = defaultdict(list)
+        cre_ids: set = set()
+        for link in self.session.query(Links).filter(Links.node.in_(node_ids)).all():
+            links_by_node[link.node].append(link)
+            cre_ids.add(link.cre)
+
+        cres_by_id: Dict[str, CRE] = {}
+        if cre_ids:
+            cres_by_id = {
+                cre.id: cre
+                for cre in self.session.query(CRE).filter(CRE.id.in_(cre_ids)).all()
+            }
+
+        nodes: List[cre_defs.Node] = []
+        for dbnode in dbnodes:
+            node = nodeFromDB(dbnode=dbnode)
+            for dbcre_link in links_by_node.get(dbnode.id, []):
+                dbcre = cres_by_id.get(dbcre_link.cre)
+                if not dbcre:
+                    logger.fatal(
+                        f"CRE {dbcre_link.cre} exists in the links but not in the cre table, database corrupt?"
+                    )
+                    raise AssertionError(
+                        f"CRE {dbcre_link.cre} exists in the links but not in the cre table, database corrupt?"
+                    )
+                if not include_only or (
+                    dbcre.external_id in include_only or dbcre.name in include_only
+                ):
+                    node.add_link(
+                        cre_defs.Link(
+                            ltype=cre_defs.LinkTypes.from_str(dbcre_link.type),
+                            document=CREfromDB(dbcre),
+                        )
+                    )
+            nodes.append(node)
+        return nodes
+
     def _hydrate_cres_batch(
         self,
         dbcres: List[CRE],

application/tests/db_test.py

@@ -812,6 +812,40 @@ def test_get_cre_by_db_id_returns_none_when_get_cres_empty(
         self.assertIsNone(result)
         get_cres_mock.assert_called_once_with(external_id=self.dbcre.external_id)
 
+    def test_get_all_nodes_and_cres_does_not_use_per_id_lookups(self) -> None:
+        with patch.object(self.collection, "get_cre_by_db_id") as per_cre:
+            with patch.object(self.collection, "get_nodes") as per_node:
+                docs = self.collection._Node_collection__get_all_nodes_and_cres()
+        per_cre.assert_not_called()
+        per_node.assert_not_called()
+        self.assertGreaterEqual(len(docs), 4)  # 2 CREs + BarStand + Unlinked
+
+    def test_get_all_nodes_and_cres_cres_only_skips_nodes(self) -> None:
+        docs = self.collection._Node_collection__get_all_nodes_and_cres(cres_only=True)
+        self.assertTrue(all(d.doctype == defs.Credoctypes.CRE for d in docs))
+        self.assertCountEqual(["CREname", "GroupName"], [d.name for d in docs])
+        cre = next(d for d in docs if d.name == "CREname")
+        self.assertTrue(
+            any(
+                link.document.name == "BarStand"
+                for link in cre.links
+                if link.document.doctype != defs.Credoctypes.CRE
+            )
+        )
+
+    def test_get_all_nodes_and_cres_hydrates_same_as_per_id(self) -> None:
+        batch = self.collection._Node_collection__get_all_nodes_and_cres()
+        batch_cres = {d.id: d for d in batch if d.doctype == defs.Credoctypes.CRE}
+        for db_id in [row[0] for row in self.collection.session.query(db.CRE.id).all()]:
+            per_id = self.collection.get_cre_by_db_id(db_id)
+            self.assertIsNotNone(per_id)
+            self.assertEqual(per_id.todict(), batch_cres[per_id.id].todict())
+        batch_nodes = [d for d in batch if d.doctype != defs.Credoctypes.CRE]
+        self.assertTrue(any(n.name == "BarStand" for n in batch_nodes))
+        self.assertTrue(any(n.name == "Unlinked" for n in batch_nodes))
+        bar = next(n for n in batch_nodes if n.name == "BarStand")
+        self.assertTrue(any(link.document.name == "CREname" for link in bar.links))
+
     def test_get_standards(self) -> None:
         """Given: a Standard 'S1' that links to cres
         return the Standard in Document format"""