postgres/src/mcp_postgres/server.py - master - man4j/mcp-serv

postgres/src/mcp_postgres/server.py
526 строк16 KB
Your Name
fixes
17 май 2026, 22:01
17 май 2026, 22:015b0d93f
# Env:
#   PGHOST, PGUSER, PGPASSWORD
# Optional:
#   PGPORT=5432
#   PGSSLMODE=disable|require|...
#   PGAPPNAME=esp-bot
#   SQL_STATEMENT_TIMEOUT_MS=15000
#   SQL_MAX_ROWS=500
#   MCP_HOST=0.0.0.0
#   MCP_PORT=8010

import os
import re
import time
import json
import logging
from pathlib import Path
from typing import Annotated

import psycopg
from fastmcp import FastMCP
from mcp_common import build_instructions_block, cached, json_default, run_http_server, run_startup_checks
from psycopg.rows import dict_row
from pydantic import Field

mcp = FastMCP("postgres-sql")
BASE_DIR = Path(__file__).resolve().parent
logger = logging.getLogger("mcp_postgres")

SQL_STATEMENT_TIMEOUT_MS = int(os.environ.get("SQL_STATEMENT_TIMEOUT_MS", "15000"))
SQL_MAX_ROWS = int(os.environ.get("SQL_MAX_ROWS", "500"))

# Шаблон вида: {helpdesk:billing_orders} или {cashdesk:v_cashdesk_license}
DDL_TEMPLATE_RE = re.compile(r"\{([a-zA-Z0-9_]+):([a-zA-Z0-9_]+)\}")


def _instruction_parts() -> tuple[str, ...]:
    parts = ["public_instructions.txt"]

    if (BASE_DIR / "private_instructions.txt").exists():
        parts.append("private_instructions.txt")

    return tuple(parts)


def _connect(dbname: str):
    host = os.environ["PGHOST"]
    port = int(os.environ.get("PGPORT", "5432"))
    user = os.environ["PGUSER"]
    password = os.environ["PGPASSWORD"]
    sslmode = os.environ.get("PGSSLMODE", "disable")
    appname = os.environ.get("PGAPPNAME", "mcp-server-bot")

    return psycopg.connect(
        host=host,
        port=port,
        dbname=dbname,
        user=user,
        password=password,
        sslmode=sslmode,
        application_name=appname,
        connect_timeout=10,
        autocommit=True,
    )


def execute_remote_sql(sql: str, dbname: str) -> dict:
    """
    Только SELECT/CTE SELECT.
    Возвращает JSON-совместимый dict:
      {columns, rows, rowcount, truncated, _elapsed_ms}
      или {error}
    """
    if not isinstance(sql, str) or not sql.strip():
        return {"error": "SQL is empty"}

    started = time.time()
    try:
        with _connect(dbname) as conn:
            with conn.cursor(row_factory=dict_row) as cur:
                cur.execute(f"SET statement_timeout = {SQL_STATEMENT_TIMEOUT_MS}")
                cur.execute(sql.strip())

                if cur.description is None:
                    return {"error": "Query returned no result set (expected SELECT/CTE SELECT only)"}

                cols = [d.name for d in cur.description]
                rows = cur.fetchmany(SQL_MAX_ROWS)

                res = {
                    "columns": cols,
                    "rows": rows,
                    "rowcount": len(rows),
                    "truncated": len(rows) >= SQL_MAX_ROWS,
                    "_elapsed_ms": int((time.time() - started) * 1000),
                }

                json.dumps(res, default=json_default)
                return res

    except Exception as e:
        return {"error": f"{type(e).__name__}: {e}"}


def get_relation_kind(dbname: str, relation_name: str, schema_name: str = "public") -> str:
    """
    Возвращает relkind:
      r - ordinary table
      p - partitioned table
      v - view
      m - materialized view
    """
    sql = """
    SELECT c.relkind
    FROM pg_class c
    JOIN pg_namespace n ON n.oid = c.relnamespace
    WHERE n.nspname = %(schema_name)s
      AND c.relname = %(relation_name)s
      AND c.relkind IN ('r', 'p', 'v', 'm')
    """

    with _connect(dbname) as conn:
        with conn.cursor() as cur:
            cur.execute(sql, {"schema_name": schema_name, "relation_name": relation_name})
            row = cur.fetchone()

            if not row:
                raise RuntimeError(
                    f"Relation not found while loading DDL: {schema_name}.{relation_name} (dbname={dbname})"
                )

            return row[0]


def fetch_table_ddl(dbname: str, table_name: str, schema_name: str = "public") -> str:
    """
    Генерирует компактный DDL таблицы из PostgreSQL:
    - колонки
    - DEFAULT
    - NOT NULL
    - PK / UNIQUE / FK / CHECK
    - COMMENT ON TABLE / COMMENT ON COLUMN, если есть
    """
    sql_columns = """
    SELECT
        a.attnum,
        a.attname AS column_name,
        pg_catalog.format_type(a.atttypid, a.atttypmod) AS data_type,
        NOT a.attnotnull AS is_nullable,
        pg_get_expr(ad.adbin, ad.adrelid) AS column_default
    FROM pg_attribute a
    JOIN pg_class c ON c.oid = a.attrelid
    JOIN pg_namespace n ON n.oid = c.relnamespace
    LEFT JOIN pg_attrdef ad
        ON ad.adrelid = a.attrelid
       AND ad.adnum = a.attnum
    WHERE n.nspname = %(schema_name)s
      AND c.relname = %(table_name)s
      AND a.attnum > 0
      AND NOT a.attisdropped
    ORDER BY a.attnum
    """

    sql_constraints = """
    SELECT
        con.conname,
        con.contype,
        pg_get_constraintdef(con.oid, true) AS condef
    FROM pg_constraint con
    JOIN pg_class c ON c.oid = con.conrelid
    JOIN pg_namespace n ON n.oid = c.relnamespace
    WHERE n.nspname = %(schema_name)s
      AND c.relname = %(table_name)s
    ORDER BY
        CASE con.contype
            WHEN 'p' THEN 1
            WHEN 'u' THEN 2
            WHEN 'f' THEN 3
            WHEN 'c' THEN 4
            ELSE 99
        END,
        con.conname
    """

    sql_table_comment = """
    SELECT obj_description(c.oid, 'pg_class') AS table_comment
    FROM pg_class c
    JOIN pg_namespace n ON n.oid = c.relnamespace
    WHERE n.nspname = %(schema_name)s
      AND c.relname = %(table_name)s
    """

    sql_column_comments = """
    SELECT
        a.attname AS column_name,
        col_description(c.oid, a.attnum) AS column_comment
    FROM pg_attribute a
    JOIN pg_class c ON c.oid = a.attrelid
    JOIN pg_namespace n ON n.oid = c.relnamespace
    WHERE n.nspname = %(schema_name)s
      AND c.relname = %(table_name)s
      AND a.attnum > 0
      AND NOT a.attisdropped
    ORDER BY a.attnum
    """

    with _connect(dbname) as conn:
        with conn.cursor(row_factory=dict_row) as cur:
            cur.execute(sql_columns, {"schema_name": schema_name, "table_name": table_name})
            columns = cur.fetchall()

            if not columns:
                raise RuntimeError(
                    f"Table not found while loading DDL: {schema_name}.{table_name} (dbname={dbname})"
                )

            cur.execute(sql_constraints, {"schema_name": schema_name, "table_name": table_name})
            constraints = cur.fetchall()

            cur.execute(sql_table_comment, {"schema_name": schema_name, "table_name": table_name})
            row = cur.fetchone()
            table_comment = row["table_comment"] if row else None

            cur.execute(sql_column_comments, {"schema_name": schema_name, "table_name": table_name})
            column_comments_rows = cur.fetchall()
            column_comments = {
                r["column_name"]: r["column_comment"]
                for r in column_comments_rows
                if r["column_comment"]
            }

    parts: list[str] = []

    if table_comment:
        parts.append(f"-- {table_comment}")

    for col in columns:
        comment = column_comments.get(col["column_name"])
        if comment:
            parts.append(f"-- {col['column_name']} - {comment}")

    ddl_lines = [f"CREATE TABLE {schema_name}.{table_name} ("]
    item_lines: list[str] = []

    for col in columns:
        line = f"  {col['column_name']} {col['data_type']}"
        if col["column_default"] is not None:
            line += f" DEFAULT {col['column_default']}"
        if not col["is_nullable"]:
            line += " NOT NULL"
        item_lines.append(line)

    for con in constraints:
        item_lines.append(f"  CONSTRAINT {con['conname']} {con['condef']}")

    ddl_lines.append(",\n".join(item_lines))
    ddl_lines.append(");")

    parts.append("\n".join(ddl_lines))
    return "\n".join(parts)


def fetch_view_signature(dbname: str, view_name: str, schema_name: str = "public") -> str:
    """
    Возвращает компактную сигнатуру view/materialized view:
    - комментарий на объект, если есть
    - комментарии на колонки, если есть
    - список колонок с типами
    Без тяжелого AS SELECT ... тела.
    """
    sql_columns = """
    SELECT
        a.attnum,
        a.attname AS column_name,
        pg_catalog.format_type(a.atttypid, a.atttypmod) AS data_type
    FROM pg_attribute a
    JOIN pg_class c ON c.oid = a.attrelid
    JOIN pg_namespace n ON n.oid = c.relnamespace
    WHERE n.nspname = %(schema_name)s
      AND c.relname = %(view_name)s
      AND a.attnum > 0
      AND NOT a.attisdropped
    ORDER BY a.attnum
    """

    sql_view_comment = """
    SELECT obj_description(c.oid, 'pg_class') AS view_comment
    FROM pg_class c
    JOIN pg_namespace n ON n.oid = c.relnamespace
    WHERE n.nspname = %(schema_name)s
      AND c.relname = %(view_name)s
    """

    sql_column_comments = """
    SELECT
        a.attname AS column_name,
        col_description(c.oid, a.attnum) AS column_comment
    FROM pg_attribute a
    JOIN pg_class c ON c.oid = a.attrelid
    JOIN pg_namespace n ON n.oid = c.relnamespace
    WHERE n.nspname = %(schema_name)s
      AND c.relname = %(view_name)s
      AND a.attnum > 0
      AND NOT a.attisdropped
    ORDER BY a.attnum
    """

    with _connect(dbname) as conn:
        with conn.cursor(row_factory=dict_row) as cur:
            cur.execute(sql_columns, {"schema_name": schema_name, "view_name": view_name})
            columns = cur.fetchall()

            if not columns:
                raise RuntimeError(
                    f"View not found while loading signature: {schema_name}.{view_name} (dbname={dbname})"
                )

            cur.execute(sql_view_comment, {"schema_name": schema_name, "view_name": view_name})
            row = cur.fetchone()
            view_comment = row["view_comment"] if row else None

            cur.execute(sql_column_comments, {"schema_name": schema_name, "view_name": view_name})
            column_comments_rows = cur.fetchall()
            column_comments = {
                r["column_name"]: r["column_comment"]
                for r in column_comments_rows
                if r["column_comment"]
            }

    parts: list[str] = []

    if view_comment:
        parts.append(f"-- {view_comment}")

    for col in columns:
        comment = column_comments.get(col["column_name"])
        if comment:
            parts.append(f"-- {col['column_name']} - {comment}")

    ddl_lines = [f"CREATE VIEW {schema_name}.{view_name} ("]
    ddl_lines.append(",\n".join(
        f"  {col['column_name']} {col['data_type']}" for col in columns
    ))
    ddl_lines.append(");")

    parts.append("\n".join(ddl_lines))
    return "\n".join(parts)


def fetch_matview_signature(dbname: str, view_name: str, schema_name: str = "public") -> str:
    """
    Возвращает компактную сигнатуру materialized view.
    """
    sql_columns = """
    SELECT
        a.attnum,
        a.attname AS column_name,
        pg_catalog.format_type(a.atttypid, a.atttypmod) AS data_type
    FROM pg_attribute a
    JOIN pg_class c ON c.oid = a.attrelid
    JOIN pg_namespace n ON n.oid = c.relnamespace
    WHERE n.nspname = %(schema_name)s
      AND c.relname = %(view_name)s
      AND a.attnum > 0
      AND NOT a.attisdropped
    ORDER BY a.attnum
    """

    sql_view_comment = """
    SELECT obj_description(c.oid, 'pg_class') AS view_comment
    FROM pg_class c
    JOIN pg_namespace n ON n.oid = c.relnamespace
    WHERE n.nspname = %(schema_name)s
      AND c.relname = %(view_name)s
    """

    sql_column_comments = """
    SELECT
        a.attname AS column_name,
        col_description(c.oid, a.attnum) AS column_comment
    FROM pg_attribute a
    JOIN pg_class c ON c.oid = a.attrelid
    JOIN pg_namespace n ON n.oid = c.relnamespace
    WHERE n.nspname = %(schema_name)s
      AND c.relname = %(view_name)s
      AND a.attnum > 0
      AND NOT a.attisdropped
    ORDER BY a.attnum
    """

    with _connect(dbname) as conn:
        with conn.cursor(row_factory=dict_row) as cur:
            cur.execute(sql_columns, {"schema_name": schema_name, "view_name": view_name})
            columns = cur.fetchall()

            if not columns:
                raise RuntimeError(
                    f"Materialized view not found while loading signature: {schema_name}.{view_name} (dbname={dbname})"
                )

            cur.execute(sql_view_comment, {"schema_name": schema_name, "view_name": view_name})
            row = cur.fetchone()
            view_comment = row["view_comment"] if row else None

            cur.execute(sql_column_comments, {"schema_name": schema_name, "view_name": view_name})
            column_comments_rows = cur.fetchall()
            column_comments = {
                r["column_name"]: r["column_comment"]
                for r in column_comments_rows
                if r["column_comment"]
            }

    parts: list[str] = []

    if view_comment:
        parts.append(f"-- {view_comment}")

    for col in columns:
        comment = column_comments.get(col["column_name"])
        if comment:
            parts.append(f"-- {col['column_name']} - {comment}")

    ddl_lines = [f"CREATE MATERIALIZED VIEW {schema_name}.{view_name} ("]
    ddl_lines.append(",\n".join(
        f"  {col['column_name']} {col['data_type']}" for col in columns
    ))
    ddl_lines.append(");")

    parts.append("\n".join(ddl_lines))
    return "\n".join(parts)


def fetch_relation_prompt_repr(dbname: str, relation_name: str, schema_name: str = "public") -> str:
    relkind = get_relation_kind(dbname, relation_name, schema_name)

    if relkind in ("r", "p"):
        return fetch_table_ddl(dbname, relation_name, schema_name)

    if relkind == "v":
        return fetch_view_signature(dbname, relation_name, schema_name)

    if relkind == "m":
        return fetch_matview_signature(dbname, relation_name, schema_name)

    raise RuntimeError(
        f"Unsupported relation type for {schema_name}.{relation_name} (dbname={dbname}, relkind={relkind})"
    )


def render_prompt_template(template: str) -> str:
    missing: list[str] = []

    def repl(match: re.Match) -> str:
        dbname = match.group(1)
        relation_name = match.group(2)

        try:
            return fetch_relation_prompt_repr(dbname, relation_name)
        except Exception as e:
            missing.append(f"{dbname}.{relation_name}: {type(e).__name__}: {e}")
            return match.group(0)

    rendered = DDL_TEMPLATE_RE.sub(repl, template)

    if missing:
        logger.error(
            "Failed to render PostgreSQL prompt because some DDL placeholders are unavailable: %s",
            "; ".join(missing),
        )
        raise RuntimeError(
            "Failed to render prompt template because some DDL placeholders could not be resolved:\n"
            + "\n".join(f" - {item}" for item in missing)
        )

    unresolved = DDL_TEMPLATE_RE.findall(rendered)
    if unresolved:
        unresolved_text = ", ".join(f"{db}.{name}" for db, name in unresolved)
        raise RuntimeError(
            "Prompt template still contains unresolved DDL placeholders after rendering: "
            + unresolved_text
        )

    return rendered


@mcp.tool(
    name="postgres_query",
    description="Execute a read-only SQL SELECT query on PostgreSQL.",
)
def postgres_query(
    sql: Annotated[str, Field(description="SQL query")],
    dbname: Annotated[str, Field(description="Database name, for example: helpdesk, cashdesk")],
) -> dict:
    return execute_remote_sql(sql, dbname)


@cached
def _sql_agent_instructions() -> str:
    return build_instructions_block(
        "postgres_query",
        BASE_DIR,
        parts=_instruction_parts(),
        transform=render_prompt_template,
    )


@mcp.prompt(
    name="sql_agent_instructions",
    description="Instructions for using the PostgreSQL SQL tool correctly.",
)
def sql_agent_instructions() -> str:
    return _sql_agent_instructions()


def main() -> None:
    run_startup_checks(
        "mcp_postgres",
        [
            ("sql_agent_instructions", _sql_agent_instructions),
        ],
    )
    run_http_server(mcp, default_port=8010)


if __name__ == "__main__":
    main()