09. Database Query Server

Chapter 9 of 22 · 20 min

KEY INSIGHT

Database MCP servers must balance query flexibility with security, using parameterized queries and strict schema definitions to prevent injection attacks. Connecting AI applications to databases through MCP requires careful design. The goal is enabling natural language queries while maintaining security and performance. **Connection Management:** ```python import asyncpg from contextlib import asynccontextmanager class DatabaseServer: def __init__(self, connection_string: str): self.connection_string = connection_string self._pool: asyncpg.Pool | None = None async def connect(self): self._pool = await asyncpg.create_pool( self.connection_string, min_size=2, max_size=10 ) async def disconnect(self): if self._pool: await self._pool.close() @asynccontextmanager async def transaction(self): async with self._pool.acquire() as conn: async with conn.transaction(): yield conn ``` **Schema Introspection:** Servers should expose table structure as resources. ```python @server.list_resources() async def list_resources() -> list[Resource]: return [ Resource( uri="db://schema", name="database_schema", description="Database schema overview", mimeType="application/json" ) ] @server.read_resource() async def read_resource(uri: str) -> str: if uri == "db://schema": async with self._pool.acquire() as conn: tables = await conn.fetch(""" SELECT table_name, column_name, data_type FROM information_schema.columns WHERE table_schema = 'public' ORDER BY table_name, ordinal_position """) return json.dumps([dict(row) for row in tables], indent=2) ``` **Query Tool Implementation:** ```python @server.call_tool() async def call_tool(name: str, args: dict) -> str: match name: case "execute_query": return await self._execute_select(args["sql"]) case "get_table_data": return await self._get_table( args["table"], args.get("limit", 100), args.get("offset", 0) ) case "get_table_schema": return await self._get_schema(args["table"]) async def _execute_select(self, sql: str) -> str: # Validate SQL is SELECT only normalized = sql.strip().upper() if not normalized.startswith("SELECT"): raise ToolError("Only SELECT queries allowed") async with self._pool.acquire() as conn: rows = await conn.fetch(sql) return json.dumps([dict(row) for row in rows], indent=2) async def _get_table(self, table: str, limit: int, offset: int) -> str: # Parameterized query prevents injection async with self._pool.acquire() as conn: rows = await conn.fetch( f'SELECT * FROM "{table}" LIMIT $1 OFFSET $2', limit, offset ) return json.dumps([dict(row) for row in rows], indent=2) ``` **Security Considerations:** - Whitelist allowed tables - Implement query timeouts - Limit result set sizes - Log all queries for audit - Validate identifier names against whitelist **Connection Failures:** ```python async def _safe_query(self, sql: str, *args): try: async with self._pool.acquire() as conn: return await conn.fetch(sql, *args) except asyncpg.PostgresConnectionError: raise ToolError("Database connection failed") except asyncpg.UndefinedTableError: raise ToolError(f"Table not found") except asyncpg.TooManyRowsError: raise ToolError("Query returned too many rows, add filters") ```

EXERCISE

Build a database server for a simple task management system with tables for users, tasks, and projects. Implement tools for querying task status, assigning tasks to users, and creating new projects with proper parameterized queries.