Merge pull request #191 from blarApp/dev

Dev

Author: v4rgas

blarify/code_hierarchy/languages/go_definitions.py

@@ -24,23 +24,13 @@ def get_parsers_for_extensions() -> Dict[str, Parser]:
     def should_create_node(node: Node) -> bool:
         return LanguageDefinitions._should_create_node_base_implementation(
             node,
-            ["type_declaration", "method_declaration", "function_declaration"],
+            ["type_spec", "type_alias", "method_declaration", "function_declaration"],
         )
 
     def get_identifier_node(node: Node) -> Node:
-        if node.type == "type_declaration":
-            for child in node.named_children:
-                if child.type == "type_spec" or "type_alias":
-                    node = child
-
         return LanguageDefinitions._get_identifier_node_base_implementation(node)
 
     def get_body_node(node: Node) -> Node:
-        if node.type == "type_declaration":
-            for child in node.named_children:
-                if child.type == "type_spec":
-                    node = child
-
         return LanguageDefinitions._get_body_node_base_implementation(node)
 
     def get_relationship_type(node: GraphNode, node_in_point_reference: Node) -> Optional[FoundRelationshipScope]:
@@ -51,7 +41,8 @@ def get_relationship_type(node: GraphNode, node_in_point_reference: Node) -> Opt
 
     def get_node_label_from_type(type: str) -> NodeLabels:
         return {
-            "type_declaration": NodeLabels.CLASS,
+            "type_spec": NodeLabels.CLASS,
+            "type_alias": NodeLabels.CLASS,
             "method_declaration": NodeLabels.FUNCTION,
             "function_declaration": NodeLabels.FUNCTION,
         }[type]

blarify/db_managers/db_manager.py

@@ -0,0 +1,20 @@
+class AbstractDbManager:
+    def close(self):
+        """Close the connection to the database."""
+        raise NotImplementedError
+
+    def save_graph(self, nodes, edges):
+        """Save nodes and edges to the database."""
+        raise NotImplementedError
+
+    def create_nodes(self, nodeList):
+        """Create nodes in the database."""
+        raise NotImplementedError
+
+    def create_edges(self, edgesList):
+        """Create edges between nodes in the database."""
+        raise NotImplementedError
+
+    def detatch_delete_nodes_with_path(self, path):
+        """Detach and delete nodes matching the given path."""
+        raise NotImplementedError

blarify/db_managers/falkordb_manager.py

@@ -0,0 +1,81 @@
+import os
+import time
+from typing import Any, List
+import logging
+
+from dotenv import load_dotenv
+from falkordb import FalkorDB, exceptions
+
+logger = logging.getLogger(__name__)
+
+load_dotenv()
+
+
+class FalkorDBManager:
+    entity_id: str
+    repo_id: str
+    db: FalkorDB
+
+    def __init__(
+        self,
+        repo_id: str = None,
+        entity_id: str = None,
+        uri: str = None,
+        user: str = None,
+        password: str = None,
+    ):
+        host = uri or os.getenv("FALKORDB_URI", "localhost")
+        port = int(os.getenv("FALKORDB_PORT", 6379))
+        user = user or os.getenv("FALKORDB_USERNAME")
+        password = password or os.getenv("FALKORDB_PASSWORD")
+
+        self.db = FalkorDB(host=host, port=port, username=user, password=password)
+
+        self.repo_id = repo_id if repo_id is not None else "default_repo"
+        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()
+
+    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
+        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})
+
+    def create_edges(self, edgesList: List[Any]):
+        # Function to create edges between nodes in the FalkorDB database
+        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"],
+                },
+            )
+
+    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

blarify/examples/graph_builder.py

@@ -1,5 +1,6 @@
 from blarify.prebuilt.graph_builder import GraphBuilder
 from blarify.db_managers.neo4j_manager import Neo4jManager
+from blarify.db_managers.falkordb_manager import FalkorDBManager
 
 import dotenv
 import os
@@ -12,7 +13,8 @@ def build(root_path: str = None):
     relationships = graph.get_relationships_as_objects()
     nodes = graph.get_nodes_as_objects()
 
-    save_to_neo4j(relationships, nodes)
+    # save_to_neo4j(relationships, nodes)
+    save_to_falkordb(relationships, nodes)
 
 
 def save_to_neo4j(relationships, nodes):
@@ -23,6 +25,14 @@ def save_to_neo4j(relationships, nodes):
     graph_manager.close()
 
 
+def save_to_falkordb(relationships, nodes):
+    graph_manager = FalkorDBManager(repo_id="repo", entity_id="organization")
+
+    print(f"Saving graph with {len(nodes)} nodes and {len(relationships)} relationships")
+    graph_manager.save_graph(nodes, relationships)
+    graph_manager.close()
+
+
 if __name__ == "__main__":
     import logging
 

poetry.lock

@@ -23,6 +23,7 @@ tree-sitter-typescript = "^0.23.2"
 tree-sitter-c-sharp = "^0.23.1"
 multilspy = {git = "https://github.com/blarApp/multilspy.git"}
 tree-sitter-go = "^0.23.4"
+falkordb = "^1.0.10"
 
 
 [build-system]