Model Context Protocol (MCP)

pip install mcp

# script_id: day_095_model_context_protocol/db_mcp_server
# db_mcp_server.py
import asyncio
import os
import sqlite3
from typing import Any
import mcp.server.stdio
import mcp.types as types
from mcp.server import NotificationOptions, Server
from mcp.server.models import InitializationOptions
from pydantic import AnyUrl

# Initialize the MCP server
server = Server("database-server")

# Your actual database connection
DB_PATH = os.environ.get("DB_PATH", "company.db")


def get_db():
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    return conn


@server.list_tools()
async def handle_list_tools() -> list[types.Tool]:
    """Declare what tools this server provides."""
    return [
        types.Tool(
            name="query_database",
            description=(
                "Execute a read-only SQL query against the company database. "
                "Use this to look up customer records, orders, products, etc. "
                "Only SELECT queries are allowed."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "sql": {
                        "type": "string",
                        "description": "The SELECT SQL query to execute",
                    },
                    "limit": {
                        "type": "integer",
                        "description": "Maximum rows to return (default 10, max 100)",
                        "default": 10,
                    },
                },
                "required": ["sql"],
            },
        ),
        types.Tool(
            name="list_tables",
            description="List all available tables in the database with their schemas.",
            inputSchema={
                "type": "object",
                "properties": {},
            },
        ),
    ]


@server.call_tool()
async def handle_call_tool(
    name: str, arguments: dict[str, Any] | None
) -> list[types.TextContent]:
    """Execute a tool call."""
    arguments = arguments or {}

    if name == "query_database":
        sql = arguments.get("sql", "").strip()
        limit = min(int(arguments.get("limit", 10)), 100)

        # Security: only allow SELECT
        if not sql.upper().startswith("SELECT"):
            return [types.TextContent(
                type="text",
                text="Error: Only SELECT queries are allowed."
            )]

        try:
            conn = get_db()
            # `limit` is safe to interpolate because we cast it to int and cap it
            # above. `sql` itself is NOT parameterized here — we only gate it to
            # SELECT. For untrusted input, use a real SQL parser/allowlist rather
            # than a prefix check, and never f-string raw user SQL.
            sql = sql.rstrip(";")
            if " LIMIT " not in sql.upper():
                sql = f"{sql} LIMIT {limit}"
            cursor = conn.execute(sql)
            rows = cursor.fetchall()
            columns = [desc[0] for desc in cursor.description]
            conn.close()

            if not rows:
                return [types.TextContent(type="text", text="No results found.")]

            # Format as markdown table
            result = "| " + " | ".join(columns) + " |\n"
            result += "| " + " | ".join(["---"] * len(columns)) + " |\n"
            for row in rows:
                result += "| " + " | ".join(str(v) for v in row) + " |\n"

            result += f"\n*{len(rows)} rows returned*"
            return [types.TextContent(type="text", text=result)]

        except Exception as e:
            return [types.TextContent(type="text", text=f"Database error: {e}")]

    elif name == "list_tables":
        try:
            conn = get_db()
            cursor = conn.execute(
                "SELECT name FROM sqlite_master WHERE type='table' ORDER BY name"
            )
            tables = [row[0] for row in cursor.fetchall()]

            result = "Available tables:\n"
            for table in tables:
                result += f"- {table}\n"

            conn.close()
            return [types.TextContent(type="text", text=result)]

        except Exception as e:
            return [types.TextContent(type="text", text=f"Error: {e}")]

    else:
        return [types.TextContent(type="text", text=f"Unknown tool: {name}")]


async def main():
    async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
        await server.run(
            read_stream,
            write_stream,
            # Boilerplate: copy as-is
            InitializationOptions(
                server_name="database-server",
                server_version="0.1.0",
                capabilities=server.get_capabilities(
                    notification_options=NotificationOptions(),
                    experimental_capabilities={},
                ),
            ),
        )


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

# script_id: day_095_model_context_protocol/db_mcp_server
@server.list_resources()
async def handle_list_resources() -> list[types.Resource]:
    """Expose readable resources."""
    return [
        types.Resource(
            uri="database://schema",
            name="Database Schema",
            description="Full schema of all database tables",
            mimeType="text/plain",
        ),
        types.Resource(
            uri="database://stats",
            name="Database Statistics",
            description="Row counts and last-updated timestamps for all tables",
            mimeType="text/plain",
        ),
    ]


@server.read_resource()
async def handle_read_resource(uri: AnyUrl) -> str:
    """Return the content of a resource by URI."""
    # Returning a plain str works today; a future mcp release will want
    # Iterable[ReadResourceContents].
    if str(uri) == "database://schema":
        conn = get_db()
        cursor = conn.execute(
            "SELECT sql FROM sqlite_master WHERE type='table' ORDER BY name"
        )
        schema = "\n\n".join(row[0] for row in cursor.fetchall() if row[0])
        conn.close()
        return schema

    elif str(uri) == "database://stats":
        conn = get_db()
        cursor = conn.execute(
            "SELECT name FROM sqlite_master WHERE type='table'"
        )
        tables = [row[0] for row in cursor.fetchall()]

        stats = []
        for table in tables:
            count = conn.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0]
            stats.append(f"{table}: {count} rows")

        conn.close()
        return "\n".join(stats)

    raise ValueError(f"Unknown resource: {uri}")

{
  "mcpServers": {
    "database": {
      "command": "python",
      "args": ["/path/to/your/db_mcp_server.py"],
      "env": {
        "DB_PATH": "/path/to/company.db"
      }
    }
  }
}

# script_id: day_095_model_context_protocol/filesystem_mcp_server
# filesystem_mcp_server.py
import os
from pathlib import Path
import mcp.server.stdio
import mcp.types as types
from mcp.server import Server, NotificationOptions
from mcp.server.models import InitializationOptions
import asyncio

server = Server("filesystem-server")

# Restrict access to a safe directory
ALLOWED_ROOT = Path(os.environ.get("ALLOWED_ROOT", "/tmp/ai_workspace")).resolve()


def safe_path(requested: str) -> Path:
    """Resolve path and verify it's within ALLOWED_ROOT."""
    resolved = (ALLOWED_ROOT / requested).resolve()
    if not resolved.is_relative_to(ALLOWED_ROOT):
        raise ValueError(f"Path traversal attempt: {requested}")
    return resolved


@server.list_tools()
async def handle_list_tools() -> list[types.Tool]:
    return [
        types.Tool(
            name="read_file",
            description="Read the contents of a file in the workspace.",
            inputSchema={
                "type": "object",
                "properties": {
                    "path": {"type": "string", "description": "Relative path to the file"},
                },
                "required": ["path"],
            },
        ),
        types.Tool(
            name="list_directory",
            description="List files and directories at a given path.",
            inputSchema={
                "type": "object",
                "properties": {
                    "path": {"type": "string", "description": "Relative path (default: root)", "default": "."},
                },
            },
        ),
        types.Tool(
            name="write_file",
            description="Write content to a file in the workspace.",
            inputSchema={
                "type": "object",
                "properties": {
                    "path": {"type": "string"},
                    "content": {"type": "string"},
                },
                "required": ["path", "content"],
            },
        ),
    ]


@server.call_tool()
async def handle_call_tool(name: str, arguments: dict | None) -> list[types.TextContent]:
    arguments = arguments or {}

    if name == "read_file":
        try:
            path = safe_path(arguments["path"])
            content = path.read_text(encoding="utf-8")
            return [types.TextContent(type="text", text=content)]
        except Exception as e:
            return [types.TextContent(type="text", text=f"Error: {e}")]

    elif name == "list_directory":
        try:
            path = safe_path(arguments.get("path", "."))
            entries = []
            for item in sorted(path.iterdir()):
                prefix = "📁 " if item.is_dir() else "📄 "
                entries.append(f"{prefix}{item.name}")
            return [types.TextContent(type="text", text="\n".join(entries) or "Empty directory")]
        except Exception as e:
            return [types.TextContent(type="text", text=f"Error: {e}")]

    elif name == "write_file":
        try:
            path = safe_path(arguments["path"])
            path.parent.mkdir(parents=True, exist_ok=True)
            path.write_text(arguments["content"], encoding="utf-8")
            return [types.TextContent(type="text", text=f"Written: {path}")]
        except Exception as e:
            return [types.TextContent(type="text", text=f"Error: {e}")]

    return [types.TextContent(type="text", text=f"Unknown tool: {name}")]


async def main():
    async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
        await server.run(
            read_stream, write_stream,
            InitializationOptions(
                server_name="filesystem-server",
                server_version="0.1.0",
                capabilities=server.get_capabilities(
                    notification_options=NotificationOptions(),
                    experimental_capabilities={},
                ),
            ),
        )

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

# Register a Claude Desktop MCP server (config sketch)
# ~/Library/Application Support/Claude/claude_desktop_config.json
{
  "mcpServers": {
    "db": { "command": "python", "args": ["db_mcp_server.py"] }
  }
}