fastmcp

1from __future__ import annotations2 3import os4import re5import sqlite36from typing import Any7 8from fastmcp import FastMCP9 10 11mcp = FastMCP("__SERVER_NAME__")12 13DATABASE_PATH = os.getenv("SQLITE_PATH", "./app.db")14MAX_ROWS = int(os.getenv("SQLITE_MAX_ROWS", "200"))15TABLE_NAME_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")16 17 18def _connect() -> sqlite3.Connection:19    return sqlite3.connect(f"file:{DATABASE_PATH}?mode=ro", uri=True)20 21 22def _reject_mutation(sql: str) -> None:23    normalized = sql.strip().lower()24    if not normalized.startswith("select"):25        raise ValueError("Only SELECT queries are allowed")26 27 28def _validate_table_name(table_name: str) -> str:29    if not TABLE_NAME_RE.fullmatch(table_name):30        raise ValueError("Invalid table name")31    return table_name32 33 34@mcp.tool35def list_tables() -> list[str]:36    """List user-defined SQLite tables."""37    with _connect() as conn:38        rows = conn.execute(39            "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%' ORDER BY name"40        ).fetchall()41    return [row[0] for row in rows]42 43 44@mcp.tool45def describe_table(table_name: str) -> list[dict[str, Any]]:46    """Describe columns for a SQLite table."""47    safe_table_name = _validate_table_name(table_name)48    with _connect() as conn:49        rows = conn.execute(f"PRAGMA table_info({safe_table_name})").fetchall()50    return [51        {52            "cid": row[0],53            "name": row[1],54            "type": row[2],55            "notnull": bool(row[3]),56            "default": row[4],57            "pk": bool(row[5]),58        }59        for row in rows60    ]61 62 63@mcp.tool64def query(sql: str, limit: int = 50) -> dict[str, Any]:65    """Run a read-only SELECT query and return rows plus column names."""66    _reject_mutation(sql)67    safe_limit = max(0, min(limit, MAX_ROWS))68    wrapped_sql = f"SELECT * FROM ({sql.strip().rstrip(';')}) LIMIT {safe_limit}"69    with _connect() as conn:70        cursor = conn.execute(wrapped_sql)71        columns = [column[0] for column in cursor.description or []]72        rows = [dict(zip(columns, row)) for row in cursor.fetchall()]73    return {"limit": safe_limit, "columns": columns, "rows": rows}74 75 76if __name__ == "__main__":77    mcp.run()78