Zero-Trust Enterprise Database Access via MCP

graph TD
    Agent[AI Agent: tools/call execute_sql] --> Server[Enterprise DB MCP Server]
    
    subgraph Perimeter 1: AST Validation Layer
        Server --> ASTParser[SQL AST Parser: node-sql-parser]
        ASTParser --> BlockCheck{Disallowed Clauses? Mutation / DDL / Multi-Stmt}
        BlockCheck -->|Violation Detected| Reject[-32004 Security Error]
    end
    
    subgraph Perimeter 2: Read-Only Transaction Pool
        BlockCheck -->|Passed| ReplicaPool[PostgreSQL Read Replica Pool]
        ReplicaPool --> SessionConfig[SET default_transaction_read_only = on]
        SessionConfig --> TimeoutEnforce[SET statement_timeout = 5000]
    end
    
    subgraph Perimeter 3: Result Sanitization & Capping
        TimeoutEnforce --> Exec[Execute Query]
        Exec --> CapRows[Enforce MAX_ROWS = 250]
        CapRows --> FormatMD[Format Dense Markdown Table]
        FormatMD --> Agent
    end

-- Vulnerability: Simple regex filters miss comment injection
SELECT * FROM users; /* DROP TABLE users; */

import { Parser, AST } from 'node-sql-parser';

const parser = new Parser();

export interface ValidationResult {
  isValid: boolean;
  error?: string;
}

export function validateSafeQuery(sql: string): ValidationResult {
  let astList: AST[] | AST;

  try {
    astList = parser.astify(sql, { database: 'postgresql' });
  } catch (err: any) {
    return { isValid: false, error: `SQL Syntax Error: ${err.message}` };
  }

  // Enforce single statement
  if (Array.isArray(astList) && astList.length > 1) {
    return { isValid: false, error: 'Multi-statement execution is strictly prohibited.' };
  }

  const ast = Array.isArray(astList) ? astList[0] : astList;

  // Enforce SELECT-only statement type
  if (ast.type !== 'select') {
    return {
      isValid: false,
      error: `Prohibited operation: '${ast.type.toUpperCase()}'. Only read-only SELECT queries are allowed.`
    };
  }

  // Inspect referenced tables to protect internal catalogs
  const disallowedTables = ['pg_shadow', 'pg_authid', 'pg_user'];
  const tableList = parser.tableList(sql, { database: 'postgresql' });

  for (const tableEntry of tableList) {
    const tableName = tableEntry.split('::')[1]?.toLowerCase();
    if (disallowedTables.includes(tableName)) {
      return {
        isValid: false,
        error: `Access to restricted system catalog '${tableName}' is denied.`
      };
    }
  }

  return { isValid: true };
}

-- 1. Create unprivileged role
CREATE ROLE mcp_analyst_ro WITH LOGIN PASSWORD 'secure_entropy_key_2026';

-- 2. Connect only to read replicas
GRANT CONNECT ON DATABASE production_analytics TO mcp_analyst_ro;
GRANT USAGE ON SCHEMA public TO mcp_analyst_ro;
GRANT SELECT ON ALL TABLES IN SCHEMA public TO mcp_analyst_ro;
ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT SELECT ON TABLES TO mcp_analyst_ro;

-- 3. Engine-level execution bounds
ALTER ROLE mcp_analyst_ro SET default_transaction_read_only = on;
ALTER ROLE mcp_analyst_ro SET statement_timeout = '5000'; -- Terminate any query taking > 5 seconds
ALTER ROLE mcp_analyst_ro SET work_mem = '64MB';         -- Cap RAM usage per query
ALTER ROLE mcp_analyst_ro SET idle_in_transaction_session_timeout = '2000';

-- Snowflake Virtual Warehouse with Auto-Suspend and Resource Monitor
CREATE ROLE mcp_snowflake_reader;
GRANT USAGE ON WAREHOUSE ANALYTICS_QUERY_WH TO ROLE mcp_snowflake_reader;
GRANT USAGE ON DATABASE PROD_DATA TO ROLE mcp_snowflake_reader;
GRANT USAGE ON ALL SCHEMAS IN DATABASE PROD_DATA TO ROLE mcp_snowflake_reader;
GRANT SELECT ON ALL TABLES IN DATABASE PROD_DATA TO ROLE mcp_snowflake_reader;

-- Restrict query execution time to 15 seconds
ALTER USER mcp_agent_user SET STATEMENT_TIMEOUT_IN_SECONDS = 15;

import { Pool } from 'pg';

export class CachedSchemaManager {
  private pool: Pool;
  private schemaCache: Map<string, { schemaData: any; expiresAt: number }> = new Map();
  private CACHE_TTL_MS = 3600 * 1000; // 1 Hour

  constructor(pool: Pool) {
    this.pool = pool;
  }

  async getTableSchema(tableName: string): Promise<any> {
    const cached = this.schemaCache.get(tableName);
    if (cached && Date.now() < cached.expiresAt) {
      return cached.schemaData;
    }

    const query = `
      SELECT column_name, data_type, is_nullable
      FROM information_schema.columns
      WHERE table_schema = 'public' AND table_name = $1
      ORDER BY ordinal_position;
    `;

    const res = await this.pool.query(query, [tableName]);
    const schemaData = res.rows;

    this.schemaCache.set(tableName, {
      schemaData,
      expiresAt: Date.now() + this.CACHE_TTL_MS
    });

    return schemaData;
  }
}

import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
import { CallToolRequestSchema, ListToolsRequestSchema } from '@modelcontextprotocol/sdk/types.js';
import { Pool } from 'pg';
import { validateSafeQuery } from './ast-validator.js';

const dbPool = new Pool({
  connectionString: process.env.READONLY_DATABASE_URL,
  max: 10,
  idleTimeoutMillis: 10000,
  connectionTimeoutMillis: 3000
});

const server = new Server(
  { name: 'enterprise-database-mcp', version: '2.0.0' },
  { capabilities: { tools: {} } }
);

// Register Available Tools
server.setRequestHandler(ListToolsRequestSchema, async () => {
  return {
    tools: [
      {
        name: 'execute_safe_query',
        description: 'Execute a read-only SQL SELECT query against the read replica. Maximum 200 rows returned.',
        inputSchema: {
          type: 'object',
          properties: {
            sql: { type: 'string', description: 'ANSI SQL SELECT query.' }
          },
          required: ['sql']
        }
      }
    ]
  };
});

// Handle Tool Execution with Full Verification
server.setRequestHandler(CallToolRequestSchema, async (request) => {
  if (request.params.name === 'execute_safe_query') {
    const sql = request.params.arguments?.sql as string;

    // 1. AST Validation
    const validation = validateSafeQuery(sql);
    if (!validation.isValid) {
      return {
        isError: true,
        content: [{ type: 'text', text: `Access Denied: ${validation.error}` }]
      };
    }

    // 2. Query Execution with Row Capping
    const client = await dbPool.connect();
    try {
      // Force read-only transaction state
      await client.query('BEGIN READ ONLY');
      const result = await client.query(`${sql} LIMIT 200`);
      await client.query('COMMIT');

      // 3. Serialize as Token-Efficient Markdown Table
      if (result.rows.length === 0) {
        return { content: [{ type: 'text', text: 'Query returned 0 records.' }] };
      }

      const headers = Object.keys(result.rows[0]);
      const headerRow = `| ${headers.join(' | ')} |`;
      const dividerRow = `| ${headers.map(() => '---').join(' | ')} |`;
      const dataRows = result.rows.map(r => `| ${headers.map(h => String(r[h] ?? 'NULL')).join(' | ')} |`).join('\n');

      return {
        content: [{ type: 'text', text: `${headerRow}\n${dividerRow}\n${dataRows}` }]
      };
    } catch (err: any) {
      await client.query('ROLLBACK');
      return {
        isError: true,
        content: [{ type: 'text', text: `Query Execution Error: ${err.message}` }]
      };
    } finally {
      client.release();
    }
  }

  throw new Error(`Unknown tool: ${request.params.name}`);
});

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

run();

{
  "jsonrpc": "2.0",
  "id": "query-509",
  "result": {
    "isError": true,
    "content": [
      {
        "type": "text",
        "text": "Access Denied: Prohibited operation: 'UPDATE'. Only read-only SELECT queries are allowed."
      }
    ]
  }
}