AI_DATA/doris-mcp-server
6
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

support Multi-Catalog

Browse Source
This commit is contained in:
FreeOnePlus
2025-06-06 14:35:53 +08:00
parent 466fcbdb45
commit 5e98e5ba41
8 changed files with 312 additions and 151 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
doris_mcp_server/tools/__init__.py
View File

@@ -6,7 +6,8 @@ from .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
)
# The __all__ list should reflect the registered tool names,
@@ -19,5 +20,6 @@ __all__ = [
"get_table_comment",
"get_table_column_comments",
"get_table_indexes",
"get_recent_audit_logs"
"get_recent_audit_logs",
"get_catalog_list"
]

88
doris_mcp_server/tools/mcp_doris_tools.py
View File

@@ -55,20 +55,31 @@ def _format_response(success: bool, result: Any = None, error: str = None, messa
]
}
async def mcp_doris_exec_query(sql: str = None, db_name: str = None, max_rows: int = 100, timeout: int = 30) -> Dict[str, Any]:
async def mcp_doris_exec_query(sql: str = None, db_name: str = None, catalog_name: str = None, max_rows: int = 100, timeout: int = 30) -> Dict[str, Any]:
"""
Executes an SQL query and returns the result.
Executes an SQL query and returns the result with catalog federation support.
Args:
sql (str): The SQL query to execute.
db_name (str, optional): Target database name. Defaults to the configured default database.
sql (str): The SQL query to execute. MUST use three-part naming for table references:
- Internal tables: internal.db_name.table_name (e.g., "SELECT * FROM internal.ssb.customer")
- External tables: catalog_name.db_name.table_name (e.g., "SELECT * FROM mysql.ssb.customer")
- Cross-catalog queries: "SELECT * FROM mysql.ssb.customer m JOIN internal.ssb.orders o ON m.id = o.customer_id"
Examples:
- Query internal catalog: "SELECT COUNT(*) FROM internal.ssb.customer"
- Query MySQL catalog: "SELECT COUNT(*) FROM mysql.ssb.customer"
- Cross-catalog join: "SELECT * FROM internal.ssb.customer c JOIN mysql.test.user_info u ON c.id = u.customer_id"
db_name (str, optional): Target database name. Only used for connection context, table names in SQL must be fully qualified.
catalog_name (str, optional): Reference catalog name for context. Does not affect SQL execution - table names in SQL must be fully qualified.
Available catalogs can be found using get_catalog_list tool.
max_rows (int, optional): Maximum number of rows to return. Defaults to 100.
timeout (int, optional): Query timeout in seconds. Defaults to 30.
Returns:
Dict[str, Any]: A dictionary containing the query result or an error.
"""
logger.info(f"MCP Tool Call: mcp_doris_exec_query, SQL: {sql}, DB: {db_name}, MaxRows: {max_rows}, Timeout: {timeout}")
logger.info(f"MCP Tool Call: mcp_doris_exec_query, SQL: {sql}, DB: {db_name}, Catalog: {catalog_name}, MaxRows: {max_rows}, Timeout: {timeout}")
try:
if not sql:
return _format_response(success=False, error="SQL statement not provided", message="Please provide the SQL statement to execute")
@@ -78,6 +89,7 @@ async def mcp_doris_exec_query(sql: str = None, db_name: str = None, max_rows: i
"params": {
"sql": sql,
"db_name": db_name,
"catalog_name": catalog_name,
"max_rows": max_rows,
"timeout": timeout
}
@@ -121,71 +133,71 @@ async def mcp_doris_exec_query(sql: str = None, db_name: str = None, max_rows: i
return _format_response(success=False, error=str(e), message="Error executing SQL query")
async def mcp_doris_get_table_schema(table_name: str, db_name: str = None) -> Dict[str, Any]:
logger.info(f"MCP Tool Call: mcp_doris_get_table_schema, Table: {table_name}, DB: {db_name}")
async def mcp_doris_get_table_schema(table_name: str, db_name: str = None, catalog_name: str = None) -> Dict[str, Any]:
logger.info(f"MCP Tool Call: mcp_doris_get_table_schema, Table: {table_name}, DB: {db_name}, Catalog: {catalog_name}")
if not table_name:
return _format_response(success=False, error="Missing table_name parameter")
try:
extractor = MetadataExtractor(db_name=db_name)
schema = extractor.get_table_schema(table_name=table_name, db_name=db_name)
extractor = MetadataExtractor(db_name=db_name, catalog_name=catalog_name)
schema = extractor.get_table_schema(table_name=table_name, db_name=db_name, catalog_name=catalog_name)
if not schema:
return _format_response(success=False, error="Table not found or has no columns", message=f"Could not get schema for table {db_name or extractor.db_name}.{table_name}")
return _format_response(success=False, error="Table not found or has no columns", message=f"Could not get schema for table {catalog_name or 'default'}.{db_name or extractor.db_name}.{table_name}")
return _format_response(success=True, result=schema)
except Exception as e:
logger.error(f"MCP tool execution failed mcp_doris_get_table_schema: {str(e)}", exc_info=True)
return _format_response(success=False, error=str(e), message="Error getting table schema")
async def mcp_doris_get_db_table_list(db_name: str = None) -> Dict[str, Any]:
logger.info(f"MCP Tool Call: mcp_doris_get_db_table_list, DB: {db_name}")
async def mcp_doris_get_db_table_list(db_name: str = None, catalog_name: str = None) -> Dict[str, Any]:
logger.info(f"MCP Tool Call: mcp_doris_get_db_table_list, DB: {db_name}, Catalog: {catalog_name}")
try:
extractor = MetadataExtractor(db_name=db_name)
tables = extractor.get_database_tables(db_name=db_name)
extractor = MetadataExtractor(db_name=db_name, catalog_name=catalog_name)
tables = extractor.get_database_tables(db_name=db_name, catalog_name=catalog_name)
return _format_response(success=True, result=tables)
except Exception as e:
logger.error(f"MCP tool execution failed mcp_doris_get_db_table_list: {str(e)}", exc_info=True)
return _format_response(success=False, error=str(e), message="Error getting database table list")
async def mcp_doris_get_db_list() -> Dict[str, Any]:
logger.info(f"MCP Tool Call: mcp_doris_get_db_list")
async def mcp_doris_get_db_list(catalog_name: str = None) -> Dict[str, Any]:
logger.info(f"MCP Tool Call: mcp_doris_get_db_list, Catalog: {catalog_name}")
try:
extractor = MetadataExtractor()
databases = extractor.get_all_databases()
extractor = MetadataExtractor(catalog_name=catalog_name)
databases = extractor.get_all_databases(catalog_name=catalog_name)
return _format_response(success=True, result=databases)
except Exception as e:
logger.error(f"MCP tool execution failed mcp_doris_get_db_list: {str(e)}", exc_info=True)
return _format_response(success=False, error=str(e), message="Error getting database list")
async def mcp_doris_get_table_comment(table_name: str, db_name: str = None) -> Dict[str, Any]:
logger.info(f"MCP Tool Call: mcp_doris_get_table_comment, Table: {table_name}, DB: {db_name}")
async def mcp_doris_get_table_comment(table_name: str, db_name: str = None, catalog_name: str = None) -> Dict[str, Any]:
logger.info(f"MCP Tool Call: mcp_doris_get_table_comment, Table: {table_name}, DB: {db_name}, Catalog: {catalog_name}")
if not table_name:
return _format_response(success=False, error="Missing table_name parameter")
try:
extractor = MetadataExtractor(db_name=db_name)
comment = extractor.get_table_comment(table_name=table_name, db_name=db_name)
extractor = MetadataExtractor(db_name=db_name, catalog_name=catalog_name)
comment = extractor.get_table_comment(table_name=table_name, db_name=db_name, catalog_name=catalog_name)
return _format_response(success=True, result=comment)
except Exception as e:
logger.error(f"MCP tool execution failed mcp_doris_get_table_comment: {str(e)}", exc_info=True)
return _format_response(success=False, error=str(e), message="Error getting table comment")
async def mcp_doris_get_table_column_comments(table_name: str, db_name: str = None) -> Dict[str, Any]:
logger.info(f"MCP Tool Call: mcp_doris_get_table_column_comments, Table: {table_name}, DB: {db_name}")
async def mcp_doris_get_table_column_comments(table_name: str, db_name: str = None, catalog_name: str = None) -> Dict[str, Any]:
logger.info(f"MCP Tool Call: mcp_doris_get_table_column_comments, Table: {table_name}, DB: {db_name}, Catalog: {catalog_name}")
if not table_name:
return _format_response(success=False, error="Missing table_name parameter")
try:
extractor = MetadataExtractor(db_name=db_name)
comments = extractor.get_column_comments(table_name=table_name, db_name=db_name)
extractor = MetadataExtractor(db_name=db_name, catalog_name=catalog_name)
comments = extractor.get_column_comments(table_name=table_name, db_name=db_name, catalog_name=catalog_name)
return _format_response(success=True, result=comments)
except Exception as e:
logger.error(f"MCP tool execution failed mcp_doris_get_table_column_comments: {str(e)}", exc_info=True)
return _format_response(success=False, error=str(e), message="Error getting column comments")
async def mcp_doris_get_table_indexes(table_name: str, db_name: str = None) -> Dict[str, Any]:
logger.info(f"MCP Tool Call: mcp_doris_get_table_indexes, Table: {table_name}, DB: {db_name}")
async def mcp_doris_get_table_indexes(table_name: str, db_name: str = None, catalog_name: str = None) -> Dict[str, Any]:
logger.info(f"MCP Tool Call: mcp_doris_get_table_indexes, Table: {table_name}, DB: {db_name}, Catalog: {catalog_name}")
if not table_name:
return _format_response(success=False, error="Missing table_name parameter")
try:
extractor = MetadataExtractor(db_name=db_name)
indexes = extractor.get_table_indexes(table_name=table_name, db_name=db_name)
extractor = MetadataExtractor(db_name=db_name, catalog_name=catalog_name)
indexes = extractor.get_table_indexes(table_name=table_name, db_name=db_name, catalog_name=catalog_name)
return _format_response(success=True, result=indexes)
except Exception as e:
logger.error(f"MCP tool execution failed mcp_doris_get_table_indexes: {str(e)}", exc_info=True)
@@ -200,3 +212,19 @@ async def mcp_doris_get_recent_audit_logs(days: int = 7, limit: int = 100) -> Di
except Exception as e:
logger.error(f"MCP tool execution failed mcp_doris_get_recent_audit_logs: {str(e)}", exc_info=True)
return _format_response(success=False, error=str(e), message="Error getting audit logs")
async def mcp_doris_get_catalog_list() -> Dict[str, Any]:
"""
Get Doris catalog list
Returns:
Dict[str, Any]: Dictionary containing catalog list or error information
"""
logger.info(f"MCP Tool Call: mcp_doris_get_catalog_list")
try:
extractor = MetadataExtractor()
catalogs = extractor.get_catalog_list()
return _format_response(success=True, result=catalogs, message="Successfully retrieved catalog list")
except Exception as e:
logger.error(f"MCP tool execution failed mcp_doris_get_catalog_list: {str(e)}", exc_info=True)
return _format_response(success=False, error=str(e), message="Error getting catalog list")

62
doris_mcp_server/tools/tool_initializer.py
View File

@@ -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")