support Multi-Catalog
This commit is contained in:
FreeOnePlus
@@ -26,7 +26,8 @@ from doris_mcp_server.tools.mcp_doris_tools import (
|
||||
mcp_doris_get_table_comment,
|
||||
mcp_doris_get_table_column_comments,
|
||||
mcp_doris_get_table_indexes,
|
||||
mcp_doris_get_recent_audit_logs
|
||||
mcp_doris_get_recent_audit_logs,
|
||||
mcp_doris_get_catalog_list
|
||||
)
|
||||
|
||||
# Get logger
|
||||
@@ -42,79 +43,86 @@ async def register_mcp_tools(mcp):
|
||||
|
||||
try:
|
||||
# Register Tool: Execute SQL Query (Using long description string including parameters)
|
||||
@mcp.tool("exec_query", description="""[Function Description]: Execute SQL query and return result command (executed by the client).\n
|
||||
@mcp.tool("exec_query", description="""[Function Description]: Execute SQL query and return result command with catalog federation support.\n
|
||||
[Parameter Content]:\n
|
||||
- random_string (string) [Required] - Unique identifier for the tool call\n
|
||||
- sql (string) [Required] - SQL statement to execute\n
|
||||
- sql (string) [Required] - SQL statement to execute. MUST use three-part naming for all table references: 'catalog_name.db_name.table_name'. For internal tables use 'internal.db_name.table_name', for external tables use 'catalog_name.db_name.table_name'\n
|
||||
- db_name (string) [Optional] - Target database name, defaults to the current database\n
|
||||
- catalog_name (string) [Optional] - Reference catalog name for context, defaults to current catalog\n
|
||||
- max_rows (integer) [Optional] - Maximum number of rows to return, default 100
|
||||
- timeout (integer) [Optional] - Query timeout in seconds, default 30""")
|
||||
async def exec_query_tool(sql: str, db_name: str = None, max_rows: int = 100, timeout: int = 30) -> Dict[str, Any]:
|
||||
async def exec_query_tool(sql: str, db_name: str = None, catalog_name: str = None, max_rows: int = 100, timeout: int = 30) -> Dict[str, Any]:
|
||||
"""Wrapper: Execute SQL query and return result command"""
|
||||
# Note: ctx parameter is no longer needed here as we receive named parameters directly
|
||||
return await mcp_doris_exec_query(sql=sql, db_name=db_name, max_rows=max_rows, timeout=timeout)
|
||||
return await mcp_doris_exec_query(sql=sql, db_name=db_name, catalog_name=catalog_name, max_rows=max_rows, timeout=timeout)
|
||||
|
||||
# Register Tool: Get Table Schema (Keep long description string including parameters)
|
||||
@mcp.tool("get_table_schema", description="""[Function Description]: Get detailed structure information of the specified table (columns, types, comments, etc.).\n
|
||||
[Parameter Content]:\n
|
||||
- random_string (string) [Required] - Unique identifier for the tool call\n
|
||||
- table_name (string) [Required] - Name of the table to query\n
|
||||
- db_name (string) [Optional] - Target database name, defaults to the current database\n""")
|
||||
async def get_table_schema_tool(table_name: str, db_name: str = None) -> Dict[str, Any]:
|
||||
- db_name (string) [Optional] - Target database name, defaults to the current database\n
|
||||
- catalog_name (string) [Optional] - Target catalog name for federation queries, defaults to current catalog\n""")
|
||||
async def get_table_schema_tool(table_name: str, db_name: str = None, catalog_name: str = None) -> Dict[str, Any]:
|
||||
"""Wrapper: Get table schema"""
|
||||
if not table_name: return {"content": [{"type": "text", "text": json.dumps({"success": False, "error": "Missing table_name parameter"})}]}
|
||||
return await mcp_doris_get_table_schema(table_name=table_name, db_name=db_name)
|
||||
return await mcp_doris_get_table_schema(table_name=table_name, db_name=db_name, catalog_name=catalog_name)
|
||||
|
||||
# Register Tool: Get Database Table List (Keep long description string including parameters)
|
||||
@mcp.tool("get_db_table_list", description="""[Function Description]: Get a list of all table names in the specified database.\n
|
||||
[Parameter Content]:\n
|
||||
- random_string (string) [Required] - Unique identifier for the tool call\n
|
||||
- db_name (string) [Optional] - Target database name, defaults to the current database\n""")
|
||||
async def get_db_table_list_tool(db_name: str = None) -> Dict[str, Any]:
|
||||
- db_name (string) [Optional] - Target database name, defaults to the current database\n
|
||||
- catalog_name (string) [Optional] - Target catalog name for federation queries, defaults to current catalog\n""")
|
||||
async def get_db_table_list_tool(db_name: str = None, catalog_name: str = None) -> Dict[str, Any]:
|
||||
"""Wrapper: Get database table list"""
|
||||
return await mcp_doris_get_db_table_list(db_name=db_name)
|
||||
return await mcp_doris_get_db_table_list(db_name=db_name, catalog_name=catalog_name)
|
||||
|
||||
# Register Tool: Get Database List (Keep long description string including parameters)
|
||||
# Note: Although the description mentions random_string, the wrapper function signature does not. See how mcp handles this.
|
||||
@mcp.tool("get_db_list", description="""[Function Description]: Get a list of all database names on the server.\n
|
||||
[Parameter Content]:\n
|
||||
- random_string (string) [Required] - Unique identifier for the tool call\n""")
|
||||
async def get_db_list_tool() -> Dict[str, Any]: # Function signature has no parameters
|
||||
- random_string (string) [Required] - Unique identifier for the tool call\n
|
||||
- catalog_name (string) [Optional] - Target catalog name for federation queries, defaults to current catalog\n""")
|
||||
async def get_db_list_tool(catalog_name: str = None) -> Dict[str, Any]: # Function signature has no parameters
|
||||
"""Wrapper: Get database list"""
|
||||
return await mcp_doris_get_db_list()
|
||||
return await mcp_doris_get_db_list(catalog_name=catalog_name)
|
||||
|
||||
# Register Tool: Get Table Comment (Keep long description string including parameters)
|
||||
@mcp.tool("get_table_comment", description="""[Function Description]: Get the comment information for the specified table.\n
|
||||
[Parameter Content]:\n
|
||||
- random_string (string) [Required] - Unique identifier for the tool call\n
|
||||
- table_name (string) [Required] - Name of the table to query\n
|
||||
- db_name (string) [Optional] - Target database name, defaults to the current database\n""")
|
||||
async def get_table_comment_tool(table_name: str, db_name: str = None) -> Dict[str, Any]:
|
||||
- db_name (string) [Optional] - Target database name, defaults to the current database\n
|
||||
- catalog_name (string) [Optional] - Target catalog name for federation queries, defaults to current catalog\n""")
|
||||
async def get_table_comment_tool(table_name: str, db_name: str = None, catalog_name: str = None) -> Dict[str, Any]:
|
||||
"""Wrapper: Get table comment"""
|
||||
if not table_name: return {"content": [{"type": "text", "text": json.dumps({"success": False, "error": "Missing table_name parameter"})}]}
|
||||
return await mcp_doris_get_table_comment(table_name=table_name, db_name=db_name)
|
||||
return await mcp_doris_get_table_comment(table_name=table_name, db_name=db_name, catalog_name=catalog_name)
|
||||
|
||||
# Register Tool: Get Table Column Comments (Keep long description string including parameters)
|
||||
@mcp.tool("get_table_column_comments", description="""[Function Description]: Get comment information for all columns in the specified table.\n
|
||||
[Parameter Content]:\n
|
||||
- random_string (string) [Required] - Unique identifier for the tool call\n
|
||||
- table_name (string) [Required] - Name of the table to query\n
|
||||
- db_name (string) [Optional] - Target database name, defaults to the current database\n""")
|
||||
async def get_table_column_comments_tool(table_name: str, db_name: str = None) -> Dict[str, Any]:
|
||||
- db_name (string) [Optional] - Target database name, defaults to the current database\n
|
||||
- catalog_name (string) [Optional] - Target catalog name for federation queries, defaults to current catalog\n""")
|
||||
async def get_table_column_comments_tool(table_name: str, db_name: str = None, catalog_name: str = None) -> Dict[str, Any]:
|
||||
"""Wrapper: Get table column comments"""
|
||||
if not table_name: return {"content": [{"type": "text", "text": json.dumps({"success": False, "error": "Missing table_name parameter"})}]}
|
||||
return await mcp_doris_get_table_column_comments(table_name=table_name, db_name=db_name)
|
||||
return await mcp_doris_get_table_column_comments(table_name=table_name, db_name=db_name, catalog_name=catalog_name)
|
||||
|
||||
# Register Tool: Get Table Indexes (Keep long description string including parameters)
|
||||
@mcp.tool("get_table_indexes", description="""[Function Description]: Get index information for the specified table.\n
|
||||
[Parameter Content]:\n
|
||||
- random_string (string) [Required] - Unique identifier for the tool call\n
|
||||
- table_name (string) [Required] - Name of the table to query\n
|
||||
- db_name (string) [Optional] - Target database name, defaults to the current database\n""")
|
||||
async def get_table_indexes_tool(table_name: str, db_name: str = None) -> Dict[str, Any]:
|
||||
- db_name (string) [Optional] - Target database name, defaults to the current database\n
|
||||
- catalog_name (string) [Optional] - Target catalog name for federation queries, defaults to current catalog\n""")
|
||||
async def get_table_indexes_tool(table_name: str, db_name: str = None, catalog_name: str = None) -> Dict[str, Any]:
|
||||
"""Wrapper: Get table indexes"""
|
||||
if not table_name: return {"content": [{"type": "text", "text": json.dumps({"success": False, "error": "Missing table_name parameter"})}]}
|
||||
return await mcp_doris_get_table_indexes(table_name=table_name, db_name=db_name)
|
||||
return await mcp_doris_get_table_indexes(table_name=table_name, db_name=db_name, catalog_name=catalog_name)
|
||||
|
||||
# Register Tool: Get Recent Audit Logs (Keep long description string including parameters)
|
||||
@mcp.tool("get_recent_audit_logs", description="""[Function Description]: Get audit log records for a recent period.\n
|
||||
@@ -131,6 +139,14 @@ async def register_mcp_tools(mcp):
|
||||
return {"content": [{"type": "text", "text": json.dumps({"success": False, "error": "days and limit parameters must be integers"})}]}
|
||||
return await mcp_doris_get_recent_audit_logs(days=days, limit=limit)
|
||||
|
||||
# Register Tool: Get Catalog List (Keep long description string including parameters)
|
||||
@mcp.tool("get_catalog_list", description="""[Function Description]: Get a list of all catalog names on the server.\n
|
||||
[Parameter Content]:\n
|
||||
- random_string (string) [Required] - Unique identifier for the tool call\n""")
|
||||
async def get_catalog_list_tool() -> Dict[str, Any]:
|
||||
"""Wrapper: Get catalog list"""
|
||||
return await mcp_doris_get_catalog_list()
|
||||
|
||||
# Get tool count
|
||||
tools_count = len(await mcp.list_tools())
|
||||
logger.info(f"Registered all MCP tools, total {tools_count} tools")
|
||||
|
||||
Reference in New Issue
Block a user