给 MariaDB 写一个生产级 MCP Server

mkdir mariadb-mcp && cd mariadb-mcp
pnpm init
pnpm add @modelcontextprotocol/sdk mysql2 zod node-sql-parser
pnpm add -D typescript tsx @types/node
echo '{"compilerOptions":{"target":"es2022","module":"esnext","moduleResolution":"bundler","strict":true,"esModuleInterop":true}}' > tsconfig.json

// src/db.ts
import { createPool, type Pool } from 'mysql2/promise';

export const readPool: Pool = createPool({
  uri: process.env.MARIADB_READ_URL!,          // 只读账号
  connectionLimit: 5,
  connectTimeout: 5_000,
  enableKeepAlive: true,
});

export const auditPool: Pool = createPool({
  uri: process.env.MARIADB_AUDIT_URL!,         // 写入审计库的专用账号
  connectionLimit: 2,
});

await auditPool.query(`
  CREATE TABLE IF NOT EXISTS mcp_audit (
    id BIGINT AUTO_INCREMENT PRIMARY KEY,
    ts DATETIME(6) DEFAULT CURRENT_TIMESTAMP(6),
    client VARCHAR(64),
    tool VARCHAR(64),
    sql_text TEXT,
    args_json JSON,
    rows_returned INT,
    error TEXT,
    INDEX idx_ts (ts),
    INDEX idx_client (client)
  ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='MCP 工具调用审计'
`);

// src/guard.ts
import { Parser } from 'node-sql-parser';
const parser = new Parser();

const DENY = new Set(['drop', 'truncate', 'alter', 'rename', 'grant', 'revoke', 'create_user']);
const WRITE = new Set(['insert', 'update', 'delete', 'replace']);

export interface GuardResult {
  ok: boolean;
  reason?: string;
  type: 'select' | 'write' | 'ddl' | 'unknown';
}

export function guard(sql: string, opts: { allowWrites: boolean; maxRowsAffected: number }): GuardResult {
  let ast;
  try {
    ast = parser.astify(sql, { database: 'MySQL' });
  } catch (e: any) {
    return { ok: false, reason: `parse error: ${e.message}`, type: 'unknown' };
  }
  const stmts = Array.isArray(ast) ? ast : [ast];
  if (stmts.length > 1) return { ok: false, reason: 'multi-statement not allowed', type: 'unknown' };
  const t: string = stmts[0].type;
  if (DENY.has(t)) return { ok: false, reason: `statement type "${t}" denied`, type: 'ddl' };

  if (WRITE.has(t)) {
    if (!opts.allowWrites) return { ok: false, reason: 'writes disabled', type: 'write' };
    if (!sql.toLowerCase().includes('where')) {
      return { ok: false, reason: 'write without WHERE clause is rejected', type: 'write' };
    }
    return { ok: true, type: 'write' };
  }
  return { ok: true, type: 'select' };
}

// src/server.ts
import { McpServer, ResourceTemplate } from '@modelcontextprotocol/sdk/server/mcp.js';
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
import { z } from 'zod';
import { readPool, auditPool } from './db.js';
import { guard } from './guard.js';

const ALLOW_WRITES = process.env.MCP_ALLOW_WRITES === '1';
const MAX_ROWS = Number(process.env.MCP_MAX_ROWS ?? 500);

const server = new McpServer({ name: 'mariadb-mcp', version: '1.0.0' });

// ---- tool: list_tables ----
server.tool('list_tables', 'List tables in current database', {}, async () => {
  const [rows] = await readPool.query<any[]>(
    `SELECT table_name AS name, table_rows, table_comment AS description
     FROM information_schema.tables
     WHERE table_schema = DATABASE()`
  );
  return { content: [{ type: 'text', text: JSON.stringify(rows, null, 2) }] };
});

// ---- tool: describe_table ----
server.tool(
  'describe_table',
  'Get a table\'s columns, indexes and comments',
  { table: z.string() },
  async ({ table }) => {
    const [cols] = await readPool.query<any[]>(
      `SELECT column_name, column_type, is_nullable, column_key, column_default, column_comment
       FROM information_schema.columns
       WHERE table_schema = DATABASE() AND table_name = ?`,
      [table],
    );
    const [idx] = await readPool.query<any[]>(`SHOW INDEXES FROM \`${table}\``);
    return {
      content: [
        { type: 'text', text: JSON.stringify({ columns: cols, indexes: idx }, null, 2) },
      ],
    };
  },
);

// ---- tool: run_query ----
server.tool(
  'run_query',
  'Run a single SQL statement (guarded)',
  { sql: z.string(), limit: z.number().int().positive().max(MAX_ROWS).default(50) },
  async ({ sql, limit }) => {
    const check = guard(sql, { allowWrites: ALLOW_WRITES, maxRowsAffected: MAX_ROWS });
    if (!check.ok) {
      await audit('run_query', sql, { limit }, 0, check.reason);
      return { content: [{ type: 'text', text: `BLOCKED: ${check.reason}` }], isError: true };
    }
    try {
      const finalSql = check.type === 'select' ? `${sql} LIMIT ${limit}` : sql;
      const [rows] = await readPool.query<any[]>(finalSql);
      const count = Array.isArray(rows) ? rows.length : (rows as any).affectedRows ?? 0;
      await audit('run_query', sql, { limit }, count);
      return { content: [{ type: 'text', text: JSON.stringify(rows, null, 2) }] };
    } catch (e: any) {
      await audit('run_query', sql, { limit }, 0, e.message);
      return { content: [{ type: 'text', text: `ERROR: ${e.message}` }], isError: true };
    }
  },
);

async function audit(
  tool: string,
  sql: string,
  args: object,
  rows: number,
  error?: string,
) {
  await auditPool.query(
    `INSERT INTO mcp_audit (client, tool, sql_text, args_json, rows_returned, error) VALUES (?, ?, ?, ?, ?, ?)`,
    [process.env.MCP_CLIENT ?? 'unknown', tool, sql, JSON.stringify(args), rows, error ?? null],
  );
}

await server.connect(new StdioServerTransport());

-- 只读账号（给 AI 用）
CREATE USER 'ai_readonly'@'%' IDENTIFIED BY 'strong-pass-1';
GRANT SELECT, SHOW VIEW ON app.* TO 'ai_readonly'@'%';

-- 审计账号
CREATE DATABASE IF NOT EXISTS audit;
CREATE USER 'ai_audit'@'%' IDENTIFIED BY 'strong-pass-2';
GRANT INSERT, SELECT ON audit.* TO 'ai_audit'@'%';

FLUSH PRIVILEGES;

# 列工具
echo '{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}' | node dist/server.js

# 跑一个查询
echo '{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"run_query","arguments":{"sql":"SELECT 1+1 AS r"}}}' | node dist/server.js

# 尝试 DROP（应该被拦截）
echo '{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"run_query","arguments":{"sql":"DROP TABLE users"}}}' | node dist/server.js
# → BLOCKED: statement type "drop" denied