hanzo-mcp 0.5.1__py3-none-any.whl → 0.5.2__py3-none-any.whl

hanzo_mcp/tools/database/database_manager.py

@@ -0,0 +1,246 @@
+"""Database manager for project-specific SQLite and graph databases."""
+
+import sqlite3
+import os
+import json
+from pathlib import Path
+from typing import Dict, Any, Optional, Tuple, List
+from datetime import datetime
+
+from hanzo_mcp.tools.common.permissions import PermissionManager
+
+
+class ProjectDatabase:
+    """Manages SQLite and graph databases for a project."""
+    
+    def __init__(self, project_path: str):
+        self.project_path = Path(project_path)
+        self.db_dir = self.project_path / ".hanzo" / "db"
+        self.db_dir.mkdir(parents=True, exist_ok=True)
+        
+        # SQLite database path
+        self.sqlite_path = self.db_dir / "project.db"
+        self.graph_path = self.db_dir / "graph.db"
+        
+        # Initialize databases
+        self._init_sqlite()
+        self._init_graph()
+        
+        # Keep graph in memory for performance
+        self.graph_conn = sqlite3.connect(":memory:")
+        self._init_graph_schema(self.graph_conn)
+        self._load_graph_from_disk()
+    
+    def _init_sqlite(self):
+        """Initialize SQLite database with common tables."""
+        conn = sqlite3.connect(self.sqlite_path)
+        try:
+            # Create metadata table
+            conn.execute('''
+                CREATE TABLE IF NOT EXISTS metadata (
+                    key TEXT PRIMARY KEY,
+                    value TEXT,
+                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+                )
+            ''')
+            
+            # Create files table
+            conn.execute('''
+                CREATE TABLE IF NOT EXISTS files (
+                    path TEXT PRIMARY KEY,
+                    content TEXT,
+                    size INTEGER,
+                    modified_at TIMESTAMP,
+                    hash TEXT,
+                    metadata TEXT
+                )
+            ''')
+            
+            # Create symbols table
+            conn.execute('''
+                CREATE TABLE IF NOT EXISTS symbols (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    file_path TEXT,
+                    name TEXT,
+                    type TEXT,
+                    line_start INTEGER,
+                    line_end INTEGER,
+                    scope TEXT,
+                    signature TEXT,
+                    FOREIGN KEY (file_path) REFERENCES files(path)
+                )
+            ''')
+            
+            # Create index for fast searches
+            conn.execute('CREATE INDEX IF NOT EXISTS idx_files_path ON files(path)')
+            conn.execute('CREATE INDEX IF NOT EXISTS idx_symbols_name ON symbols(name)')
+            conn.execute('CREATE INDEX IF NOT EXISTS idx_symbols_type ON symbols(type)')
+            
+            conn.commit()
+        finally:
+            conn.close()
+    
+    def _init_graph(self):
+        """Initialize graph database on disk."""
+        conn = sqlite3.connect(self.graph_path)
+        try:
+            self._init_graph_schema(conn)
+            conn.commit()
+        finally:
+            conn.close()
+    
+    def _init_graph_schema(self, conn: sqlite3.Connection):
+        """Initialize graph database schema."""
+        # Nodes table
+        conn.execute('''
+            CREATE TABLE IF NOT EXISTS nodes (
+                id TEXT PRIMARY KEY,
+                type TEXT NOT NULL,
+                properties TEXT,
+                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+            )
+        ''')
+        
+        # Edges table
+        conn.execute('''
+            CREATE TABLE IF NOT EXISTS edges (
+                source TEXT NOT NULL,
+                target TEXT NOT NULL,
+                relationship TEXT NOT NULL,
+                weight REAL DEFAULT 1.0,
+                properties TEXT,
+                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                PRIMARY KEY (source, target, relationship),
+                FOREIGN KEY (source) REFERENCES nodes(id),
+                FOREIGN KEY (target) REFERENCES nodes(id)
+            )
+        ''')
+        
+        # Indexes for graph traversal
+        conn.execute('CREATE INDEX IF NOT EXISTS idx_edges_source ON edges(source)')
+        conn.execute('CREATE INDEX IF NOT EXISTS idx_edges_target ON edges(target)')
+        conn.execute('CREATE INDEX IF NOT EXISTS idx_edges_relationship ON edges(relationship)')
+        conn.execute('CREATE INDEX IF NOT EXISTS idx_nodes_type ON nodes(type)')
+    
+    def _load_graph_from_disk(self):
+        """Load graph from disk into memory."""
+        disk_conn = sqlite3.connect(self.graph_path)
+        try:
+            # Copy nodes
+            nodes = disk_conn.execute('SELECT * FROM nodes').fetchall()
+            self.graph_conn.executemany(
+                'INSERT OR REPLACE INTO nodes VALUES (?, ?, ?, ?)', 
+                nodes
+            )
+            
+            # Copy edges
+            edges = disk_conn.execute('SELECT * FROM edges').fetchall()
+            self.graph_conn.executemany(
+                'INSERT OR REPLACE INTO edges VALUES (?, ?, ?, ?, ?, ?)', 
+                edges
+            )
+            
+            self.graph_conn.commit()
+        finally:
+            disk_conn.close()
+    
+    def _save_graph_to_disk(self):
+        """Save in-memory graph to disk."""
+        disk_conn = sqlite3.connect(self.graph_path)
+        try:
+            # Clear existing data
+            disk_conn.execute('DELETE FROM edges')
+            disk_conn.execute('DELETE FROM nodes')
+            
+            # Copy nodes
+            nodes = self.graph_conn.execute('SELECT * FROM nodes').fetchall()
+            disk_conn.executemany(
+                'INSERT INTO nodes VALUES (?, ?, ?, ?)', 
+                nodes
+            )
+            
+            # Copy edges
+            edges = self.graph_conn.execute('SELECT * FROM edges').fetchall()
+            disk_conn.executemany(
+                'INSERT INTO edges VALUES (?, ?, ?, ?, ?, ?)', 
+                edges
+            )
+            
+            disk_conn.commit()
+        finally:
+            disk_conn.close()
+    
+    def get_sqlite_connection(self) -> sqlite3.Connection:
+        """Get SQLite connection."""
+        conn = sqlite3.connect(self.sqlite_path)
+        conn.row_factory = sqlite3.Row
+        return conn
+    
+    def get_graph_connection(self) -> sqlite3.Connection:
+        """Get in-memory graph connection."""
+        return self.graph_conn
+    
+    def close(self):
+        """Close connections and save graph to disk."""
+        self._save_graph_to_disk()
+        self.graph_conn.close()
+
+
+class DatabaseManager:
+    """Manages databases for multiple projects."""
+    
+    def __init__(self, permission_manager: PermissionManager):
+        self.permission_manager = permission_manager
+        self.projects: Dict[str, ProjectDatabase] = {}
+        self.search_paths: List[str] = []
+    
+    def add_search_path(self, path: str):
+        """Add a path to search for projects."""
+        if path not in self.search_paths:
+            self.search_paths.append(path)
+    
+    def get_project_db(self, project_path: str) -> ProjectDatabase:
+        """Get or create database for a project."""
+        project_path = os.path.abspath(project_path)
+        
+        # Check permissions
+        if not self.permission_manager.has_permission(project_path):
+            raise PermissionError(f"No permission to access: {project_path}")
+        
+        # Create database if not exists
+        if project_path not in self.projects:
+            self.projects[project_path] = ProjectDatabase(project_path)
+        
+        return self.projects[project_path]
+    
+    def get_project_for_path(self, file_path: str) -> Optional[ProjectDatabase]:
+        """Find the project database for a given file path."""
+        file_path = os.path.abspath(file_path)
+        
+        # Check if file is in a known project
+        for project_path in self.projects:
+            if file_path.startswith(project_path):
+                return self.projects[project_path]
+        
+        # Search up the directory tree for a project
+        current = Path(file_path)
+        if current.is_file():
+            current = current.parent
+            
+        while current != current.parent:
+            # Check for project markers
+            if (current / ".git").exists() or (current / "LLM.md").exists():
+                return self.get_project_db(str(current))
+            current = current.parent
+        
+        # No project found, use the directory of the file
+        if Path(file_path).is_file():
+            return self.get_project_db(str(Path(file_path).parent))
+        else:
+            return self.get_project_db(file_path)
+    
+    def close_all(self):
+        """Close all project databases."""
+        for db in self.projects.values():
+            db.close()
+        self.projects.clear()

hanzo_mcp/tools/database/graph_add.py

@@ -0,0 +1,257 @@
+"""Graph add tool for adding nodes and edges to the graph database."""
+
+import json
+from typing import Annotated, Optional, TypedDict, Unpack, final, override
+
+from fastmcp import Context as MCPContext
+from pydantic import Field
+
+from hanzo_mcp.tools.common.base import BaseTool
+from hanzo_mcp.tools.common.context import create_tool_context
+from hanzo_mcp.tools.common.permissions import PermissionManager
+from hanzo_mcp.tools.database.database_manager import DatabaseManager
+
+
+NodeId = Annotated[
+    Optional[str],
+    Field(
+        description="Node ID to add (required for nodes)",
+        default=None,
+    ),
+]
+
+NodeType = Annotated[
+    Optional[str],
+    Field(
+        description="Node type (e.g., 'file', 'function', 'class')",
+        default=None,
+    ),
+]
+
+Properties = Annotated[
+    Optional[dict],
+    Field(
+        description="Properties as JSON object",
+        default=None,
+    ),
+]
+
+Source = Annotated[
+    Optional[str],
+    Field(
+        description="Source node ID (required for edges)",
+        default=None,
+    ),
+]
+
+Target = Annotated[
+    Optional[str],
+    Field(
+        description="Target node ID (required for edges)",
+        default=None,
+    ),
+]
+
+Relationship = Annotated[
+    Optional[str],
+    Field(
+        description="Edge relationship type (e.g., 'imports', 'calls', 'inherits')",
+        default=None,
+    ),
+]
+
+Weight = Annotated[
+    float,
+    Field(
+        description="Edge weight (default 1.0)",
+        default=1.0,
+    ),
+]
+
+ProjectPath = Annotated[
+    Optional[str],
+    Field(
+        description="Project path (defaults to current directory)",
+        default=None,
+    ),
+]
+
+
+class GraphAddParams(TypedDict, total=False):
+    """Parameters for graph add tool."""
+
+    node_id: Optional[str]
+    node_type: Optional[str]
+    properties: Optional[dict]
+    source: Optional[str]
+    target: Optional[str]
+    relationship: Optional[str]
+    weight: float
+    project_path: Optional[str]
+
+
+@final
+class GraphAddTool(BaseTool):
+    """Tool for adding nodes and edges to graph database."""
+
+    def __init__(self, permission_manager: PermissionManager, db_manager: DatabaseManager):
+        """Initialize the graph add tool.
+
+        Args:
+            permission_manager: Permission manager for access control
+            db_manager: Database manager instance
+        """
+        self.permission_manager = permission_manager
+        self.db_manager = db_manager
+
+    @property
+    @override
+    def name(self) -> str:
+        """Get the tool name."""
+        return "graph_add"
+
+    @property
+    @override
+    def description(self) -> str:
+        """Get the tool description."""
+        return """Add nodes and edges to the project's graph database.
+
+To add a node:
+- Provide node_id and node_type
+- Optionally add properties as JSON
+
+To add an edge:
+- Provide source, target, and relationship
+- Optionally add weight and properties
+
+Common node types:
+- file, function, class, module, variable
+
+Common relationships:
+- imports, calls, inherits, contains, references, depends_on
+
+Examples:
+- graph_add --node-id "main.py" --node-type "file" --properties '{"size": 1024}'
+- graph_add --node-id "MyClass" --node-type "class" --properties '{"file": "main.py"}'
+- graph_add --source "main.py" --target "utils.py" --relationship "imports"
+- graph_add --source "func1" --target "func2" --relationship "calls" --weight 5.0
+"""
+
+    @override
+    async def call(
+        self,
+        ctx: MCPContext,
+        **params: Unpack[GraphAddParams],
+    ) -> str:
+        """Add nodes or edges to graph.
+
+        Args:
+            ctx: MCP context
+            **params: Tool parameters
+
+        Returns:
+            Result of add operation
+        """
+        tool_ctx = create_tool_context(ctx)
+        await tool_ctx.set_tool_info(self.name)
+
+        # Extract parameters
+        node_id = params.get("node_id")
+        node_type = params.get("node_type")
+        properties = params.get("properties", {})
+        source = params.get("source")
+        target = params.get("target")
+        relationship = params.get("relationship")
+        weight = params.get("weight", 1.0)
+        project_path = params.get("project_path")
+
+        # Determine if adding node or edge
+        is_node = node_id is not None
+        is_edge = source is not None and target is not None
+
+        if not is_node and not is_edge:
+            return "Error: Must provide either (node_id and node_type) for a node, or (source, target, relationship) for an edge"
+
+        if is_node and is_edge:
+            return "Error: Cannot add both node and edge in one operation"
+
+        # Get project database
+        try:
+            if project_path:
+                project_db = self.db_manager.get_project_db(project_path)
+            else:
+                import os
+                project_db = self.db_manager.get_project_for_path(os.getcwd())
+                
+            if not project_db:
+                return "Error: Could not find project database"
+
+        except PermissionError as e:
+            return str(e)
+        except Exception as e:
+            return f"Error accessing project database: {str(e)}"
+
+        # Get graph connection
+        graph_conn = project_db.get_graph_connection()
+
+        try:
+            if is_node:
+                # Add node
+                if not node_type:
+                    return "Error: node_type is required when adding a node"
+
+                await tool_ctx.info(f"Adding node: {node_id} (type: {node_type})")
+
+                # Serialize properties
+                properties_json = json.dumps(properties) if properties else None
+
+                # Insert or update node
+                graph_conn.execute("""
+                    INSERT OR REPLACE INTO nodes (id, type, properties)
+                    VALUES (?, ?, ?)
+                """, (node_id, node_type, properties_json))
+                
+                graph_conn.commit()
+                
+                return f"Successfully added node '{node_id}' of type '{node_type}'"
+
+            else:
+                # Add edge
+                if not relationship:
+                    return "Error: relationship is required when adding an edge"
+
+                await tool_ctx.info(f"Adding edge: {source} --[{relationship}]--> {target}")
+
+                # Check if nodes exist
+                cursor = graph_conn.cursor()
+                cursor.execute("SELECT id FROM nodes WHERE id IN (?, ?)", (source, target))
+                existing = [row[0] for row in cursor.fetchall()]
+                
+                if source not in existing:
+                    return f"Error: Source node '{source}' does not exist"
+                if target not in existing:
+                    return f"Error: Target node '{target}' does not exist"
+
+                # Serialize properties
+                properties_json = json.dumps(properties) if properties else None
+
+                # Insert or update edge
+                graph_conn.execute("""
+                    INSERT OR REPLACE INTO edges (source, target, relationship, weight, properties)
+                    VALUES (?, ?, ?, ?, ?)
+                """, (source, target, relationship, weight, properties_json))
+                
+                graph_conn.commit()
+                
+                # Save to disk
+                project_db._save_graph_to_disk()
+                
+                return f"Successfully added edge: {source} --[{relationship}]--> {target} (weight: {weight})"
+
+        except Exception as e:
+            await tool_ctx.error(f"Failed to add to graph: {str(e)}")
+            return f"Error adding to graph: {str(e)}"
+
+    def register(self, mcp_server) -> None:
+        """Register this tool with the MCP server."""
+        pass