516 lines
21 KiB
Python
516 lines
21 KiB
Python
"""
|
|
Apache Doris MCP Server - Enterprise Database Service Implementation
|
|
|
|
Based on Apache Doris official MCP Server architecture design, providing complete MCP protocol support
|
|
Supports independent encapsulation implementation of Resources, Tools, and Prompts
|
|
Supports both stdio and streamable HTTP startup modes
|
|
"""
|
|
|
|
import argparse
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
from typing import Any
|
|
|
|
from mcp.server import Server
|
|
from mcp.server.models import InitializationOptions
|
|
|
|
from mcp.types import (
|
|
Prompt,
|
|
Resource,
|
|
TextContent,
|
|
Tool,
|
|
)
|
|
|
|
from .tools.tools_manager import DorisToolsManager
|
|
from .tools.prompts_manager import DorisPromptsManager
|
|
from .tools.resources_manager import DorisResourcesManager
|
|
from .utils.config import DorisConfig
|
|
from .utils.db import DorisConnectionManager
|
|
from .utils.security import DorisSecurityManager
|
|
|
|
# Configure logging
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class DorisServer:
|
|
"""Apache Doris MCP Server main class"""
|
|
|
|
def __init__(self, config: DorisConfig):
|
|
self.config = config
|
|
self.server = Server("doris-mcp-server")
|
|
|
|
# Initialize security manager
|
|
self.security_manager = DorisSecurityManager(config)
|
|
|
|
# Initialize connection manager, pass in security manager
|
|
self.connection_manager = DorisConnectionManager(config, self.security_manager)
|
|
|
|
# Initialize independent managers
|
|
self.resources_manager = DorisResourcesManager(self.connection_manager)
|
|
self.tools_manager = DorisToolsManager(self.connection_manager)
|
|
self.prompts_manager = DorisPromptsManager(self.connection_manager)
|
|
|
|
self.logger = logging.getLogger(f"{__name__}.DorisServer")
|
|
self._setup_handlers()
|
|
|
|
def _setup_handlers(self):
|
|
"""Setup MCP protocol handlers"""
|
|
|
|
@self.server.list_resources()
|
|
async def handle_list_resources() -> list[Resource]:
|
|
"""Handle resource list request"""
|
|
try:
|
|
self.logger.info("Handling resource list request")
|
|
resources = await self.resources_manager.list_resources()
|
|
self.logger.info(f"Returning {len(resources)} resources")
|
|
return resources
|
|
except Exception as e:
|
|
self.logger.error(f"Failed to handle resource list request: {e}")
|
|
return []
|
|
|
|
@self.server.read_resource()
|
|
async def handle_read_resource(uri: str) -> str:
|
|
"""Handle resource read request"""
|
|
try:
|
|
self.logger.info(f"Handling resource read request: {uri}")
|
|
content = await self.resources_manager.read_resource(uri)
|
|
return content
|
|
except Exception as e:
|
|
self.logger.error(f"Failed to handle resource read request: {e}")
|
|
return json.dumps(
|
|
{"error": f"Failed to read resource: {str(e)}", "uri": uri},
|
|
ensure_ascii=False,
|
|
indent=2,
|
|
)
|
|
|
|
@self.server.list_tools()
|
|
async def handle_list_tools() -> list[Tool]:
|
|
"""Handle tool list request"""
|
|
try:
|
|
self.logger.info("Handling tool list request")
|
|
tools = await self.tools_manager.list_tools()
|
|
self.logger.info(f"Returning {len(tools)} tools")
|
|
return tools
|
|
except Exception as e:
|
|
self.logger.error(f"Failed to handle tool list request: {e}")
|
|
return []
|
|
|
|
@self.server.call_tool()
|
|
async def handle_call_tool(
|
|
name: str, arguments: dict[str, Any]
|
|
) -> list[TextContent]:
|
|
"""Handle tool call request"""
|
|
try:
|
|
self.logger.info(f"Handling tool call request: {name}")
|
|
result = await self.tools_manager.call_tool(name, arguments)
|
|
|
|
return [TextContent(type="text", text=result)]
|
|
except Exception as e:
|
|
self.logger.error(f"Failed to handle tool call request: {e}")
|
|
error_result = json.dumps(
|
|
{
|
|
"error": f"Tool call failed: {str(e)}",
|
|
"tool_name": name,
|
|
"arguments": arguments,
|
|
},
|
|
ensure_ascii=False,
|
|
indent=2,
|
|
)
|
|
|
|
return [TextContent(type="text", text=error_result)]
|
|
|
|
@self.server.list_prompts()
|
|
async def handle_list_prompts() -> list[Prompt]:
|
|
"""Handle prompt list request"""
|
|
try:
|
|
self.logger.info("Handling prompt list request")
|
|
prompts = await self.prompts_manager.list_prompts()
|
|
self.logger.info(f"Returning {len(prompts)} prompts")
|
|
return prompts
|
|
except Exception as e:
|
|
self.logger.error(f"Failed to handle prompt list request: {e}")
|
|
return []
|
|
|
|
@self.server.get_prompt()
|
|
async def handle_get_prompt(name: str, arguments: dict[str, Any]) -> str:
|
|
"""Handle prompt get request"""
|
|
try:
|
|
self.logger.info(f"Handling prompt get request: {name}")
|
|
result = await self.prompts_manager.get_prompt(name, arguments)
|
|
return result
|
|
except Exception as e:
|
|
self.logger.error(f"Failed to handle prompt get request: {e}")
|
|
error_result = json.dumps(
|
|
{
|
|
"error": f"Failed to get prompt: {str(e)}",
|
|
"prompt_name": name,
|
|
"arguments": arguments,
|
|
},
|
|
ensure_ascii=False,
|
|
indent=2,
|
|
)
|
|
return error_result
|
|
|
|
async def start_stdio(self):
|
|
"""Start stdio transport mode"""
|
|
self.logger.info("Starting Doris MCP Server (stdio mode)")
|
|
|
|
try:
|
|
# Ensure connection manager is initialized
|
|
await self.connection_manager.initialize()
|
|
self.logger.info("Connection manager initialization completed")
|
|
|
|
# Start stdio server - using simpler approach
|
|
from mcp.server.stdio import stdio_server
|
|
|
|
self.logger.info("Creating stdio_server transport...")
|
|
|
|
# Try different startup approaches
|
|
try:
|
|
async with stdio_server() as streams:
|
|
read_stream, write_stream = streams
|
|
self.logger.info("stdio_server streams created successfully")
|
|
|
|
# Create initialization options
|
|
# MCP 1.8.0 requires parameters for get_capabilities
|
|
from mcp.server.lowlevel.server import NotificationOptions
|
|
|
|
capabilities = self.server.get_capabilities(
|
|
notification_options=NotificationOptions(
|
|
prompts_changed=True,
|
|
resources_changed=True,
|
|
tools_changed=True
|
|
),
|
|
experimental_capabilities={}
|
|
)
|
|
|
|
init_options = InitializationOptions(
|
|
server_name="doris-mcp-server",
|
|
server_version="1.0.0",
|
|
capabilities=capabilities,
|
|
)
|
|
self.logger.info("Initialization options created successfully")
|
|
|
|
# Run server
|
|
self.logger.info("Starting to run MCP server...")
|
|
await self.server.run(read_stream, write_stream, init_options)
|
|
|
|
except Exception as inner_e:
|
|
self.logger.error(f"stdio_server internal error: {inner_e}")
|
|
self.logger.error(f"Error type: {type(inner_e)}")
|
|
|
|
# Try to get more error information
|
|
import traceback
|
|
self.logger.error("Complete error stack:")
|
|
self.logger.error(traceback.format_exc())
|
|
|
|
# If it's ExceptionGroup, try to parse
|
|
if hasattr(inner_e, 'exceptions'):
|
|
self.logger.error(f"ExceptionGroup contains {len(inner_e.exceptions)} exceptions:")
|
|
for i, exc in enumerate(inner_e.exceptions):
|
|
self.logger.error(f" Exception {i+1}: {type(exc).__name__}: {exc}")
|
|
|
|
raise inner_e
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"stdio server startup failed: {e}")
|
|
self.logger.error(f"Error type: {type(e)}")
|
|
raise
|
|
|
|
|
|
|
|
async def start_http(self, host: str = "localhost", port: int = 3000):
|
|
"""Start Streamable HTTP transport mode"""
|
|
self.logger.info(f"Starting Doris MCP Server (Streamable HTTP mode) - {host}:{port}")
|
|
|
|
try:
|
|
# Ensure connection manager is initialized
|
|
await self.connection_manager.initialize()
|
|
|
|
# Use Starlette and StreamableHTTPSessionManager according to official example
|
|
import uvicorn
|
|
import contextlib
|
|
from collections.abc import AsyncIterator
|
|
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
|
|
from starlette.applications import Starlette
|
|
from starlette.routing import Mount, Route
|
|
from starlette.responses import JSONResponse, Response
|
|
from starlette.types import Receive, Scope, Send
|
|
|
|
# Create session manager
|
|
session_manager = StreamableHTTPSessionManager(
|
|
app=self.server,
|
|
json_response=True, # Enable JSON response
|
|
stateless=False # Maintain session state
|
|
)
|
|
|
|
self.logger.info(f"StreamableHTTP session manager created, will start at http://{host}:{port}")
|
|
|
|
# Health check endpoint
|
|
async def health_check(request):
|
|
return JSONResponse({"status": "healthy", "service": "doris-mcp-server"})
|
|
|
|
# Lifecycle manager - simplified since we manage session_manager externally
|
|
@contextlib.asynccontextmanager
|
|
async def lifespan(app: Starlette) -> AsyncIterator[None]:
|
|
"""Context manager for managing application lifecycle"""
|
|
self.logger.info("Application started!")
|
|
try:
|
|
yield
|
|
finally:
|
|
self.logger.info("Application is shutting down...")
|
|
|
|
# Create ASGI application - use direct session manager as ASGI app
|
|
starlette_app = Starlette(
|
|
debug=True,
|
|
routes=[
|
|
Route("/health", health_check, methods=["GET"]),
|
|
],
|
|
lifespan=lifespan,
|
|
)
|
|
|
|
# Custom ASGI app that handles both /mcp and /mcp/ without redirects
|
|
async def mcp_app(scope, receive, send):
|
|
# Handle lifespan events
|
|
if scope["type"] == "lifespan":
|
|
await starlette_app(scope, receive, send)
|
|
return
|
|
|
|
# Handle HTTP requests
|
|
if scope["type"] == "http":
|
|
path = scope.get("path", "")
|
|
self.logger.info(f"Received request for path: {path}")
|
|
|
|
try:
|
|
# Handle health check
|
|
if path.startswith("/health"):
|
|
await starlette_app(scope, receive, send)
|
|
return
|
|
|
|
# Handle MCP requests - both /mcp and /mcp/ go to session manager
|
|
if path == "/mcp" or path.startswith("/mcp/"):
|
|
self.logger.info(f"Handling MCP request for path: {path}")
|
|
# Log request details for debugging
|
|
method = scope.get("method", "UNKNOWN")
|
|
headers = dict(scope.get("headers", []))
|
|
self.logger.info(f"MCP Request - Method: {method}")
|
|
self.logger.info(f"MCP Request - Headers: {headers}")
|
|
|
|
# Handle Dify compatibility for GET requests
|
|
if method == "GET":
|
|
accept_header = headers.get(b'accept', b'').decode('utf-8')
|
|
user_agent = headers.get(b'user-agent', b'').decode('utf-8')
|
|
|
|
|
|
|
|
# For other GET requests, try to add application/json to Accept header
|
|
if 'text/event-stream' in accept_header and 'application/json' not in accept_header:
|
|
self.logger.info("Adding application/json to Accept header for GET request")
|
|
# Modify headers to include both content types
|
|
new_headers = []
|
|
for name, value in scope.get("headers", []):
|
|
if name == b'accept':
|
|
# Add application/json to the accept header
|
|
new_value = value.decode('utf-8') + ', application/json'
|
|
new_headers.append((name, new_value.encode('utf-8')))
|
|
else:
|
|
new_headers.append((name, value))
|
|
# Update scope with modified headers
|
|
scope = dict(scope)
|
|
scope["headers"] = new_headers
|
|
self.logger.info(f"Modified Accept header to: {new_value}")
|
|
|
|
await session_manager.handle_request(scope, receive, send)
|
|
return
|
|
|
|
# 404 for other paths
|
|
self.logger.info(f"Path not found: {path}")
|
|
response = Response("Not Found", status_code=404)
|
|
await response(scope, receive, send)
|
|
except Exception as e:
|
|
self.logger.error(f"Error handling request for {path}: {e}")
|
|
import traceback
|
|
self.logger.error(traceback.format_exc())
|
|
response = Response("Internal Server Error", status_code=500)
|
|
await response(scope, receive, send)
|
|
else:
|
|
# For other scope types, just return
|
|
self.logger.warning(f"Unsupported scope type: {scope['type']}")
|
|
return
|
|
|
|
# Start uvicorn server with session manager lifecycle
|
|
config = uvicorn.Config(
|
|
app=mcp_app,
|
|
host=host,
|
|
port=port,
|
|
log_level="info"
|
|
)
|
|
server = uvicorn.Server(config)
|
|
|
|
# Run session manager and server together
|
|
async with session_manager.run():
|
|
self.logger.info("Session manager started, now starting HTTP server")
|
|
await server.serve()
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Streamable HTTP server startup failed: {e}")
|
|
import traceback
|
|
self.logger.error("Complete error stack:")
|
|
self.logger.error(traceback.format_exc())
|
|
|
|
# If it's ExceptionGroup, try to parse
|
|
if hasattr(e, 'exceptions'):
|
|
self.logger.error(f"ExceptionGroup contains {len(e.exceptions)} exceptions:")
|
|
for i, exc in enumerate(e.exceptions):
|
|
self.logger.error(f" Exception {i+1}: {type(exc).__name__}: {exc}")
|
|
raise
|
|
|
|
async def shutdown(self):
|
|
"""Shutdown server"""
|
|
self.logger.info("Shutting down Doris MCP Server")
|
|
try:
|
|
await self.connection_manager.close()
|
|
self.logger.info("Doris MCP Server has been shut down")
|
|
except Exception as e:
|
|
self.logger.error(f"Error occurred while shutting down server: {e}")
|
|
|
|
|
|
def create_arg_parser():
|
|
"""Create command line argument parser"""
|
|
parser = argparse.ArgumentParser(
|
|
description="Apache Doris MCP Server - Enterprise Database Service",
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
epilog="""
|
|
Transport Modes:
|
|
stdio - Standard input/output (for local process communication)
|
|
http - Streamable HTTP mode (MCP 2025-03-26 protocol)
|
|
|
|
Examples:
|
|
python -m doris_mcp_server --transport stdio
|
|
python -m doris_mcp_server --transport http --host 0.0.0.0 --port 3000
|
|
"""
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--transport",
|
|
type=str,
|
|
choices=["stdio", "http"],
|
|
default="stdio",
|
|
help="Transport protocol type: stdio (local), http (Streamable HTTP)",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--host",
|
|
type=str,
|
|
default="localhost",
|
|
help="Host address for HTTP mode (default: localhost)",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--port", type=int, default=3000, help="Port number for HTTP mode (default: 3000)"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--db-host",
|
|
type=str,
|
|
default="localhost",
|
|
help="Doris database host address (default: localhost)",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--db-port", type=int, default=9030, help="Doris database port number (default: 9030)"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--db-user", type=str, default="root", help="Doris database username (default: root)"
|
|
)
|
|
|
|
parser.add_argument("--db-password", type=str, default="", help="Doris database password")
|
|
|
|
parser.add_argument(
|
|
"--db-database",
|
|
type=str,
|
|
default="information_schema",
|
|
help="Doris database name (default: information_schema)",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--log-level",
|
|
type=str,
|
|
choices=["DEBUG", "INFO", "WARNING", "ERROR"],
|
|
default="INFO",
|
|
help="Log level (default: INFO)",
|
|
)
|
|
|
|
return parser
|
|
|
|
|
|
async def main():
|
|
"""Main function"""
|
|
parser = create_arg_parser()
|
|
args = parser.parse_args()
|
|
|
|
# Set log level
|
|
logging.getLogger().setLevel(getattr(logging, args.log_level))
|
|
|
|
# Create configuration - priority: command line arguments > .env file > default values
|
|
config = DorisConfig.from_env() # First load from .env file and environment variables
|
|
|
|
# Command line arguments override configuration (if provided)
|
|
if args.db_host != "localhost": # If not default value, use command line argument
|
|
config.database.host = args.db_host
|
|
if args.db_port != 9030:
|
|
config.database.port = args.db_port
|
|
if args.db_user != "root":
|
|
config.database.user = args.db_user
|
|
if args.db_password: # Use password if provided
|
|
config.database.password = args.db_password
|
|
if args.db_database != "information_schema":
|
|
config.database.database = args.db_database
|
|
if args.log_level != "INFO":
|
|
config.logging.level = args.log_level
|
|
|
|
# Create server instance
|
|
server = DorisServer(config)
|
|
|
|
try:
|
|
if args.transport == "stdio":
|
|
await server.start_stdio()
|
|
elif args.transport == "http":
|
|
await server.start_http(args.host, args.port)
|
|
else:
|
|
logger.error(f"Unsupported transport protocol: {args.transport}")
|
|
await server.shutdown()
|
|
return 1
|
|
|
|
except KeyboardInterrupt:
|
|
logger.info("Received interrupt signal, shutting down server...")
|
|
except Exception as e:
|
|
logger.error(f"Server runtime error: {e}")
|
|
# Clean up resources even in case of exception
|
|
try:
|
|
await server.shutdown()
|
|
except Exception as shutdown_error:
|
|
logger.error(f"Error occurred while shutting down server: {shutdown_error}")
|
|
return 1
|
|
finally:
|
|
# Cleanup in case of normal shutdown
|
|
try:
|
|
await server.shutdown()
|
|
except Exception as shutdown_error:
|
|
logger.error(f"Error occurred while shutting down server: {shutdown_error}")
|
|
|
|
return 0
|
|
|
|
|
|
def main_sync():
|
|
"""Synchronous main function for entry point"""
|
|
exit_code = asyncio.run(main())
|
|
exit(exit_code)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main_sync()
|