Understanding MCP Servers: The Future of AI Tool Integration

from mcp.server import Server, NotificationOptions
from mcp.server.models import InitializationOptions
from mcp.types import Resource, Tool, TextContent
import asyncpg
import json

class DatabaseMCPServer:
    def __init__(self):
        self.server = Server("postgresql-mcp")
        self.db_pool = None
        self.setup_handlers()
    
    def setup_handlers(self):
        @self.server.list_resources()
        async def handle_list_resources() -> list[Resource]:
            """List available database tables as resources"""
            tables = await self.get_tables()
            return [
                Resource(
                    uri=f"postgresql://table/{table}",
                    name=f"Table: {table}",
                    description=f"PostgreSQL table {table}",
                    mimeType="application/json"
                )
                for table in tables
            ]
        
        @self.server.read_resource()
        async def handle_read_resource(uri: str) -> str:
            """Read data from a specific table"""
            table_name = uri.split('/')[-1]
            data = await self.query_table(table_name)
            return json.dumps(data, indent=2)
        
        @self.server.list_tools()
        async def handle_list_tools() -> list[Tool]:
            """List available database operations"""
            return [
                Tool(
                    name="execute_query",
                    description="Execute a SQL query on the database",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "query": {
                                "type": "string",
                                "description": "SQL query to execute"
                            }
                        },
                        "required": ["query"]
                    }
                )
            ]
        
        @self.server.call_tool()
        async def handle_call_tool(name: str, arguments: dict) -> list[TextContent]:
            """Execute database tools"""
            if name == "execute_query":
                query = arguments["query"]
                result = await self.execute_query(query)
                return [
                    TextContent(
                        type="text",
                        text=f"Query executed successfully:\n{json.dumps(result, indent=2)}"
                    )
                ]
    
    async def get_tables(self):
        """Get list of tables in the database"""
        async with self.db_pool.acquire() as conn:
            rows = await conn.fetch("""
                SELECT table_name 
                FROM information_schema.tables 
                WHERE table_schema = 'public'
            """)
            return [row['table_name'] for row in rows]
    
    async def execute_query(self, query: str):
        """Execute a SQL query safely"""
        async with self.db_pool.acquire() as conn:
            try:
                rows = await conn.fetch(query)
                return [dict(row) for row in rows]
            except Exception as e:
                return {"error": str(e)}

async def main():
    server = DatabaseMCPServer()
    
    # Initialize database connection
    server.db_pool = await asyncpg.create_pool(
        "postgresql://user:password@localhost/mydb"
    )
    
    # Run the MCP server
    async with server.server.run_stdio() as (read_stream, write_stream):
        await server.server.run(
            read_stream, write_stream,
            InitializationOptions(
                server_name="postgresql-mcp",
                server_version="1.0.0",
                capabilities=server.server.get_capabilities(
                    notification_options=NotificationOptions(),
                    experimental_capabilities={}
                )
            )
        )

if __name__ == "__main__":
    import asyncio
    asyncio.run(main())