Developing an MCP Server for Live Databases with Sandboxing

[AI Agent] -> ToolCall(execute_query)
     |
     v
[MCP Server Middleware] -> AST Parser: Проверка типов стейтментов (SELECT only)
     |
     v
[Transaction Sandbox] -> BEGIN -> SET TRANSACTION READ ONLY -> Выполнение -> JSON
     |
     v
(Завершение сессии / Ошибка / Disconnect) -> ROLLBACK -> DISCARD ALL -> Возврат в пул

import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
import { z } from "zod";
import { Pool, PoolClient } from "pg";
import { parse } from "pgsql-ast-parser";

const pool = new Pool({
  connectionString: process.env.DATABASE_URL,
  max: 10,
  idleTimeoutMillis: 30000,
});

const server = new McpServer({ name: "PostgresSandboxMCP", version: "1.0.0" });

// Хранилище активных транзакционных сессий агентов
const activeSessions = new Map<string, PoolClient>();

/**
 * Валидация SQL-запроса через AST-парсер.
 * Блокирует любые деструктивные DDL и мутирующие DML операции.
 */
function isQuerySafe(sql: string): boolean {
  try {
    const ast = parse(sql);
    // Проверяем, что все стейтменты в пакете представляют собой исключительно выборку данных
    return ast.length > 0 && ast.every((statement) => statement.type === "select");
  } catch {
    return false; // При ошибке парсинга запрос безусловно отклоняется
  }
}

server.tool(
  "execute_read_only_query",
  "Выполняет безопасный SELECT SQL-запрос к PostgreSQL базе данных с лимитом выборки.",
  { 
    sessionId: z.string().describe("Уникальный идентификатор диалоговой сессии агента"),
    sql: z.string().describe("SQL-запрос на чтение данных") 
  },
  async ({ sessionId, sql }) => {
    if (!isQuerySafe(sql)) {
      return { 
        content: [{ type: "text", text: "Security Error: Разрешены исключительно операторы SELECT." }], 
        isError: true 
      };
    }

    let client = activeSessions.get(sessionId);

    if (!client) {
      // Инициализация песочницы: выделение клиента из пула и перевод транзакции в Read-Only
      client = await pool.connect();
      await client.query("BEGIN");
      await client.query("SET TRANSACTION READ ONLY");
      activeSessions.set(sessionId, client);
    }

    try {
      // Принудительное ограничение размера выборки для защиты контекста модели
      const result = await client.query(sql);
      const rows = result.rows.slice(0, 100);

      return {
        content: [{ type: "text", text: JSON.stringify(rows, null, 2) }]
      };
    } catch (error) {
      // При синтаксической ошибке транзакция в Postgres переходит в статус 'aborted'.
      // Выполняем автоматическое восстановление сессии для продолжения диалога.
      await client.query("ROLLBACK");
      await client.query("BEGIN");
      await client.query("SET TRANSACTION READ ONLY");
      
      return { 
        content: [{ type: "text", text: `Postgres Execution Error: ${(error as Error).message}` }], 
        isError: true 
      };
    }
  }
);

// Очистка ресурсов при завершении процесса
process.on("SIGINT", async () => {
  for (const [, client] of activeSessions.entries()) {
    try {
      await client.query("ROLLBACK");
      client.release();
    } catch {
      // Подавление ошибок очистки при завершении
    }
  }
  await pool.end();
  process.exit(0);
});

async function main() {
  const transport = new StdioServerTransport();
  await server.connect(transport);
}

main().catch(console.error);