From ce4b2ef977b712e4d2c976982165f4a504e09718 Mon Sep 17 00:00:00 2001 From: berrazuriz1 Date: Mon, 25 May 2026 15:52:56 -0400 Subject: [PATCH] Batch node creation to avoid oversized Bolt transactions --- blarify/db_managers/neo4j_manager.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/blarify/db_managers/neo4j_manager.py b/blarify/db_managers/neo4j_manager.py index e3c67d4f..5360c3b8 100644 --- a/blarify/db_managers/neo4j_manager.py +++ b/blarify/db_managers/neo4j_manager.py @@ -52,11 +52,19 @@ def save_graph(self, nodes: List[Any], edges: List[Any]): self.create_edges(edges) def create_nodes(self, nodeList: List[Any]): - # Function to create nodes in the Neo4j database + batch_size: int = 5000 + total_nodes: int = len(nodeList) + total_batches: int = (total_nodes + batch_size - 1) // batch_size + logger.info(f"Creating {total_nodes} nodes in batches of {batch_size}") + with self.driver.session() as session: - session.write_transaction( - self._create_nodes_txn, nodeList, 100, repoId=self.repo_id, entityId=self.entity_id - ) + for i in range(0, total_nodes, batch_size): + batch: List[Any] = nodeList[i:i + batch_size] + batch_num: int = i // batch_size + 1 + logger.info(f"Processing nodes batch {batch_num}/{total_batches} ({i}/{total_nodes})") + session.write_transaction( + self._create_nodes_txn, batch, 100, repoId=self.repo_id, entityId=self.entity_id + ) def create_edges(self, edgesList: List[Any]): # Function to create edges between nodes in the Neo4j database