157 lines
9.9 KiB
Python
157 lines
9.9 KiB
Python
#!/usr/bin/env python
|
|
# -*- coding: utf-8 -*-
|
|
|
|
"""
|
|
Tool Initialization Module
|
|
|
|
Centralized initialization of all tools, ensuring they are correctly registered with MCP
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
from typing import List, Dict, Any, Optional
|
|
import json
|
|
from datetime import datetime
|
|
import traceback
|
|
|
|
# Import Context
|
|
from mcp.server.fastmcp import Context
|
|
|
|
# Import doris mcp tools
|
|
from doris_mcp_server.tools.mcp_doris_tools import (
|
|
mcp_doris_exec_query,
|
|
mcp_doris_get_table_schema,
|
|
mcp_doris_get_db_table_list,
|
|
mcp_doris_get_db_list,
|
|
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_catalog_list
|
|
)
|
|
|
|
# Get logger
|
|
logger = logging.getLogger("doris-mcp-tools-initializer")
|
|
|
|
async def register_mcp_tools(mcp):
|
|
"""Register MCP tool functions
|
|
|
|
Args:
|
|
mcp: FastMCP instance
|
|
"""
|
|
logger.info("Starting to register MCP tools...")
|
|
|
|
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 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. 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, 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, 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
|
|
- 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, 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
|
|
- 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, 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
|
|
- 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(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
|
|
- 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, 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
|
|
- 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, 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
|
|
- 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, 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
|
|
[Parameter Content]:\n
|
|
- random_string (string) [Required] - Unique identifier for the tool call\n
|
|
- days (integer) [Optional] - Number of recent days of logs to retrieve, default is 7\n
|
|
- limit (integer) [Optional] - Maximum number of records to return, default is 100\n""")
|
|
async def get_recent_audit_logs_tool(days: int = 7, limit: int = 100) -> Dict[str, Any]:
|
|
"""Wrapper: Get recent audit logs"""
|
|
try:
|
|
days = int(days)
|
|
limit = int(limit)
|
|
except (ValueError, TypeError):
|
|
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")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Error registering MCP tools: {str(e)}")
|
|
logger.error(traceback.format_exc())
|
|
return False |