Skip to content

Commit ed2de61

Browse files
author
依诺
committed
add sql related tools
1 parent eb0e529 commit ed2de61

File tree

3 files changed

+154
-29
lines changed

3 files changed

+154
-29
lines changed

‎pyproject.toml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,14 @@
11
[project]
22
name = "alibabacloud-dms-mcp-server"
3-
version = "0.1.15"
3+
version = "0.1.16"
44
description = "MCP Server for AlibabaCloud DMS"
55
readme = "README.md"
66
authors = [
77
{ name = "AlibabaCloud DMS" }
88
]
99
requires-python = ">=3.10"
1010
dependencies = [
11-
"alibabacloud-dms-enterprise20181101>=1.75.0",
11+
"alibabacloud-dms-enterprise20181101>=1.76.0",
1212
"httpx>=0.28.1",
1313
"mcp[cli]>=1.8.1",
1414
]

‎src/alibabacloud_dms_mcp_server/server.py‎

Lines changed: 147 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,15 @@ class SqlResult(MyBaseModel):
118118
sql: Optional[str] = Field(description="The generated SQL query")
119119

120120

121+
DATABASE_ID_DESCRIPTION = (
122+
"If you don't know the databaseId, first use getDatabase or searchDatabase to retrieve it.\n"
123+
"(1) If you have the exact host, port, and database name, use getDatabase.\n"
124+
"(2) If you only know the database name, use searchDatabase.\n"
125+
"(3) If you don't know any information, ask the user to provide the necessary details.\n"
126+
"Note: searchDatabase may return multiple databases. In this case, let the user choose which one to use."
127+
)
128+
129+
121130
# --- Aliyun Client Creation ---
122131
def create_client() -> dms_enterprise20181101Client:
123132
config = open_api_models.Config(
@@ -269,7 +278,7 @@ async def get_database(
269278
) -> DatabaseDetail:
270279
client = create_client()
271280
req = dms_enterprise_20181101_models.GetDatabaseRequest(host=host, port=port, schema_name=schema_name)
272-
281+
273282
if sid:
274283
req.sid = sid
275284
if mcp.state.real_login_uid:
@@ -433,7 +442,8 @@ async def nl2sql(
433442
database_id: str = Field(description="DMS databaseId"),
434443
question: str = Field(description="Natural language question"),
435444
knowledge: Optional[str] = Field(default=None, description="Optional: additional context"),
436-
model: Optional[str] = Field(default=None, description="Optional: if a specific model is desired, it can be specified here")
445+
model: Optional[str] = Field(default=None,
446+
description="Optional: if a specific model is desired, it can be specified here")
437447
) -> SqlResult:
438448
client = create_client()
439449
req = dms_enterprise_20181101_models.GenerateSqlFromNLRequest(db_id=database_id, question=question)
@@ -454,6 +464,79 @@ async def nl2sql(
454464
raise
455465

456466

467+
async def answer_sql_syntax(
468+
database_id: str = Field(description="DMS databaseId"),
469+
question: str = Field(description="Natural language question"),
470+
model: Optional[str] = Field(default=None,
471+
description="Optional: if a specific model is desired, it can be specified here")
472+
) -> Dict[str, Any]:
473+
client = create_client()
474+
req = dms_enterprise_20181101_models.AnswerSqlSyntaxByMetaAgentRequest(db_id=database_id, query=question)
475+
# if mcp.state.real_login_uid:
476+
# req.real_login_user_uid = mcp.state.real_login_uid
477+
if model:
478+
req.model = model
479+
try:
480+
resp = client.answer_sql_syntax_by_meta_agent(req)
481+
if not resp or not resp.body:
482+
return None
483+
data = resp.body.to_map()
484+
return data
485+
except Exception as e:
486+
logger.error(f"Error in ask_sql_syntax: {e}")
487+
raise
488+
489+
490+
async def fix_sql_statement(
491+
database_id: str = Field(description="DMS databaseId"),
492+
question: Optional[str] = Field(default=None, description="Natural language question"),
493+
sql: str = Field(description="The SQL that caused an error"),
494+
error: str = Field(description="SQL error message"),
495+
model: Optional[str] = Field(default=None,
496+
description="Optional: if a specific model is desired, it can be specified here")
497+
) -> Dict[str, Any]:
498+
client = create_client()
499+
req = dms_enterprise_20181101_models.FixSqlByMetaAgentRequest(db_id=database_id, query=question, sql=sql,
500+
error=error)
501+
# if mcp.state.real_login_uid:
502+
# req.real_login_user_uid = mcp.state.real_login_uid
503+
if model:
504+
req.model = model
505+
try:
506+
resp = client.fix_sql_by_meta_agent(req)
507+
if not resp or not resp.body:
508+
return None
509+
data = resp.body.to_map()
510+
return data
511+
except Exception as e:
512+
logger.error(f"Error in fix_sql_statement: {e}")
513+
raise
514+
515+
516+
async def optimize_sql(
517+
database_id: str = Field(description="DMS databaseId"),
518+
question: Optional[str] = Field(default=None, description="Natural language question"),
519+
sql: str = Field(description="SQL statement"),
520+
model: Optional[str] = Field(default=None,
521+
description="Optional: if a specific model is desired, it can be specified here")
522+
) -> Any:
523+
client = create_client()
524+
req = dms_enterprise_20181101_models.OptimizeSqlByMetaAgentRequest(db_id=database_id, query=question, sql=sql)
525+
# if mcp.state.real_login_uid:
526+
# req.real_login_user_uid = mcp.state.real_login_uid
527+
if model:
528+
req.model = model
529+
try:
530+
resp = client.optimize_sql_by_meta_agent(req)
531+
if not resp or not resp.body:
532+
return None
533+
data = resp.body.to_map()
534+
return data
535+
except Exception as e:
536+
logger.error(f"Error in optimize_sql: {e}")
537+
raise
538+
539+
457540
# --- ToolRegistry Class ---
458541
class ToolRegistry:
459542
def __init__(self, mcp: FastMCP):
@@ -526,7 +609,8 @@ async def ask_database_configured(
526609
description="Your question in natural language about the pre-configured database."),
527610
knowledge: Optional[str] = Field(default=None,
528611
description="Optional: additional context to help formulate the SQL query."),
529-
model: Optional[str] = Field(default=None, description="Optional: if a specific model is desired, it can be specified here")
612+
model: Optional[str] = Field(default=None,
613+
description="Optional: if a specific model is desired, it can be specified here")
530614
) -> AskDatabaseResult:
531615
sql_result_obj = await nl2sql(database_id=self.default_database_id, question=question,
532616
knowledge=knowledge, model=model)
@@ -547,6 +631,45 @@ async def ask_database_configured(
547631
return AskDatabaseResult(executed_sql=generated_sql,
548632
execution_result=f"Error: An issue occurred while executing the query: {str(e)}")
549633

634+
@self.mcp.tool(name="answerSqlSyntax",
635+
description="Answer syntax-related questions for the corresponding database engine ",
636+
annotations={"title": "SQL语法回答", "readOnlyHint": True, "destructiveHint": False})
637+
async def answer_sql_syntax_configured(
638+
question: str = Field(description="Natural language question"),
639+
model: Optional[str] = Field(default=None,
640+
description="Optional: if a specific model is desired, it can be specified here")
641+
) -> Dict[str, Any]:
642+
result_obj = await answer_sql_syntax(database_id=self.default_database_id, question=question,
643+
model=model)
644+
return result_obj
645+
646+
@self.mcp.tool(name="fixSql",
647+
description="Analyze and fix the SQL error based on the provided SQL statement and error message.",
648+
annotations={"title": "SQL修复", "readOnlyHint": True, "destructiveHint": False})
649+
async def fix_sql_configured(
650+
question: Optional[str] = Field(default=None, description="Natural language question"),
651+
sql: str = Field(description="The SQL that caused an error"),
652+
error: str = Field(description="SQL error message"),
653+
model: Optional[str] = Field(default=None,
654+
description="Optional: if a specific model is desired, it can be specified here")
655+
) -> Dict[str, Any]:
656+
result_obj = await fix_sql_statement(database_id=self.default_database_id, question=question, sql=sql,
657+
error=error, model=model)
658+
return result_obj
659+
660+
@self.mcp.tool(name="optimizeSql",
661+
description="Analyze and optimize SQL performance based on the provided SQL statement",
662+
annotations={"title": "SQL优化", "readOnlyHint": True, "destructiveHint": False})
663+
async def optimize_sql_configured(
664+
question: Optional[str] = Field(default=None, description="Natural language question"),
665+
sql: str = Field(description="SQL statement"),
666+
model: Optional[str] = Field(default=None,
667+
description="Optional: if a specific model is desired, it can be specified here")
668+
) -> Any:
669+
result_obj = await optimize_sql(database_id=self.default_database_id, question=question, sql=sql,
670+
model=model)
671+
return result_obj
672+
550673
def _register_full_toolset(self):
551674
self.mcp.tool(name="addInstance",
552675
description="Add an instance to DMS. The username and password are required. "
@@ -567,25 +690,17 @@ def _register_full_toolset(self):
567690
description="Obtain detailed information about a specific database in DMS when the host and port are provided.",
568691
annotations={"title": "获取DMS数据库详情", "readOnlyHint": True})(get_database)
569692
self.mcp.tool(name="listTables",
570-
description="Search for tables by databaseId and (optional) table name. "
571-
"If you don't know the databaseId, first use getDatabase or searchDatabase to retrieve it."
572-
"(1)If you have the exact host, port, and database name, use getDatabase."
573-
"(2)If you only know the database name, use searchDatabase."
574-
"(3)If you don't know any information, ask the user to provide the necessary details."
575-
"Note: searchDatabase may return multiple databases. In this case, let the user choose which one to use.",
693+
description=f"Search for tables by databaseId and (optional) table name. "
694+
f"{DATABASE_ID_DESCRIPTION}",
576695
annotations={"title": "列出DMS表", "readOnlyHint": True})(list_tables)
577696
self.mcp.tool(name="getTableDetailInfo",
578697
description="Retrieve detailed metadata information about a specific database table including "
579698
"schema and index details. If you don't know the table_guid parameter, retrieve it using listTables.",
580699
annotations={"title": "获取DMS表详细信息", "readOnlyHint": True})(get_meta_table_detail_info)
581700

582701
@self.mcp.tool(name="executeScript",
583-
description="Execute SQL script against a database in DMS and return structured results."
584-
"If you don't know the databaseId, first use getDatabase or searchDatabase to retrieve it."
585-
"(1)If you have the exact host, port, and database name, use getDatabase."
586-
"(2)If you only know the database name, use searchDatabase."
587-
"(3)If you don't know any information, ask the user to provide the necessary details."
588-
"Note: searchDatabase may return multiple databases. In this case, let the user choose which one to use.",
702+
description=f"Execute SQL script against a database in DMS and return structured results."
703+
f"{DATABASE_ID_DESCRIPTION}",
589704
annotations={"title": "在DMS中执行SQL脚本", "readOnlyHint": False, "destructiveHint": True})
590705
async def execute_script_full_wrapper(
591706
database_id: str = Field(description="Required DMS databaseId. Obtained via getDatabase tool"),
@@ -596,13 +711,10 @@ async def execute_script_full_wrapper(
596711
return str(result_obj)
597712

598713
@self.mcp.tool(name="createDataChangeOrder",
599-
description="Execute SQL changes through a data change order, and a corresponding order ID will be returned. "
600-
"Prefer using the executeScript tool for SQL execution; only use this tool when explicitly instructed to perform the operation via a order."
601-
"If you don't know the databaseId, first use getDatabase or searchDatabase to retrieve it."
602-
"(1)If you have the exact host, port, and database name, use getDatabase."
603-
"(2)If you only know the database name, use searchDatabase."
604-
"(3)If you don't know any information, ask the user to provide the necessary details."
605-
"Note: searchDatabase may return multiple databases. In this case, let the user choose which one to use.",
714+
description=f"Execute SQL changes through a data change order, and a corresponding order ID will be returned. "
715+
f"Prefer using the executeScript tool for SQL execution;"
716+
f"only use this tool when explicitly instructed to perform the operation via a order."
717+
f"{DATABASE_ID_DESCRIPTION}",
606718
annotations={"title": "在DMS中创建数据变更工单", "readOnlyHint": False, "destructiveHint": True})
607719
async def create_data_change_order_wrapper(
608720
database_id: str = Field(description="Required DMS databaseId. Obtained via getDatabase tool"),
@@ -621,6 +733,19 @@ async def create_data_change_order_wrapper(
621733

622734
self.mcp.tool(name="generateSql", description="Generate SELECT-type SQL queries from natural language input.",
623735
annotations={"title": "自然语言转SQL (DMS)", "readOnlyHint": True})(nl2sql)
736+
self.mcp.tool(name="fixSql", description=f"Analyze and fix the SQL error based on the provided "
737+
f"SQL statement, error message, and database ID."
738+
f"{DATABASE_ID_DESCRIPTION}",
739+
annotations={"title": "SQL修复", "readOnlyHint": True})(fix_sql_statement)
740+
self.mcp.tool(name="answerSqlSyntax", description=f"Answer syntax-related questions "
741+
f"for the corresponding database engine "
742+
f"based on the database ID."
743+
f"{DATABASE_ID_DESCRIPTION}",
744+
annotations={"title": "SQL语法回答", "readOnlyHint": True})(answer_sql_syntax)
745+
self.mcp.tool(name="optimizeSql", description=f"Analyze and optimize SQL performance "
746+
f"based on the provided SQL statement and database ID"
747+
f"{DATABASE_ID_DESCRIPTION}",
748+
annotations={"title": "SQL优化", "readOnlyHint": True})(optimize_sql)
624749

625750

626751
# --- Lifespan Function ---

‎uv.lock‎

Lines changed: 5 additions & 5 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)