code-graph-builder 0.2.0__py3-none-any.whl

code_graph_builder/tools/graph_query.py

@@ -0,0 +1,495 @@
+"""Graph query layer for unified access to Kuzu and Memgraph backends.
+
+This module provides a unified interface for querying graph data from
+different backends (Memgraph and Kuzu), enabling seamless integration
+with vector store search results.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
+
+from loguru import logger
+
+if TYPE_CHECKING:
+    from ..services import IngestorProtocol, QueryProtocol
+    from ..types import ResultRow
+
+
+@dataclass
+class GraphNode:
+    """Represents a node in the code graph.
+
+    Attributes:
+        node_id: Unique node identifier
+        qualified_name: Fully qualified name (e.g., "module.Class.method")
+        name: Simple name (e.g., "method")
+        type: Node type (Function, Class, Method, Module, etc.)
+        path: File path
+        start_line: Start line number in file
+        end_line: End line number in file
+        docstring: Documentation string if available
+        properties: Additional node properties
+    """
+
+    node_id: int
+    qualified_name: str
+    name: str
+    type: str
+    path: str | None = None
+    start_line: int | None = None
+    end_line: int | None = None
+    docstring: str | None = None
+    properties: dict | None = None
+
+
+@dataclass
+class GraphRelationship:
+    """Represents a relationship in the code graph.
+
+    Attributes:
+        rel_type: Relationship type (CALLS, DEFINES, INHERITS, etc.)
+        source_id: Source node ID
+        target_id: Target node ID
+        properties: Relationship properties
+    """
+
+    rel_type: str
+    source_id: int
+    target_id: int
+    properties: dict | None = None
+
+
+@runtime_checkable
+class GraphQueryProtocol(Protocol):
+    """Protocol for graph query operations."""
+
+    def fetch_nodes_by_ids(self, node_ids: list[int]) -> list[GraphNode]: ...
+
+    def fetch_node_by_qualified_name(self, qualified_name: str) -> GraphNode | None: ...
+
+    def fetch_callers(self, function_name: str) -> list[GraphNode]: ...
+
+    def fetch_callees(self, function_name: str) -> list[GraphNode]: ...
+
+    def fetch_related_nodes(
+        self, node_id: int, relationship_types: list[str] | None = None
+    ) -> list[tuple[GraphNode, str]]: ...
+
+    def execute_cypher(self, query: str, params: dict | None = None) -> list[ResultRow]: ...
+
+
+class GraphQueryService:
+    """Unified service for querying code graph data.
+
+    Supports both Memgraph and Kuzu backends through a common interface.
+
+    Example:
+        >>> from code_graph_builder.services import MemgraphIngestor
+        >>> from code_graph_builder.tools.graph_query import GraphQueryService
+        >>>
+        >>> with MemgraphIngestor("localhost", 7687) as ingestor:
+        ...     query_service = GraphQueryService(ingestor)
+        ...     node = query_service.fetch_node_by_qualified_name("myproject.utils.foo")
+        ...     callers = query_service.fetch_callers("foo")
+    """
+
+    def __init__(self, graph_service: QueryProtocol, backend: str = "memgraph"):
+        """Initialize graph query service.
+
+        Args:
+            graph_service: Graph service instance (MemgraphIngestor or KuzuIngestor)
+            backend: Backend type ("memgraph" or "kuzu")
+        """
+        self.graph_service = graph_service
+        self.backend = backend.lower()
+
+    def fetch_nodes_by_ids(self, node_ids: list[int]) -> list[GraphNode]:
+        """Fetch multiple nodes by their IDs.
+
+        Args:
+            node_ids: List of node identifiers
+
+        Returns:
+            List of GraphNode objects
+        """
+        if not node_ids:
+            return []
+
+        query = self._build_nodes_by_id_query()
+
+        try:
+            results = self.graph_service.fetch_all(query, {"node_ids": node_ids})
+            return [self._row_to_node(row) for row in results if self._extract_node_id(row) in node_ids]
+        except Exception as e:
+            logger.error(f"Failed to fetch nodes by IDs: {e}")
+            return []
+
+    def fetch_node_by_qualified_name(self, qualified_name: str) -> GraphNode | None:
+        """Fetch a single node by its qualified name.
+
+        Args:
+            qualified_name: Fully qualified name (e.g., "module.Class.method")
+
+        Returns:
+            GraphNode if found, None otherwise
+        """
+        query = """
+            MATCH (n)
+            WHERE n.qualified_name = $qualified_name
+            RETURN n,
+                   n.node_id AS node_id,
+                   n.id AS id,
+                   n.qualified_name AS qualified_name,
+                   n.name AS name,
+                   labels(n) AS labels,
+                   n.path AS path,
+                   n.start_line AS start_line,
+                   n.end_line AS end_line,
+                   n.docstring AS docstring
+            LIMIT 1
+        """
+
+        try:
+            results = self.graph_service.fetch_all(query, {"qualified_name": qualified_name})
+            if results:
+                return self._row_to_node(results[0])
+        except Exception as e:
+            logger.error(f"Failed to fetch node {qualified_name}: {e}")
+
+        return None
+
+    def fetch_callers(self, function_name: str) -> list[GraphNode]:
+        """Find all functions that call the given function.
+
+        Args:
+            function_name: Function name or qualified name
+
+        Returns:
+            List of caller GraphNodes
+        """
+        # Try qualified name first
+        query = """
+            MATCH (caller:Function)-[:CALLS]->(callee)
+            WHERE callee.qualified_name = $name
+               OR callee.name = $name
+            RETURN caller,
+                   caller.node_id AS node_id,
+                   caller.id AS id,
+                   caller.qualified_name AS qualified_name,
+                   caller.name AS name,
+                   labels(caller) AS labels,
+                   caller.path AS path,
+                   caller.start_line AS start_line,
+                   caller.end_line AS end_line
+        """
+
+        try:
+            results = self.graph_service.fetch_all(query, {"name": function_name})
+            return [self._row_to_node(row) for row in results]
+        except Exception as e:
+            logger.error(f"Failed to fetch callers of {function_name}: {e}")
+            return []
+
+    def fetch_callees(self, function_name: str) -> list[GraphNode]:
+        """Find all functions called by the given function.
+
+        Args:
+            function_name: Function name or qualified name
+
+        Returns:
+            List of callee GraphNodes
+        """
+        query = """
+            MATCH (caller:Function)-[:CALLS]->(callee:Function)
+            WHERE caller.qualified_name = $name
+               OR caller.name = $name
+            RETURN callee,
+                   callee.node_id AS node_id,
+                   callee.id AS id,
+                   callee.qualified_name AS qualified_name,
+                   callee.name AS name,
+                   labels(callee) AS labels,
+                   callee.path AS path,
+                   callee.start_line AS start_line,
+                   callee.end_line AS end_line
+        """
+
+        try:
+            results = self.graph_service.fetch_all(query, {"name": function_name})
+            return [self._row_to_node(row) for row in results]
+        except Exception as e:
+            logger.error(f"Failed to fetch callees of {function_name}: {e}")
+            return []
+
+    def fetch_related_nodes(
+        self,
+        node_id: int,
+        relationship_types: list[str] | None = None,
+        direction: str = "both",
+    ) -> list[tuple[GraphNode, str]]:
+        """Fetch nodes related to the given node.
+
+        Args:
+            node_id: Node identifier
+            relationship_types: Optional filter for relationship types
+            direction: Relationship direction ("in", "out", or "both")
+
+        Returns:
+            List of (GraphNode, relationship_type) tuples
+        """
+        if direction == "in":
+            pattern = "(related)-[r]->(n)"
+        elif direction == "out":
+            pattern = "(n)-[r]->(related)"
+        else:
+            pattern = "(n)-[r]-(related)"
+
+        rel_filter = ""
+        if relationship_types:
+            rel_types = "|".join(f":{rt}" for rt in relationship_types)
+            rel_filter = f"AND type(r) IN {relationship_types}"
+
+        query = f"""
+            MATCH {pattern}
+            WHERE n.node_id = $node_id
+               OR n.id = $node_id
+               OR id(n) = $node_id
+            {rel_filter}
+            RETURN related,
+                   related.node_id AS node_id,
+                   related.id AS id,
+                   related.qualified_name AS qualified_name,
+                   related.name AS name,
+                   labels(related) AS labels,
+                   related.path AS path,
+                   related.start_line AS start_line,
+                   related.end_line AS end_line,
+                   type(r) AS rel_type
+        """
+
+        try:
+            results = self.graph_service.fetch_all(query, {"node_id": node_id})
+            return [
+                (self._row_to_node(row), str(row.get("rel_type", "UNKNOWN")))
+                for row in results
+            ]
+        except Exception as e:
+            logger.error(f"Failed to fetch related nodes for {node_id}: {e}")
+            return []
+
+    def fetch_class_hierarchy(self, class_name: str) -> dict:
+        """Fetch class hierarchy information.
+
+        Args:
+            class_name: Class name or qualified name
+
+        Returns:
+            Dictionary with superclass and subclasses
+        """
+        query = """
+            MATCH (c:Class)
+            WHERE c.qualified_name = $name OR c.name = $name
+            OPTIONAL MATCH (c)-[:INHERITS]->(super:Class)
+            OPTIONAL MATCH (sub:Class)-[:INHERITS]->(c)
+            RETURN c,
+                   super.qualified_name AS superclass,
+                   collect(sub.qualified_name) AS subclasses
+        """
+
+        try:
+            results = self.graph_service.fetch_all(query, {"name": class_name})
+            if results:
+                return {
+                    "class": results[0].get("c"),
+                    "superclass": results[0].get("superclass"),
+                    "subclasses": results[0].get("subclasses", []),
+                }
+        except Exception as e:
+            logger.error(f"Failed to fetch class hierarchy for {class_name}: {e}")
+
+        return {}
+
+    def execute_cypher(self, query: str, params: dict | None = None) -> list[ResultRow]:
+        """Execute a raw Cypher query.
+
+        Args:
+            query: Cypher query string
+            params: Query parameters
+
+        Returns:
+            Query results as list of dictionaries
+        """
+        try:
+            return self.graph_service.fetch_all(query, params or {})
+        except Exception as e:
+            logger.error(f"Cypher query failed: {e}")
+            return []
+
+    def _build_nodes_by_id_query(self) -> str:
+        """Build query to fetch nodes by IDs.
+
+        Compatible with both Memgraph and Kuzu.
+        """
+        return """
+            MATCH (n)
+            WHERE n.node_id IN $node_ids
+               OR n.id IN $node_ids
+               OR id(n) IN $node_ids
+            RETURN n,
+                   n.node_id AS node_id,
+                   n.id AS id,
+                   n.qualified_name AS qualified_name,
+                   n.name AS name,
+                   labels(n) AS labels,
+                   n.path AS path,
+                   n.start_line AS start_line,
+                   n.end_line AS end_line,
+                   n.docstring AS docstring
+        """
+
+    def _extract_node_id(self, row: ResultRow) -> int:
+        """Extract node ID from query result."""
+        for key in ["node_id", "id", "n.node_id", "n.id"]:
+            if key in row:
+                val = row[key]
+                if isinstance(val, int):
+                    return val
+                try:
+                    return int(val)
+                except (ValueError, TypeError):
+                    continue
+        return 0
+
+    def _extract_type(self, row: ResultRow) -> str:
+        """Extract node type from query result."""
+        # Try labels first
+        labels = row.get("labels")
+        if labels:
+            if isinstance(labels, list) and labels:
+                return labels[0]
+            return str(labels)
+
+        # Try type property
+        node_type = row.get("type")
+        if node_type:
+            return str(node_type)
+
+        # Extract from node object if available
+        node = row.get("n") or row.get("caller") or row.get("callee") or row.get("related")
+        if node and isinstance(node, dict):
+            node_labels = node.get("_label") or node.get("labels")
+            if node_labels:
+                if isinstance(node_labels, list) and node_labels:
+                    return node_labels[0]
+                return str(node_labels)
+            node_type = node.get("type")
+            if node_type:
+                return str(node_type)
+
+        return "Unknown"
+
+    def _row_to_node(self, row: ResultRow) -> GraphNode:
+        """Convert query result row to GraphNode."""
+        node_id = self._extract_node_id(row)
+        qualified_name = str(row.get("qualified_name", ""))
+        name = str(row.get("name", qualified_name.split(".")[-1] if qualified_name else ""))
+        node_type = self._extract_type(row)
+
+        return GraphNode(
+            node_id=node_id,
+            qualified_name=qualified_name,
+            name=name,
+            type=node_type,
+            path=str(row.get("path")) if row.get("path") else None,
+            start_line=int(row["start_line"]) if row.get("start_line") is not None else None,
+            end_line=int(row["end_line"]) if row.get("end_line") is not None else None,
+            docstring=str(row["docstring"]) if row.get("docstring") else None,
+            properties={k: v for k, v in row.items() if k not in [
+                "node_id", "id", "qualified_name", "name", "labels", "type",
+                "path", "start_line", "end_line", "docstring"
+            ]},
+        )
+
+
+# Convenience factory functions
+
+
+def create_graph_query_service(
+    graph_service: QueryProtocol,
+    backend: str = "memgraph",
+) -> GraphQueryService:
+    """Create graph query service with auto-detected backend.
+
+    Args:
+        graph_service: Graph service instance
+        backend: Backend type ("memgraph" or "kuzu")
+
+    Returns:
+        Configured GraphQueryService
+    """
+    return GraphQueryService(graph_service, backend)
+
+
+def query_nodes_by_vector_results(
+    vector_results: list,
+    graph_service: QueryProtocol,
+) -> list[GraphNode]:
+    """Query graph nodes corresponding to vector search results.
+
+    This is the main integration point between vector store and graph database.
+
+    Args:
+        vector_results: Results from VectorStore.search_similar()
+        graph_service: Graph service to query
+
+    Returns:
+        List of GraphNode objects
+    """
+    service = GraphQueryService(graph_service)
+    node_ids = [vr.node_id for vr in vector_results]
+    return service.fetch_nodes_by_ids(node_ids)
+
+
+def get_function_with_context(
+    qualified_name: str,
+    graph_service: QueryProtocol,
+    include_callers: bool = True,
+    include_callees: bool = True,
+) -> dict:
+    """Get comprehensive information about a function including its context.
+
+    Args:
+        qualified_name: Function qualified name
+        graph_service: Graph service
+        include_callers: Whether to include calling functions
+        include_callees: Whether to include called functions
+
+    Returns:
+        Dictionary with function info, callers, and callees
+    """
+    service = GraphQueryService(graph_service)
+
+    result = {
+        "function": None,
+        "callers": [],
+        "callees": [],
+        "related": [],
+    }
+
+    # Get main function
+    func = service.fetch_node_by_qualified_name(qualified_name)
+    if func:
+        result["function"] = func
+
+    # Get callers
+    if include_callers:
+        result["callers"] = service.fetch_callers(qualified_name)
+
+    # Get callees
+    if include_callees:
+        result["callees"] = service.fetch_callees(qualified_name)
+
+    return result