added 3 tools

execute_select_query
execute_update_query
execute_delete_query

Author: SHAHARYAR1255

README-zh.md

@@ -308,6 +308,27 @@ ORDERS 表定义了哪些约束？
 ORDERS 表与哪些表有关联？
 ```
 
+#### `execute_select_query`
+执行 SELECT SQL 查询并返回格式化的结果表。仅允许以 SELECT 开头的查询。
+示例：
+```
+运行：SELECT * FROM EMPLOYEES WHERE DEPARTMENT_ID = 10;
+```
+
+#### `execute_update_query`
+执行 UPDATE SQL 查询并返回受影响的行数。仅允许以 UPDATE 开头的查询。
+示例：
+```
+将部门 10 的所有员工工资上调 10%：UPDATE EMPLOYEES SET SALARY = SALARY * 1.1 WHERE DEPARTMENT_ID = 10;
+```
+
+#### `execute_delete_query`
+执行 DELETE SQL 查询并返回受影响的行数。仅允许以 DELETE 开头的查询。
+示例：
+```
+删除部门 10 的所有员工：DELETE FROM EMPLOYEES WHERE DEPARTMENT_ID = 10;
+```
+
 ## 架构
 
 本 MCP 服务器采用三层架构，针对大型 Oracle 数据库进行了优化：

README.md

@@ -308,6 +308,27 @@ Example:
 What tables are related to the ORDERS table?
 ```
 
+#### `execute_select_query`
+Execute a SELECT SQL query and return the results as a formatted table. Only queries starting with SELECT are allowed.
+Example:
+```
+Run: SELECT * FROM EMPLOYEES WHERE DEPARTMENT_ID = 10;
+```
+
+#### `execute_update_query`
+Execute an UPDATE SQL query and return the number of affected rows. Only queries starting with UPDATE are allowed.
+Example:
+```
+Update the salary for all employees in department 10: UPDATE EMPLOYEES SET SALARY = SALARY * 1.1 WHERE DEPARTMENT_ID = 10;
+```
+
+#### `execute_delete_query`
+Execute a DELETE SQL query and return the number of affected rows. Only queries starting with DELETE are allowed.
+Example:
+```
+Delete all employees in department 10: DELETE FROM EMPLOYEES WHERE DEPARTMENT_ID = 10;
+```
+
 ## Architecture
 
 This MCP server employs a three-layer architecture optimized for large-scale Oracle databases:

db_context/__init__.py

@@ -135,4 +135,16 @@ async def get_related_tables(self, table_name: str) -> Dict[str, List[str]]:
 
     async def explain_query_plan(self, query: str) -> Dict[str, Any]:
         """Get execution plan for an SQL query with optimization suggestions"""
-        return await self.db_connector.explain_query_plan(query)
+        return await self.db_connector.explain_query_plan(query)
+
+    async def execute_select_query(self, query: str) -> list:
+        """Execute a SELECT query and return the results as a list of dicts."""
+        return await self.db_connector.execute_select_query(query)
+
+    async def execute_update_query(self, query: str) -> int:
+        """Execute an UPDATE query and return the number of affected rows."""
+        return await self.db_connector.execute_update_query(query)
+
+    async def execute_delete_query(self, query: str) -> int:
+        """Execute a DELETE query and return the number of affected rows."""
+        return await self.db_connector.execute_delete_query(query)

db_context/database.py

@@ -811,6 +811,46 @@ def _analyze_query_for_optimization(self, query: str) -> List[str]:
 
         return suggestions
 
+    async def execute_select_query(self, query: str) -> list:
+        """Execute a SELECT query and return the results as a list of dicts. Only allows queries starting with SELECT."""
+        conn = await self.get_connection()
+        try:
+            cursor = conn.cursor()
+            await cursor.execute(query) if not self.thick_mode else cursor.execute(query)
+            columns = [col[0] for col in cursor.description]
+            if self.thick_mode:
+                rows = cursor.fetchall()
+            else:
+                rows = await cursor.fetchall()
+            result = [dict(zip(columns, row)) for row in rows]
+            return result
+        finally:
+            await self._close_connection(conn)
+
+    async def execute_update_query(self, query: str) -> int:
+        """Execute an UPDATE query and return the number of affected rows. Only allows queries starting with UPDATE."""
+        conn = await self.get_connection()
+        try:
+            cursor = conn.cursor()
+            await cursor.execute(query) if not self.thick_mode else cursor.execute(query)
+            rowcount = cursor.rowcount
+            await self._commit(conn)
+            return rowcount
+        finally:
+            await self._close_connection(conn)
+
+    async def execute_delete_query(self, query: str) -> int:
+        """Execute a DELETE query and return the number of affected rows. Only allows queries starting with DELETE."""
+        conn = await self.get_connection()
+        try:
+            cursor = conn.cursor()
+            await cursor.execute(query) if not self.thick_mode else cursor.execute(query)
+            rowcount = cursor.rowcount
+            await self._commit(conn)
+            return rowcount
+        finally:
+            await self._close_connection(conn)
+
     async def _close_connection(self, conn):
         """Helper method to close connection based on mode"""
         try:

db_context/schema/formatter.py

@@ -359,4 +359,18 @@ def _format_relationship_groups(groups: List[Dict[str, Any]], result: List[str])
         else:
             result.append(f"    - {group['pattern']}:")
             for pattern in sorted(group['column_patterns']):
-                result.append(f"      {pattern}")
+                result.append(f"      {pattern}")
+
+def format_select_results(rows: list) -> str:
+    """Format SELECT query results (list of dicts) as a readable table string."""
+    if not rows:
+        return "No results."
+    columns = list(rows[0].keys())
+    # Calculate max width for each column
+    col_widths = {col: max(len(str(col)), max(len(str(row.get(col, ''))) for row in rows)) for col in columns}
+    # Header
+    header = " | ".join(col.ljust(col_widths[col]) for col in columns)
+    sep = "-+-".join("-" * col_widths[col] for col in columns)
+    # Rows
+    row_lines = [" | ".join(str(row.get(col, '')).ljust(col_widths[col]) for col in columns) for row in rows]
+    return f"{header}\n{sep}\n" + "\n".join(row_lines)

main.py

@@ -9,6 +9,7 @@
 from dotenv import load_dotenv
 
 from db_context import DatabaseContext
+from db_context.schema.formatter import format_select_results
 
 # Load environment variables from .env file
 load_dotenv()
@@ -607,5 +608,59 @@ async def get_related_tables(table_name: str, ctx: Context) -> str:
     except Exception as e:
         return f"Error getting related tables: {str(e)}"
 
+@mcp.tool()
+async def execute_select_query(query: str, ctx: Context) -> str:
+    """
+    Execute a SELECT SQL query and return the results as a formatted table. Only queries starting with SELECT are allowed.
+    Args:
+        query: The SQL SELECT query to execute. Must start with SELECT (case-insensitive).
+    Returns:
+        A formatted string table of the query results, or an error message if the query is invalid or fails.
+    """
+    db_context: DatabaseContext = ctx.request_context.lifespan_context
+    if not query.strip().lower().startswith("select"):
+        return "Error: Only SELECT queries are allowed."
+    try:
+        rows = await db_context.execute_select_query(query)
+        return format_select_results(rows)
+    except Exception as e:
+        return f"Error executing query: {str(e)}"
+
+@mcp.tool()
+async def execute_update_query(query: str, ctx: Context) -> str:
+    """
+    Execute an UPDATE SQL query and return the number of affected rows. Only queries starting with UPDATE are allowed.
+    Args:
+        query: The SQL UPDATE query to execute. Must start with UPDATE (case-insensitive).
+    Returns:
+        A string indicating the number of affected rows, or an error message if the query is invalid or fails.
+    """
+    db_context: DatabaseContext = ctx.request_context.lifespan_context
+    if not query.strip().lower().startswith("update"):
+        return "Error: Only UPDATE queries are allowed."
+    try:
+        count = await db_context.execute_update_query(query)
+        return f"Rows affected: {count}"
+    except Exception as e:
+        return f"Error executing update: {str(e)}"
+
+@mcp.tool()
+async def execute_delete_query(query: str, ctx: Context) -> str:
+    """
+    Execute a DELETE SQL query and return the number of affected rows. Only queries starting with DELETE are allowed.
+    Args:
+        query: The SQL DELETE query to execute. Must start with DELETE (case-insensitive).
+    Returns:
+        A string indicating the number of affected rows, or an error message if the query is invalid or fails.
+    """
+    db_context: DatabaseContext = ctx.request_context.lifespan_context
+    if not query.strip().lower().startswith("delete"):
+        return "Error: Only DELETE queries are allowed."
+    try:
+        count = await db_context.execute_delete_query(query)
+        return f"Rows affected: {count}"
+    except Exception as e:
+        return f"Error executing delete: {str(e)}"
+
 if __name__ == "__main__":
     mcp.run()