Build a read-only database MCP server with scoped roles

"""Create and remove a throwaway PostgreSQL 17 cluster for this lesson.

Usage:
  python lesson_cluster.py create DATA_DIR [--port 54329]   # prints the reader DSN
  python lesson_cluster.py destroy DATA_DIR

The cluster lives entirely in DATA_DIR, listens on localhost only, has no Unix
socket and uses scram-sha-256 passwords. It never touches another cluster or a
Homebrew service. Set PG_BIN if the PostgreSQL 17 binaries are not in
/opt/homebrew/opt/postgresql@17/bin.
"""
import argparse
import os
import re
import secrets
import shutil
import socket
import subprocess
from pathlib import Path

import psycopg
from psycopg import sql

HERE = Path(__file__).resolve().parent
PG_BIN = Path(os.environ.get("PG_BIN", "/opt/homebrew/opt/postgresql@17/bin"))
DATABASE = "docs_lesson"
ADMIN, OWNER, READER = "lesson_admin", "docs_owner", "mcp_reader"
ENV = {**os.environ, "LC_ALL": "C"}  # the macOS postmaster refuses to start without a valid locale


def dsn(user, password, port, database=DATABASE):
    return f"host=localhost port={port} dbname={database} user={user} password={password}"


def free_port():
    with socket.socket() as probe:
        probe.bind(("127.0.0.1", 0))
        return probe.getsockname()[1]


def server_version():
    output = subprocess.run([PG_BIN / "postgres", "--version"], capture_output=True, text=True, check=True).stdout
    match = re.search(r"\b(\d+)\.(\d+)", output)
    if not match or match[1] != "17":
        raise SystemExit(f"This lesson needs PostgreSQL 17; found: {output.strip()}")
    return f"{match[1]}.{match[2]}"


class LessonCluster:
    """A cluster in data_dir. create() starts and seeds it; destroy() stops it and deletes data_dir."""

    def __init__(self, data_dir, port=None):
        if not str(data_dir).strip():
            raise SystemExit("give the cluster a directory; an empty path is refused")
        self.data_dir = Path(data_dir).expanduser().resolve()
        if self.data_dir in (Path.cwd(), Path.home(), Path(self.data_dir.anchor)):
            raise SystemExit(f"refusing to use {self.data_dir} as a cluster directory")
        self.port = port or free_port()
        self.admin_password = secrets.token_urlsafe(18)
        self.reader_password = secrets.token_urlsafe(18)

    @property
    def reader_dsn(self):
        return dsn(READER, self.reader_password, self.port)

    def admin_dsn(self, database=DATABASE):
        return dsn(ADMIN, self.admin_password, self.port, database)

    def create(self):
        self.version = server_version()
        if self.data_dir.exists():
            raise SystemExit(f"{self.data_dir} already exists; destroy it first")
        try:
            self._create()
        except BaseException:
            self.destroy()  # never leave a running server behind a deleted data directory
            raise
        return self

    def _create(self):
        self.data_dir.mkdir(parents=True)
        password_file = self.data_dir / "admin-password"
        password_file.write_text(self.admin_password + "\n")
        password_file.chmod(0o600)
        subprocess.run([PG_BIN / "initdb", "-D", self.data_dir / "data", "-U", ADMIN, "--auth=scram-sha-256",
                        f"--pwfile={password_file}", "-E", "UTF8", "--locale=C"],
                       check=True, capture_output=True, env=ENV)
        options = f"-p {self.port} -c listen_addresses=localhost -c unix_socket_directories=''"
        subprocess.run([PG_BIN / "pg_ctl", "-D", self.data_dir / "data", "-o", options,
                        "-l", self.data_dir / "server.log", "-w", "start"], check=True, capture_output=True, env=ENV)
        self.seed()

    def seed(self):
        with psycopg.connect(self.admin_dsn("postgres"), autocommit=True) as admin:
            admin.execute(f"CREATE ROLE {OWNER} NOLOGIN")
            admin.execute(sql.SQL("CREATE ROLE {} LOGIN CONNECTION LIMIT 5 PASSWORD {}").format(
                sql.Identifier(READER), sql.Literal(self.reader_password)))
            admin.execute(f"CREATE DATABASE {DATABASE} OWNER {OWNER}")
        with psycopg.connect(self.admin_dsn(), autocommit=True) as admin:
            admin.execute(f"SET ROLE {OWNER}")
            admin.execute((HERE / "schema.sql").read_text())
            admin.execute("RESET ROLE")
            admin.execute((HERE / "grants.sql").read_text())

    def destroy(self):
        """Stop the server and delete data_dir, but only if data_dir looks like a lesson cluster."""
        if not self.data_dir.exists():
            return
        if any(self.data_dir.iterdir()) and not (self.data_dir / "admin-password").is_file():
            raise SystemExit(f"refusing to delete {self.data_dir}: it is not a lesson cluster (no admin-password file)")
        if (self.data_dir / "data" / "postmaster.pid").exists():
            subprocess.run([PG_BIN / "pg_ctl", "-D", self.data_dir / "data", "-m", "fast", "-w", "stop"],
                           check=True, capture_output=True, env=ENV)
        shutil.rmtree(self.data_dir)

    def __enter__(self):
        return self.create()

    def __exit__(self, *exc):
        self.destroy()


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("command", choices=["create", "destroy"])
    parser.add_argument("data_dir")
    parser.add_argument("--port", type=int, default=54329)
    options = parser.parse_args(argv)
    cluster = LessonCluster(options.data_dir, options.port)
    if options.command == "create":
        cluster.create()
        dsn_file = cluster.data_dir / "reader.dsn"
        dsn_file.write_text(cluster.reader_dsn + "\n")
        dsn_file.chmod(0o600)
        print(f"PostgreSQL {cluster.version} on localhost:{cluster.port}; reader DSN saved in {dsn_file}")
    else:
        cluster.destroy()
        print(f"stopped and removed {cluster.data_dir}")


if __name__ == "__main__":
    main()

-- Lesson schema and seed rows. lesson_cluster.py runs this as the owner role, docs_owner.
CREATE SCHEMA docs;

CREATE TABLE docs.editors (
    editor_id text PRIMARY KEY,
    email     text NOT NULL
);

CREATE TABLE docs.pages (
    slug      text PRIMARY KEY,
    editor_id text NOT NULL REFERENCES docs.editors,
    section   text NOT NULL CHECK (section IN ('guides', 'reference', 'changelog')),
    words     integer NOT NULL CHECK (words >= 0)
);

INSERT INTO docs.editors VALUES
    ('e-1', 'ana@example.com'),
    ('e-2', 'bo@example.com');

-- 30 deterministic pages: page-01 ... page-30.
INSERT INTO docs.pages
SELECT format('page-%s', lpad(n::text, 2, '0')),
       CASE WHEN n % 2 = 0 THEN 'e-1' ELSE 'e-2' END,
       (ARRAY['guides', 'reference', 'changelog'])[n % 3 + 1],
       200 + n * 25
FROM generate_series(1, 30) AS n;

-- Privileges for the MCP tool's login role. lesson_cluster.py runs this as the cluster administrator.
-- The reader may connect to this database, see schema docs and read docs.pages. It can still
-- read the system catalogs (pg_catalog), as every role can; it cannot read other tables or write.
REVOKE ALL ON DATABASE docs_lesson FROM PUBLIC;
REVOKE CONNECT ON DATABASE postgres FROM PUBLIC;
REVOKE CONNECT ON DATABASE template1 FROM PUBLIC;
GRANT CONNECT ON DATABASE docs_lesson TO mcp_reader;
GRANT USAGE ON SCHEMA docs TO mcp_reader;
GRANT SELECT ON docs.pages TO mcp_reader;

-- Session defaults for the reader. A session can change these with SET, so the tool also
-- sets them for every transaction; the GRANTs above are the boundary.
ALTER ROLE mcp_reader IN DATABASE docs_lesson SET default_transaction_read_only = on;
ALTER ROLE mcp_reader IN DATABASE docs_lesson SET statement_timeout = '500ms';

"""Run one model-written SQL query under read-only limits.

run_readonly_query() opens a new connection as the read-only role, starts a
READ ONLY transaction, sets a statement timeout for that transaction, and runs
the query through a named server-side cursor. PostgreSQL keeps the result on
the server; the tool fetches at most max_rows + 1 rows. The role's privileges
are the real boundary; these limits narrow what a query can do within them.
"""
import datetime
import decimal
import json

import psycopg

MAX_ROWS_LIMIT = 200
DEFAULT_MAX_ROWS = 50
TIMEOUT_MS = 500
MAX_RESULT_BYTES = 200_000  # JSON size of the returned rows

# SQLSTATE codes this tool expects, with the reason it reports.
REASONS = {
    "25006": "the transaction is read-only",
    "42501": "the role lacks the privilege",
    "57014": f"the query ran longer than {TIMEOUT_MS} ms",
    "42601": "only one SELECT or VALUES query is allowed",
    "0A000": "the query uses a feature a cursor does not allow, such as a data-modifying WITH",
    "42P01": "the table does not exist or is not visible",
    "42703": "a column does not exist",
}

WRITE_CHECK = """
SELECT r.rolsuper, r.rolcreaterole, r.rolcreatedb,
       pg_has_role(current_user, 'pg_write_all_data', 'USAGE'),
       EXISTS (SELECT 1 FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace
               WHERE n.nspname NOT IN ('pg_catalog', 'information_schema') AND n.nspname NOT LIKE 'pg_toast%'
                 AND c.relkind IN ('r', 'p', 'v', 'm', 'f')
                 AND (has_table_privilege(current_user, c.oid, 'INSERT, UPDATE, DELETE, TRUNCATE')
                      OR has_any_column_privilege(current_user, c.oid, 'INSERT, UPDATE'))),
       EXISTS (SELECT 1 FROM pg_namespace n
               WHERE n.nspname NOT IN ('pg_catalog', 'information_schema') AND n.nspname NOT LIKE 'pg_toast%'
                 AND has_schema_privilege(current_user, n.oid, 'CREATE')),
       has_database_privilege(current_user, current_database(), 'CREATE, TEMP'),
       EXISTS (SELECT 1 FROM pg_auth_members m WHERE m.member = r.oid)
FROM pg_roles r WHERE r.rolname = current_user
"""
WRITE_PROBLEMS = ["is a superuser", "can create roles", "can create databases", "is a member of pg_write_all_data",
                  "can change rows or columns of a table", "can create objects in a schema",
                  "can create schemas or temporary tables in the database",
                  "is a member of another role, which SET ROLE could switch to"]


class QueryRejected(Exception):
    """The query did not run, or PostgreSQL stopped it."""

    def __init__(self, reason, sqlstate=None):
        super().__init__(reason)
        self.reason, self.sqlstate = reason, sqlstate


def to_json_value(value):
    """Turn a database value into something json.dumps can write."""
    if isinstance(value, (datetime.date, datetime.datetime, decimal.Decimal)):
        return str(value)
    return value


def run_readonly_query(reader_dsn, sql, max_rows=DEFAULT_MAX_ROWS):
    if not isinstance(sql, str) or not sql.strip():
        raise QueryRejected("the query is empty")
    if not isinstance(max_rows, int) or isinstance(max_rows, bool) or not 1 <= max_rows <= MAX_ROWS_LIMIT:
        raise QueryRejected(f"max_rows must be an integer from 1 to {MAX_ROWS_LIMIT}")
    try:
        # A new connection per call: a SET or set_config() from an earlier query cannot carry over.
        with psycopg.connect(reader_dsn, connect_timeout=5) as conn:
            with conn.transaction():
                conn.execute("SET TRANSACTION READ ONLY")
                conn.execute("SELECT set_config('statement_timeout', %s, true)", (f"{TIMEOUT_MS}ms",))
                # A named cursor is a server-side DECLARE ... CURSOR FOR <query>: PostgreSQL accepts one
                # SELECT or VALUES query, and keeps the result on the server until it is fetched.
                with conn.cursor(name="run_query") as cursor:
                    cursor.execute(sql)
                    rows = cursor.fetchmany(max_rows + 1)
                    columns = [column.name for column in cursor.description]
    except psycopg.Error as error:
        state = error.sqlstate
        raise QueryRejected(REASONS.get(state, "PostgreSQL rejected the query"), state) from None
    result_rows, size = [], 0
    for row in rows[:max_rows]:
        values = [to_json_value(value) for value in row]
        size += len(json.dumps(values, default=str))
        if size > MAX_RESULT_BYTES:
            # The rows were already fetched, so this bounds what the model receives, not what PostgreSQL sent.
            raise QueryRejected(f"the result is larger than {MAX_RESULT_BYTES} bytes; select fewer or shorter values")
        result_rows.append(values)
    return {"columns": columns, "rows": result_rows, "truncated": len(rows) > max_rows}


def visible_columns(reader_dsn):
    """Columns the reader can see. information_schema shows only objects the role has privileges on."""
    with psycopg.connect(reader_dsn, connect_timeout=5) as conn:
        rows = conn.execute(
            "SELECT table_schema, table_name, column_name, data_type FROM information_schema.columns "
            "WHERE table_schema NOT IN ('pg_catalog', 'information_schema') "
            "ORDER BY table_schema, table_name, ordinal_position").fetchall()
    return [{"table": f"{schema}.{table}", "column": column, "type": kind} for schema, table, column, kind in rows]


def write_problems(reader_dsn):
    """Ways the connected role could change data or objects in this database; an empty list means none found."""
    with psycopg.connect(reader_dsn, connect_timeout=5) as conn:
        flags = conn.execute(WRITE_CHECK).fetchone()
    return [problem for problem, flag in zip(WRITE_PROBLEMS, flags) if flag]


def check_reader(reader_dsn):
    """Refuse to serve with a role that could write."""
    problems = write_problems(reader_dsn)
    if problems:
        raise SystemExit("refusing to serve: the database role " + "; ".join(problems))

"""An MCP server that lets a model run read-only SQL against the lesson database.

Usage: DOCS_READER_DSN="host=localhost port=54329 ..." python postgres_server.py
The DSN must name the read-only role created by lesson_cluster.py. At start-up the
server refuses a role with any write capability its check knows about: superuser,
role or database creation, pg_write_all_data, table or column write privileges, or
CREATE or TEMP privileges. Built with the official MCP
Python SDK (mcp 2.3.0) and psycopg 3.
"""
import json
import os
import sys
from typing import Annotated, Any, TypedDict

from pydantic import Field

from mcp.server import MCPServer
from mcp.server.mcpserver.exceptions import ToolError
from mcp.types import ToolAnnotations

from readonly_query import DEFAULT_MAX_ROWS, MAX_ROWS_LIMIT, QueryRejected, check_reader, run_readonly_query, visible_columns


class QueryResult(TypedDict):
    columns: list[str]
    rows: list[list[Any]]
    truncated: bool


def build_server(reader_dsn):
    server = MCPServer("docs-sql", version="1.0.0",
                       instructions="Read-only SQL over schema docs. Read postgres://docs/schema for the visible columns.")

    @server.tool(title="Run a read-only SQL query", annotations=ToolAnnotations(readOnlyHint=True, openWorldHint=False))
    def run_query(sql: Annotated[str, Field(min_length=1, max_length=10_000)],
                  max_rows: Annotated[int, Field(ge=1, le=MAX_ROWS_LIMIT, strict=True)] = DEFAULT_MAX_ROWS) -> QueryResult:
        """Run one SELECT or VALUES query on the docs database. Returns at most max_rows rows (200 at most) and 200,000 bytes; each query stops after 500 ms."""
        try:
            return run_readonly_query(reader_dsn, sql, max_rows)
        except QueryRejected as rejected:
            code = f" (SQLSTATE {rejected.sqlstate})" if rejected.sqlstate else ""
            raise ToolError(f"Query rejected: {rejected.reason}{code}.") from None

    @server.resource("postgres://docs/schema", title="Visible tables and columns", mime_type="application/json")
    def schema() -> str:
        """The tables and columns the read-only role can see."""
        return json.dumps(visible_columns(reader_dsn), indent=2)

    return server


if __name__ == "__main__":
    dsn = os.environ.get("DOCS_READER_DSN")
    if not dsn:
        sys.exit("set DOCS_READER_DSN to the read-only role's connection string")
    check_reader(dsn)
    build_server(dsn).run()

"""Create a throwaway cluster, run the lesson's queries and role checks, save the results, remove the cluster.

Usage: python record_results.py [--write query-results.json]
Needs PostgreSQL 17 (see lesson_cluster.py) and the pinned packages.
"""
import argparse
import asyncio
import json
import tempfile
from pathlib import Path

import psycopg
from mcp import Client

import postgres_server
from lesson_cluster import LessonCluster
from readonly_query import QueryRejected, run_readonly_query, visible_columns, write_problems

QUERIES = [
    ("first five pages", "SELECT slug, section, words FROM docs.pages ORDER BY slug", 5),
    ("pages per section", "SELECT section, count(*) AS pages FROM docs.pages GROUP BY section ORDER BY section", 50),
    ("one row of a million", "SELECT n FROM generate_series(1, 1000000) AS n", 1),
    ("insert a page", "INSERT INTO docs.pages VALUES ('page-99', 'e-1', 'guides', 100)", 50),
    ("delete inside WITH", "WITH gone AS (DELETE FROM docs.pages RETURNING slug) SELECT * FROM gone", 50),
    ("function that writes", "SELECT lo_create(0)", 50),
    ("read editor emails", "SELECT email FROM docs.editors", 50),
    ("slow query", "SELECT pg_sleep(2)", 50),
    ("two statements", "SELECT 1; SELECT 2", 50),
    ("turn off read-only for the session", "SELECT set_config('default_transaction_read_only', 'off', false)", 50),
    ("function that writes, next call", "SELECT lo_create(0)", 50),
    ("one 300,000-character value", "SELECT repeat('x', 300000) AS big", 1),
    ("send a notification", "SELECT pg_notify('docs', 'hello')", 50),
    ("too many rows requested", "SELECT slug FROM docs.pages", 1000),
]


def tool_results(dsn):
    results = {}
    for label, sql, max_rows in QUERIES:
        try:
            results[label] = {"sql": sql, "max_rows": max_rows, "result": run_readonly_query(dsn, sql, max_rows)}
        except QueryRejected as rejected:
            results[label] = {"sql": sql, "max_rows": max_rows, "rejected": rejected.reason, "sqlstate": rejected.sqlstate}
    return results


def role_checks(dsn):
    """Bypass the tool's limits on purpose: what does the role itself allow?"""
    checks = {}
    with psycopg.connect(dsn, autocommit=True) as conn:
        conn.execute("SET default_transaction_read_only = off")  # a session may change its own default
        checks["read-only default after SET"] = conn.execute("SHOW default_transaction_read_only").fetchone()[0]
        for label, sql in [("insert with read-only off", "INSERT INTO docs.pages VALUES ('page-99', 'e-1', 'guides', 100)"),
                           ("delete with read-only off", "DELETE FROM docs.pages"),
                           ("read editors with read-only off", "SELECT email FROM docs.editors"),
                           ("create a large object with read-only off", "SELECT lo_create(0)")]:
            try:
                conn.execute(sql)
                checks[label] = "allowed"
            except psycopg.Error as error:
                checks[label] = f"{error.sqlstate}: {error.diag.message_primary}"
    return checks


async def mcp_results(dsn):
    async with Client(postgres_server.build_server(dsn)) as client:
        tools = await client.list_tools()
        found = await client.call_tool("run_query", {"sql": "SELECT count(*) AS pages FROM docs.pages"})
        denied = await client.call_tool("run_query", {"sql": "SELECT lo_create(0)"})
        tool = tools.tools[0].model_dump(by_alias=True, exclude_none=True, mode="json")
        return {"tool": tool,
                "count": {"isError": found.is_error, "structuredContent": found.structured_content},
                "write": {"isError": denied.is_error, "text": denied.content[0].text}}


def record():
    with tempfile.TemporaryDirectory() as directory:
        with LessonCluster(Path(directory) / "cluster") as cluster:
            dsn = cluster.reader_dsn
            return {"postgres": cluster.version, "visible_columns": visible_columns(dsn), "write_problems": write_problems(dsn),
                    "queries": tool_results(dsn), "role_checks": role_checks(dsn), "mcp": asyncio.run(mcp_results(dsn))}


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--write", type=Path, help="save the results to this JSON file")
    options = parser.parse_args(argv)
    results = record()
    print(f"PostgreSQL {results['postgres']}")
    for label, outcome in results["queries"].items():
        if "rejected" in outcome:
            shown = f"rejected: {outcome['rejected']}" + (f" ({outcome['sqlstate']})" if outcome["sqlstate"] else "")
        else:
            shown = f"{len(outcome['result']['rows'])} rows, truncated={outcome['result']['truncated']}"
        print(f"{label:36} {shown}")
    for label, outcome in results["role_checks"].items():
        print(f"role: {label:30} {outcome}")
    if options.write:
        options.write.write_text(json.dumps(results, indent=2) + "\n", encoding="utf-8")
        print(f"wrote {options.write}")


if __name__ == "__main__":
    main()