diff --git a/application/database/db.py b/application/database/db.py index 3101a9ca7..18e5d2bf0 100644 --- a/application/database/db.py +++ b/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], diff --git a/application/tests/db_test.py b/application/tests/db_test.py index 1d6ed7ead..68362cfc2 100644 --- a/application/tests/db_test.py +++ b/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"""