db-connect-mcp 0.1.0__py3-none-any.whl

db_connect_mcp/server.py

@@ -0,0 +1,496 @@
+"""Multi-database MCP Server
+
+A Model Context Protocol (MCP) server providing database analysis and querying
+capabilities for PostgreSQL, MySQL, and ClickHouse databases.
+"""
+
+import asyncio
+import json
+import logging
+import os
+from typing import Any, Optional
+
+from dotenv import load_dotenv
+from mcp.server import Server
+from mcp.types import TextContent, Tool
+
+from db_connect_mcp.adapters import create_adapter
+from db_connect_mcp.core import (
+    DatabaseConnection,
+    MetadataInspector,
+    QueryExecutor,
+    StatisticsAnalyzer,
+)
+from db_connect_mcp.models.config import DatabaseConfig
+
+# Load environment variables
+load_dotenv()
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class DatabaseMCPServer:
+    """MCP server for multi-database operations."""
+
+    def __init__(self, config: DatabaseConfig):
+        """
+        Initialize database MCP server.
+
+        Args:
+            config: Database configuration
+        """
+        self.config = config
+        self.connection = DatabaseConnection(config)
+        self.adapter = create_adapter(config)
+        self.inspector: Optional[MetadataInspector] = None
+        self.executor: Optional[QueryExecutor] = None
+        self.analyzer: Optional[StatisticsAnalyzer] = None
+        self.server = Server("db-mcp")
+
+    async def initialize(self) -> None:
+        """Initialize all components."""
+        await self.connection.initialize()
+
+        self.inspector = MetadataInspector(self.connection, self.adapter)
+        self.executor = QueryExecutor(self.connection, self.adapter)
+        self.analyzer = StatisticsAnalyzer(self.connection, self.adapter)
+
+        # Register MCP tool handlers
+        await self._register_tools()
+
+        logger.info(
+            f"Initialized {self.config.dialect} MCP server "
+            f"({len(self.adapter.capabilities.get_supported_features())} features)"
+        )
+
+    async def _register_tools(self) -> None:
+        """Register MCP tools based on database capabilities."""
+        # Note: Tools are registered via the list_tools decorator, not add_tool
+        # This method is kept for initializing any tool-related state
+        pass
+
+    def _create_get_database_info_tool(self) -> Tool:
+        """Create get_database_info tool."""
+        return Tool(
+            name="get_database_info",
+            description="Get database information including version, size, and capabilities",
+            inputSchema={
+                "type": "object",
+                "properties": {},
+                "required": [],
+            },
+        )
+
+    def _create_list_schemas_tool(self) -> Tool:
+        """Create list_schemas tool."""
+        return Tool(
+            name="list_schemas",
+            description="List all schemas/databases in the database instance",
+            inputSchema={
+                "type": "object",
+                "properties": {},
+                "required": [],
+            },
+        )
+
+    def _create_list_tables_tool(self) -> Tool:
+        """Create list_tables tool."""
+        return Tool(
+            name="list_tables",
+            description="List all tables and views in a schema",
+            inputSchema={
+                "type": "object",
+                "properties": {
+                    "schema": {
+                        "type": "string",
+                        "description": "Schema name (optional, uses default if not specified)",
+                    },
+                    "include_views": {
+                        "type": "boolean",
+                        "description": "Whether to include views (default: true)",
+                        "default": True,
+                    },
+                },
+                "required": [],
+            },
+        )
+
+    def _create_describe_table_tool(self) -> Tool:
+        """Create describe_table tool."""
+        return Tool(
+            name="describe_table",
+            description="Get comprehensive table information including columns, indexes, and constraints",
+            inputSchema={
+                "type": "object",
+                "properties": {
+                    "table": {"type": "string", "description": "Table name"},
+                    "schema": {
+                        "type": "string",
+                        "description": "Schema name (optional)",
+                    },
+                },
+                "required": ["table"],
+            },
+        )
+
+    def _create_execute_query_tool(self) -> Tool:
+        """Create execute_query tool."""
+        return Tool(
+            name="execute_query",
+            description="Execute a read-only SQL query (SELECT, WITH, EXPLAIN)",
+            inputSchema={
+                "type": "object",
+                "properties": {
+                    "query": {"type": "string", "description": "SQL query to execute"},
+                    "limit": {
+                        "type": "integer",
+                        "description": "Maximum number of rows to return (default: 1000)",
+                        "default": 1000,
+                    },
+                },
+                "required": ["query"],
+            },
+        )
+
+    def _create_sample_data_tool(self) -> Tool:
+        """Create sample_data tool."""
+        return Tool(
+            name="sample_data",
+            description="Sample data from a table efficiently",
+            inputSchema={
+                "type": "object",
+                "properties": {
+                    "table": {"type": "string", "description": "Table name"},
+                    "schema": {
+                        "type": "string",
+                        "description": "Schema name (optional)",
+                    },
+                    "limit": {
+                        "type": "integer",
+                        "description": "Number of rows to sample (default: 100)",
+                        "default": 100,
+                    },
+                },
+                "required": ["table"],
+            },
+        )
+
+    def _create_get_relationships_tool(self) -> Tool:
+        """Create get_table_relationships tool."""
+        return Tool(
+            name="get_table_relationships",
+            description="Get foreign key relationships for a table",
+            inputSchema={
+                "type": "object",
+                "properties": {
+                    "table": {"type": "string", "description": "Table name"},
+                    "schema": {
+                        "type": "string",
+                        "description": "Schema name (optional)",
+                    },
+                },
+                "required": ["table"],
+            },
+        )
+
+    def _create_analyze_column_tool(self) -> Tool:
+        """Create analyze_column tool."""
+        return Tool(
+            name="analyze_column",
+            description="Get comprehensive column statistics including percentiles and distributions",
+            inputSchema={
+                "type": "object",
+                "properties": {
+                    "table": {"type": "string", "description": "Table name"},
+                    "column": {"type": "string", "description": "Column name"},
+                    "schema": {
+                        "type": "string",
+                        "description": "Schema name (optional)",
+                    },
+                },
+                "required": ["table", "column"],
+            },
+        )
+
+    def _create_explain_query_tool(self) -> Tool:
+        """Create explain_query tool."""
+        return Tool(
+            name="explain_query",
+            description="Get query execution plan to analyze performance",
+            inputSchema={
+                "type": "object",
+                "properties": {
+                    "query": {"type": "string", "description": "SQL query to explain"},
+                    "analyze": {
+                        "type": "boolean",
+                        "description": "Whether to execute the query (EXPLAIN ANALYZE)",
+                        "default": False,
+                    },
+                },
+                "required": ["query"],
+            },
+        )
+
+    def _create_profile_database_tool(self) -> Tool:
+        """Create profile_database tool."""
+        return Tool(
+            name="profile_database",
+            description="Get database-wide profiling information (size, table counts, etc.)",
+            inputSchema={
+                "type": "object",
+                "properties": {},
+                "required": [],
+            },
+        )
+
+    # Tool handlers
+    async def handle_get_database_info(
+        self, arguments: dict[str, Any]
+    ) -> list[TextContent]:
+        """Handle get_database_info request."""
+        assert self.inspector is not None
+
+        version = await self.connection.get_version()
+
+        # Sanitize connection URL (remove password)
+        url_parts = self.config.url.split("@")
+        if len(url_parts) > 1:
+            sanitized_url = f"<credentials>@{url_parts[-1]}"
+        else:
+            sanitized_url = self.config.url
+
+        from db_connect_mcp.models.database import DatabaseInfo
+
+        db_info = DatabaseInfo(
+            name=self.config.url.split("/")[-1],  # Extract DB name from URL
+            dialect=self.config.dialect,
+            version=version,
+            size_bytes=None,
+            schema_count=None,
+            table_count=None,
+            capabilities=self.adapter.capabilities,
+            server_encoding=None,
+            collation=None,
+            connection_url=sanitized_url,
+            read_only=self.config.read_only,
+        )
+
+        return [
+            TextContent(type="text", text=json.dumps(db_info.model_dump(), indent=2))
+        ]
+
+    async def handle_list_schemas(self, arguments: dict[str, Any]) -> list[TextContent]:
+        """Handle list_schemas request."""
+        assert self.inspector is not None
+
+        schemas = await self.inspector.get_schemas()
+        schemas_data = [s.model_dump() for s in schemas]
+
+        return [TextContent(type="text", text=json.dumps(schemas_data, indent=2))]
+
+    async def handle_list_tables(self, arguments: dict[str, Any]) -> list[TextContent]:
+        """Handle list_tables request."""
+        assert self.inspector is not None
+
+        schema = arguments.get("schema")
+        include_views = arguments.get("include_views", True)
+
+        tables = await self.inspector.get_tables(schema, include_views)
+        tables_data = [t.model_dump() for t in tables]
+
+        return [TextContent(type="text", text=json.dumps(tables_data, indent=2))]
+
+    async def handle_describe_table(
+        self, arguments: dict[str, Any]
+    ) -> list[TextContent]:
+        """Handle describe_table request."""
+        assert self.inspector is not None
+
+        table = arguments["table"]
+        schema = arguments.get("schema")
+
+        table_info = await self.inspector.describe_table(table, schema)
+
+        return [
+            TextContent(type="text", text=json.dumps(table_info.model_dump(), indent=2))
+        ]
+
+    async def handle_execute_query(
+        self, arguments: dict[str, Any]
+    ) -> list[TextContent]:
+        """Handle execute_query request."""
+        assert self.executor is not None
+
+        query = arguments["query"]
+        limit = arguments.get("limit", 1000)
+
+        result = await self.executor.execute_query(query, limit=limit)
+
+        return [
+            TextContent(type="text", text=json.dumps(result.model_dump(), indent=2))
+        ]
+
+    async def handle_sample_data(self, arguments: dict[str, Any]) -> list[TextContent]:
+        """Handle sample_data request."""
+        assert self.executor is not None
+
+        table = arguments["table"]
+        schema = arguments.get("schema")
+        limit = arguments.get("limit", 100)
+
+        result = await self.executor.sample_data(table, schema, limit)
+
+        return [
+            TextContent(type="text", text=json.dumps(result.model_dump(), indent=2))
+        ]
+
+    async def handle_get_relationships(
+        self, arguments: dict[str, Any]
+    ) -> list[TextContent]:
+        """Handle get_table_relationships request."""
+        assert self.inspector is not None
+
+        table = arguments["table"]
+        schema = arguments.get("schema")
+
+        relationships = await self.inspector.get_relationships(table, schema)
+        relationships_data = [r.model_dump() for r in relationships]
+
+        return [TextContent(type="text", text=json.dumps(relationships_data, indent=2))]
+
+    async def handle_analyze_column(
+        self, arguments: dict[str, Any]
+    ) -> list[TextContent]:
+        """Handle analyze_column request."""
+        assert self.analyzer is not None
+
+        table = arguments["table"]
+        column = arguments["column"]
+        schema = arguments.get("schema")
+
+        stats = await self.analyzer.analyze_column(table, column, schema)
+
+        return [TextContent(type="text", text=json.dumps(stats.model_dump(), indent=2))]
+
+    async def handle_explain_query(
+        self, arguments: dict[str, Any]
+    ) -> list[TextContent]:
+        """Handle explain_query request."""
+        assert self.executor is not None
+
+        query = arguments["query"]
+        analyze = arguments.get("analyze", False)
+
+        plan = await self.executor.explain_query(query, analyze)
+
+        return [TextContent(type="text", text=json.dumps(plan.model_dump(), indent=2))]
+
+    async def handle_profile_database(
+        self, arguments: dict[str, Any]
+    ) -> list[TextContent]:
+        """Handle profile_database request."""
+        # This would be implemented with adapter-specific profiling queries
+        return [
+            TextContent(
+                type="text",
+                text=json.dumps(
+                    {"message": "Database profiling not yet implemented"}, indent=2
+                ),
+            )
+        ]
+
+    async def cleanup(self) -> None:
+        """Cleanup resources."""
+        await self.connection.dispose()
+        logger.info("Database MCP server cleaned up")
+
+
+async def main() -> None:
+    """Main entry point for the MCP server."""
+    # Get database URL from environment
+    database_url = os.getenv("DATABASE_URL")
+    if not database_url:
+        raise ValueError("DATABASE_URL environment variable must be set")
+
+    # Create configuration
+    config = DatabaseConfig(url=database_url)
+
+    # Create and initialize server
+    mcp_server = DatabaseMCPServer(config)
+
+    try:
+        await mcp_server.initialize()
+
+        # Register list_tools handler
+        @mcp_server.server.list_tools()
+        async def list_tools() -> list[Tool]:
+            """List available tools based on database capabilities."""
+            tools = [
+                mcp_server._create_get_database_info_tool(),
+                mcp_server._create_list_schemas_tool(),
+                mcp_server._create_list_tables_tool(),
+                mcp_server._create_describe_table_tool(),
+                mcp_server._create_execute_query_tool(),
+                mcp_server._create_sample_data_tool(),
+            ]
+
+            # Add conditional tools
+            if mcp_server.adapter.capabilities.foreign_keys:
+                tools.append(mcp_server._create_get_relationships_tool())
+
+            if mcp_server.adapter.capabilities.advanced_stats:
+                tools.append(mcp_server._create_analyze_column_tool())
+
+            if mcp_server.adapter.capabilities.explain_plans:
+                tools.append(mcp_server._create_explain_query_tool())
+
+            if mcp_server.adapter.capabilities.profiling:
+                tools.append(mcp_server._create_profile_database_tool())
+
+            return tools
+
+        # Register tool call handlers
+        @mcp_server.server.call_tool()
+        async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]:
+            """Handle tool calls."""
+            handlers = {
+                "get_database_info": mcp_server.handle_get_database_info,
+                "list_schemas": mcp_server.handle_list_schemas,
+                "list_tables": mcp_server.handle_list_tables,
+                "describe_table": mcp_server.handle_describe_table,
+                "execute_query": mcp_server.handle_execute_query,
+                "sample_data": mcp_server.handle_sample_data,
+                "get_table_relationships": mcp_server.handle_get_relationships,
+                "analyze_column": mcp_server.handle_analyze_column,
+                "explain_query": mcp_server.handle_explain_query,
+                "profile_database": mcp_server.handle_profile_database,
+            }
+
+            handler = handlers.get(name)
+            if handler is None:
+                raise ValueError(f"Unknown tool: {name}")
+
+            return await handler(arguments)
+
+        # Run the server
+        from mcp.server.stdio import stdio_server
+
+        async with stdio_server() as (read_stream, write_stream):
+            await mcp_server.server.run(
+                read_stream,
+                write_stream,
+                mcp_server.server.create_initialization_options(),
+            )
+
+    finally:
+        await mcp_server.cleanup()
+
+
+if __name__ == "__main__":
+    # Windows-specific event loop policy
+    if os.name == "nt":
+        asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())  # type: ignore[attr-defined]
+
+    asyncio.run(main())