diff --git a/blarify/db_managers/falkordb_manager.py b/blarify/db_managers/falkordb_manager.py index 6b65eebc..77fef097 100644 --- a/blarify/db_managers/falkordb_manager.py +++ b/blarify/db_managers/falkordb_manager.py @@ -35,47 +35,42 @@ def __init__( self.entity_id = entity_id if entity_id is not None else "default_user" def close(self): - # Close the connection to the database - self.db.close() + pass def save_graph(self, nodes: List[Any], edges: List[Any]): self.create_nodes(nodes) self.create_edges(edges) - def create_nodes(self, nodeList: List[Any]): - # Function to create nodes in the FalkorDB database + def create_nodes(self, nodeList: List[dict]): graph = self.db.select_graph(self.repo_id) - for node in nodeList: - labels = ":".join(node.get("extra_labels", []) + [node["type"], "NODE"]) - attributes = node.get("attributes", {}) - attributes.update({"repoId": self.repo_id, "entityId": self.entity_id}) - # Construct parameterized query - cypher_query = f"CREATE (n:{labels} $props)" - graph.query(cypher_query, params={"props": attributes}) + cypher_query = """ + UNWIND $nodes AS node + CREATE (n) + SET n = node.attributes + WITH n AS created_node, node.extra_labels AS labels + UNWIND labels AS label + SET created_node:label + RETURN created_node + """ + graph.query( + cypher_query, + params={"nodes": nodeList}, + ) - def create_edges(self, edgesList: List[Any]): - # Function to create edges between nodes in the FalkorDB database + def create_edges(self, edgesList: List[dict]): graph = self.db.select_graph(self.repo_id) - for edge in edgesList: - # Construct parameterized query - cypher_query = ( - "MATCH (a:NODE {node_id: $sourceId}), " - "(b:NODE {node_id: $targetId}) " - "CREATE (a)-[r:$type {scopeText: $scopeText}]->(b)" - ) - graph.query( - cypher_query, - params={ - "sourceId": edge["sourceId"], - "targetId": edge["targetId"], - "type": edge["type"], - "scopeText": edge["scopeText"], - }, - ) + cypher_query = """ + UNWIND $edges AS edge + MATCH (a {node_id: edge.sourceId}), (b {node_id: edge.targetId}) + CREATE (a)-[:`$edge.type` {scopeText: edge.scopeText}]->(b) + """ + graph.query( + cypher_query, + params={"edges": edgesList}, + ) def detach_delete_nodes_with_path(self, path: str): graph = self.db.select_graph(self.repo_id) - # Construct parameterized query cypher_query = "MATCH (n {path: $path}) DETACH DELETE n" result = graph.query(cypher_query, params={"path": path}) return result.result_set