从零开始：MCP数据库助手（二）- 核心功能实现

为什么要实现这3个核心功能？

上一篇我们搭建了基础框架，AI已经能"看到"数据库中有哪些表了。但这还远远不够，今天咱们要让AI变得更聪明，实现3个核心功能：

• 能智能搜索相关的表
• 能深度理解表的结构
• 能安全地执行SQL查询

第一步：filter_table_names() - 智能搜索表名

第1周我们已经能获取所有表名了，但当数据库有几十个表的时候，你得能快速找到相关的表。这就像在图书馆里找书，你不可能一本本翻，得按分类找。

@mcp.tool(description="根据关键词搜索相关的表名。支持模糊匹配，不区分大小写。")
def filter_table_names(keyword: str) -> str:"""根据关键词搜索相关的表名"""try:if not keyword or not keyword.strip():return "请提供搜索关键词"db_mgr = get_database_manager()matching_tables = db_mgr.filter_table_names(keyword.strip())if matching_tables:table_list = ", ".join(matching_tables)result = f"搜索关键词 '{keyword}' 找到 {len(matching_tables)} 个相关表:\n"result += f"匹配的表: {table_list}"logger.info(f"关键词 '{keyword}' 匹配到: {table_list}")return resultelse:all_tables = db_mgr.get_table_names()all_table_list = ", ".join(all_tables) if all_tables else "无"result = f"搜索关键词 '{keyword}' 没有找到匹配的表。\n"result += f"数据库中的所有表: {all_table_list}\n"result += f"建议尝试其他关键词，如表名的一部分。"logger.info(f"关键词 '{keyword}' 没有找到匹配的表。")return resultexcept Exception as e:error_msg = f"搜索表名失败: {str(e)}"logger.error(error_msg)return f"{error_msg}"

数据库管理器中的搜索逻辑：

def filter_table_names(self, keyword: str) -> List[str]:"""根据关键词搜索相关表名"""try:all_tables = self.get_table_names()# 不区分大小写的模糊匹配matching_tables = [table for table in all_tables if keyword.lower() in table.lower()]return matching_tablesexcept Exception as e:logger.error(f"搜索表名失败: {e}")raise

这个功能在大型项目中特别有用，想象一下在一个有100个表的电商系统中快速找到所有支付相关的表！

第二步：schema_definitions() - 深度解析表结构

知道表名还不够，得知道表里有什么字段、什么类型、什么关系。

@mcp.tool(description="获取指定表的详细结构信息，包括列、类型、主键、索引、外键等。支持同时查询多个表。")
def schema_info(table_names: List[str]) -> str:"""获取指定表的详细结构信息"""try:if not table_names:return "请提供要查询的表名"db_mgr = get_database_manager()result_parts = []for table_name in table_names:try:schema_info = db_mgr.get_table_schema(table_name)# 格式化表结构信息table_section = f"\n{'='*50}\n"table_section += f"表名：{table_name}\n"table_section += f"列：{schema_info['columns']}\n"table_section += f"{'='*50}\n"# 列信息column_section = "列信息：\n"for col in schema_info['columns']:col_type = col['type']nullable = "可空" if col['nullable'] else "不可空"pk_mark = " 主键" if col['is_primary_key'] else ""default_info = f" (默认: {col['default']})" if col['default'] else ""table_section += f"  • {col['name']}: {col_type} - {nullable}{pk_mark}{default_info}\n"# 主键信息if schema_info['primary_keys']:table_section += f"\n主键: {', '.join(schema_info['primary_keys'])}\n"# 索引信息if schema_info['indexes']:table_section += "\n索引信息:\n"for idx in schema_info['indexes']:unique_mark = "唯一索引" if idx.get('unique', False) else "普通索引"columns = ', '.join(idx['column_names'])table_section += f"  • {idx['name']}: {unique_mark} ({columns})\n"# 外键信息if schema_info['foreign_keys']:table_section += "\n外键关系:\n"for fk in schema_info['foreign_keys']:local_cols = ', '.join(fk['constrained_columns'])ref_table = fk['referred_table']ref_cols = ', '.join(fk['referred_columns'])table_section += f"  • {local_cols} → {ref_table}.{ref_cols}\n"result_parts.append(table_section)except ValueError as e:result_parts.append(f"\n表 '{table_name}': {str(e)}\n")except Exception as e:result_parts.append(f"\n表 '{table_name}' 解析失败: {str(e)}\n")final_result = '\n'.join(result_parts)logger.info(f"成功解析 {len(table_names)} 个表的结构")return final_resultexcept Exception as e:error_msg = f"获取表结构信息失败: {str(e)}"logger.error(error_msg)return f"{error_msg}"

数据库管理器中的解析逻辑：

    def get_table_schema(self, table_name: str) -> Dict[str, Any]:"""获取指定表的详细结构信息"""try:with self.get_connection() as conn:inspector = inspect(conn)# 验证表是否存在if table_name not in inspector.get_table_names():raise ValueError(f"表 '{table_name}' 不存在")# 获取列信息columns = inspector.get_columns(table_name)# 获取主键信息pk_constraint = inspector.get_pk_constraint(table_name)primary_keys = set(pk_constraint["constrained_columns"])# 获取索引信息indexes = inspector.get_indexes(table_name)# 获取外键信息foreign_keys = inspector.get_foreign_keys(table_name)# 格式化列信息formatted_columns = []for col in columns:col_info = {"name": col["name"],"type": str(col["type"]),"nullable": col["nullable"],"default": col.get("default"),"is_primary_key": col["name"] in primary_keys}formatted_columns.append(col_info)schema_info = {"table_name": table_name,"columns": formatted_columns,"primary_keys": list(primary_keys),"indexes": indexes,"foreign_keys": foreign_keys,"column_count": len(columns)}return schema_infoexcept Exception as e:logger.error(f"获取表结构信息失败: {e}")raise

第三步：execute_query() - 安全执行SQL查询

最后这个功能是核心中的核心！让AI能执行SQL查询，但必须确保安全性。我们只允许SELECT查询，并且有完整的防护措施。

@mcp.tool(description="安全执行SQL查询语句。只支持SELECT查询，自动添加结果限制，支持参数化查询防止SQL注入。")
def execute_query(query: str, params: dict = None) -> str:try:if not query or not query.strip():return "请提供查询语句"db_mgr = get_database_manager()# 如果没有提供参数，设为Nonequery_params = params if params else None# 执行查询result = db_mgr.execute_query(query.strip(), query_params)if not result:return "查询结果为空"#格式化返结果result_text = "查询结果:\n"for row in result:result_text += f"{row}\n"return result_textexcept ValueError as e:# 安全验证失败error_msg = f"查询安全验证失败: {str(e)}"logger.warning(error_msg)return f"{error_msg}\n\n 提示：本工具只支持 SELECT 查询，且会自动添加安全限制。"except Exception as e:error_msg = f"执行查询失败: {str(e)}"logger.error(error_msg)return f"{error_msg}"

数据库管理器中也要有验证：

    def _validate_query(self, query: str) -> None:"""验证查询安全性"""query_upper = query.upper().strip()# 检查危险的SQL操作dangerous_keywords = ["DROP", "DELETE", "UPDATE", "INSERT", "ALTER", "CREATE", "TRUNCATE"]for keyword in dangerous_keywords:if keyword in query_upper:raise ValueError(f"查询包含危险操作: {keyword}")# 检查是否是SELECT查询if not query_upper.startswith("SELECT"):raise ValueError("只允许执行SELECT查询")def _add_limit_to_query(self, query: str, limit: int) -> str:"""为查询添加LIMIT子句"""query_upper = query.upper()# 如果已经有LIMIT，不再添加if "LIMIT" in query_upper:return query# 添加LIMIT子句return f"{query.rstrip(';')} LIMIT {limit}"def execute_query(self, query: str, params: dict = None) -> List[Dict[str, Any]]:"""安全执行SQL查询"""try:# 验证查询安全性self._validate_query(query)# 添加默认限制limited_query = self._add_limit_to_query(query, 1000)with self.get_connection() as conn:if params:# 使用参数化查询result = conn.execute(text(limited_query), params)else:# 直接执行查询result = conn.execute(text(limited_query))# 获取列名columns = result.keys()# 转换为字典列表rows = []for row in result:row_dict = dict(zip(columns, row))rows.append(row_dict)return rowsexcept Exception as e:logger.error(f"执行查询失败: {e}")raise

第四步：功能验证测试

为了确保所有功能都正常工作，我们创建一个简单的验证脚本。这不是复杂的单元测试，而是帮助读者快速验证功能的工具。

创建验证脚本

"""
tests/verify_functions.py - 功能验证脚本
"""import os
import sys
from pathlib import Path# 添加项目路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root / "src"))def verify_database_connection():"""验证数据库连接"""print("=== 数据库连接验证 ===")# 设置数据库URLdb_path = project_root / "data" / "test.db"if not db_path.exists():print("测试数据库不存在！")return Falseos.environ['DB_URL'] = f'sqlite:///{db_path}'try:from mcp_datatools.database import DatabaseManagerdb_mgr = DatabaseManager()if db_mgr.test_connection():print("数据库连接成功")return Trueelse:print("数据库连接失败")return Falseexcept Exception as e:print(f"数据库连接错误: {e}")return Falsedef verify_basic_functions():"""验证基础功能"""print("\n=== 基础功能验证 ===")try:from mcp_datatools.database import DatabaseManagerdb_mgr = DatabaseManager()# 测试1: 获取表名print("1. 测试 list_tables() 功能...")tables = db_mgr.get_table_names()print(f"   找到 {len(tables)} 个表: {', '.join(tables)}")# 测试2: 搜索表名print("2. 测试 filter_table_names() 功能...")user_tables = db_mgr.filter_table_names('user')print(f"   搜索 'user' 找到: {user_tables}")# 测试3: 获取表结构print("3. 测试 get_table_schema() 功能...")if tables:schema = db_mgr.get_table_schema(tables[0])print(f"  表 '{tables[0]}' 有 {schema['column_count']} 个列")# 测试4: 执行查询print("4. 测试 execute_query() 功能...")result = db_mgr.execute_query("SELECT COUNT(*) as count FROM users")print(f"   查询结果: {result}")return Trueexcept Exception as e:print(f"功能验证失败: {e}")return Falsedef main():"""主验证流程"""print("MCP DataTools 功能验证")print("=" * 50)# 验证数据库连接db_ok = verify_database_connection()if db_ok:# 验证基础功能func_ok = verify_basic_functions()if func_ok:print("\n 所有基础功能验证通过！")else:print("\n部分功能验证失败，请检查代码")else:print("\n数据库连接失败，请先创建测试数据库")print("\n" + "=" * 50)print("验证完成！现在可以在Cursor中测试MCP功能了。")if __name__ == "__main__":main()

运行验证

# 1. 确保有测试数据
python data/init_sqlite_db.py# 2. 运行功能验证
python tests/verify_functions.py

验证结果示例

MCP DataTools 功能验证
==================================================
=== 数据库连接验证 ===
数据库连接成功=== 基础功能验证 ===
1. 测试 list_tables() 功能...找到 3 个表: orders, products, users
2. 测试 filter_table_names() 功能...搜索 'user' 找到: ['users']
3. 测试 get_table_schema() 功能...表 'orders' 有 3 个列
4. 测试 execute_query() 功能...查询结果: [{'count': 2}]所有基础功能验证通过！==================================================
验证完成！现在可以在Cursor中测试MCP功能了。

第五步：功能演示效果

现在让我们看看这些功能在实际使用中的效果：

1. 智能搜索表名

试试问AI：“找找用户相关的表”

或者问：“有哪些订单相关的表？”

2. 深度解析表结构

问AI：“users表的结构是什么样的？”

3. 安全执行SQL查询

简单查询： “查询users表中的所有数据”
AI执行：SELECT * FROM users

条件查询： “查询价格大于1000的产品”
AI执行：SELECT * FROM products WHERE price > 1000

复杂查询： “查询每个用户的订单数量”
AI可以执行JOIN查询

总结与预告

第2篇我们成功实现了3个核心功能：表搜索、结构分析、安全SQL执行，让AI能够理解数据库结构并执行复杂查询。第3篇将实现PostgreSQL、MySQL等主流数据库支持，加上连接池优化和Docker部署，让这个工具真正走向生产环境。