"""SQL query tool for direct database queries.""" import sqlite3 from typing import Unpack, Optional, Annotated, TypedDict, final, override from pydantic import Field from mcp.server.fastmcp import Context as MCPContext from hanzo_tools.core import ( BaseTool, PermissionManager, auto_timeout, create_tool_context, ) from .database_manager import DatabaseManager Query = Annotated[ str, Field( description="SQL query to execute", min_length=1, ), ] ProjectPath = Annotated[ Optional[str], Field( description="Project path (defaults to current directory)", default=None, ), ] ReadOnly = Annotated[ bool, Field( description="Execute in read-only mode (no INSERT/UPDATE/DELETE)", default=True, ), ] class SqlQueryParams(TypedDict, total=False): """Parameters for SQL query tool.""" query: str project_path: Optional[str] read_only: bool @final class SqlQueryTool(BaseTool): """Tool for executing SQL queries on project databases.""" def __init__( self, permission_manager: PermissionManager, db_manager: DatabaseManager ): """Initialize the SQL query 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 "sql_query" @property @override def description(self) -> str: """Get the tool description.""" return """Execute SQL queries on the project's embedded SQLite database. Each project has its own SQLite database with tables: - metadata: Key-value store for project metadata - files: File information and content - symbols: Code symbols (functions, classes, etc.) Features: - Direct SQL query execution - Read-only mode by default (safety) - Returns results in tabular format - Automatic project detection Examples: - sql_query --query "SELECT * FROM files LIMIT 10" - sql_query --query "SELECT name, type FROM symbols WHERE type='function'" - sql_query --query "INSERT INTO metadata (key, value) VALUES ('version', '1.0')" --read-only false Note: Use sql_search for text search operations.""" @override @auto_timeout("sql_query") async def call( self, ctx: MCPContext, **params: Unpack[SqlQueryParams], ) -> str: """Execute SQL query. Args: ctx: MCP context **params: Tool parameters Returns: Query results """ tool_ctx = create_tool_context(ctx) await tool_ctx.set_tool_info(self.name) # Extract parameters query = params.get("query") if not query: return "Error: query is required" project_path = params.get("project_path") read_only = params.get("read_only", True) # 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)}" # Check if query is read-only using token-based parsing if read_only: import re # Remove comments and string literals to prevent bypass clean_query = re.sub( r"--.*$", "", query, flags=re.MULTILINE ) # Remove -- comments clean_query = re.sub( r"/\*.*?\*/", "", clean_query, flags=re.DOTALL ) # Remove /* */ comments clean_query = re.sub( r"'[^']*'", "''", clean_query ) # Neutralize string literals clean_query = re.sub( r'"[^"]*"', '""', clean_query ) # Neutralize identifiers # Check for write operations using word boundaries write_keywords = [ "INSERT", "UPDATE", "DELETE", "DROP", "CREATE", "ALTER", "TRUNCATE", "REPLACE", ] for keyword in write_keywords: if re.search(rf"\b{keyword}\b", clean_query, re.IGNORECASE): return f"Error: Query contains {keyword} operation. Set --read-only false to allow write operations." await tool_ctx.info( f"Executing SQL query on project: {project_db.project_path}" ) # Execute query conn = None try: conn = project_db.get_sqlite_connection() cursor = conn.cursor() # Execute the query cursor.execute(query) # Handle different query types if query.strip().upper().startswith("SELECT"): # Fetch results results = cursor.fetchall() if not results: return "No results found." # Get column names columns = [desc[0] for desc in cursor.description] # Format as table output = self._format_results_table(columns, results) return f"Query executed successfully. Found {len(results)} row(s).\n\n{output}" else: # For non-SELECT queries, commit and return affected rows conn.commit() affected = cursor.rowcount return f"Query executed successfully. Affected {affected} row(s)." except sqlite3.Error as e: await tool_ctx.error(f"SQL error: {str(e)}") return f"SQL error: {str(e)}" except Exception as e: await tool_ctx.error(f"Unexpected error: {str(e)}") return f"Error executing query: {str(e)}" finally: if conn: conn.close() def _format_results_table(self, columns: list[str], rows: list[tuple]) -> str: """Format query results as a table.""" if not rows: return "No results" # Calculate column widths col_widths = [] for i, col in enumerate(columns): max_width = len(col) for row in rows[:100]: # Check first 100 rows val_str = str(row[i]) if row[i] is not None else "NULL" max_width = max(max_width, len(val_str)) col_widths.append(min(max_width, 50)) # Cap at 50 chars # Build header header = " | ".join( col.ljust(width) for col, width in zip(columns, col_widths, strict=False) ) separator = "-+-".join("-" * width for width in col_widths) # Build rows output_rows = [] for row in rows[:1000]: # Limit to 1000 rows row_str = " | ".join( self._truncate(str(val) if val is not None else "NULL", width).ljust( width ) for val, width in zip(row, col_widths, strict=False) ) output_rows.append(row_str) # Combine output = [header, separator] + output_rows if len(rows) > 1000: output.append(f"\n... and {len(rows) - 1000} more rows") return "\n".join(output) def _truncate(self, text: str, max_width: int) -> str: """Truncate text to max width.""" if len(text) <= max_width: return text return text[: max_width - 3] + "..." def register(self, mcp_server) -> None: """Register this tool with the MCP server.""" pass