@kohlarnhin/mcp-server-postgres
Version:
MCP server for interacting with PostgreSQL databases based on Node
701 lines • 29 kB
JavaScript
import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
import { CallToolRequestSchema, ListResourcesRequestSchema, ListToolsRequestSchema, ReadResourceRequestSchema, } from "@modelcontextprotocol/sdk/types.js";
import pkg from 'pg';
const { Pool } = pkg;
import * as dotenv from "dotenv";
import SqlParser from 'node-sql-parser';
import { log } from './utils/index.js';
// @INFO: Load environment variables from .env file
dotenv.config();
log('info', 'Starting MCP server...');
// @INFO: Update the environment setup to ensure database is correctly set
if (process.env.NODE_ENV === 'test' && !process.env.POSTGRES_DB) {
process.env.POSTGRES_DB = 'mcp_test_db'; // @INFO: Ensure we have a database name for tests
}
// Write operation flags (global defaults)
const ALLOW_INSERT_OPERATION = process.env.ALLOW_INSERT_OPERATION === 'true';
const ALLOW_UPDATE_OPERATION = process.env.ALLOW_UPDATE_OPERATION === 'true';
const ALLOW_DELETE_OPERATION = process.env.ALLOW_DELETE_OPERATION === 'true';
const ALLOW_DDL_OPERATION = process.env.ALLOW_DDL_OPERATION === 'true';
// Schema-specific permissions
const SCHEMA_INSERT_PERMISSIONS = parseSchemaPermissions(process.env.SCHEMA_INSERT_PERMISSIONS);
const SCHEMA_UPDATE_PERMISSIONS = parseSchemaPermissions(process.env.SCHEMA_UPDATE_PERMISSIONS);
const SCHEMA_DELETE_PERMISSIONS = parseSchemaPermissions(process.env.SCHEMA_DELETE_PERMISSIONS);
const SCHEMA_DDL_PERMISSIONS = parseSchemaPermissions(process.env.SCHEMA_DDL_PERMISSIONS);
// Check if we're in multi-DB mode (no specific DB set)
const isMultiDbMode = !process.env.POSTGRES_DB || process.env.POSTGRES_DB.trim() === '';
// Force read-only mode in multi-DB mode unless explicitly configured otherwise
if (isMultiDbMode && process.env.MULTI_DB_WRITE_MODE !== 'true') {
log('error', 'Multi-DB mode detected - enabling read-only mode for safety');
}
// @INFO: Check if running in test mode
const isTestEnvironment = process.env.NODE_ENV === 'test' || process.env.VITEST;
// @INFO: Safe way to exit process (not during tests)
function safeExit(code) {
if (!isTestEnvironment) {
process.exit(code);
}
else {
log('error', `[Test mode] Would have called process.exit(${code})`);
}
}
// Function to parse schema-specific permissions from environment variables
function parseSchemaPermissions(permissionsString) {
const permissions = {};
if (!permissionsString) {
return permissions;
}
// Format: "schema1:true,schema2:false"
const permissionPairs = permissionsString.split(',');
for (const pair of permissionPairs) {
const [schema, value] = pair.split(':');
if (schema && value) {
permissions[schema.trim()] = value.trim() === 'true';
}
}
return permissions;
}
// Schema permission checking functions
function isInsertAllowedForSchema(schema) {
if (!schema) {
return ALLOW_INSERT_OPERATION;
}
return schema in SCHEMA_INSERT_PERMISSIONS
? SCHEMA_INSERT_PERMISSIONS[schema]
: ALLOW_INSERT_OPERATION;
}
function isUpdateAllowedForSchema(schema) {
if (!schema) {
return ALLOW_UPDATE_OPERATION;
}
return schema in SCHEMA_UPDATE_PERMISSIONS
? SCHEMA_UPDATE_PERMISSIONS[schema]
: ALLOW_UPDATE_OPERATION;
}
function isDeleteAllowedForSchema(schema) {
if (!schema) {
return ALLOW_DELETE_OPERATION;
}
return schema in SCHEMA_DELETE_PERMISSIONS
? SCHEMA_DELETE_PERMISSIONS[schema]
: ALLOW_DELETE_OPERATION;
}
function isDDLAllowedForSchema(schema) {
if (!schema) {
return ALLOW_DDL_OPERATION;
}
return schema in SCHEMA_DDL_PERMISSIONS
? SCHEMA_DDL_PERMISSIONS[schema]
: ALLOW_DDL_OPERATION;
}
// Extract schema from SQL query
function extractSchemaFromQuery(sql) {
// Default schema from environment
const defaultSchema = process.env.POSTGRES_DB || null;
// If we have a default schema and not in multi-DB mode, return it
if (defaultSchema && !isMultiDbMode) {
return defaultSchema;
}
// Try to extract schema from query
// Case 1: SET search_path statement
const useMatch = sql.match(/SET\s+search_path\s+TO\s+([a-zA-Z0-9_]+)/i);
if (useMatch && useMatch[1]) {
return useMatch[1];
}
// Case 2: database.table notation
const dbTableMatch = sql.match(/([a-zA-Z0-9_]+)\.([a-zA-Z0-9_]+)/i);
if (dbTableMatch && dbTableMatch[1]) {
return dbTableMatch[1];
}
// Return default if we couldn't find a schema in the query
return defaultSchema;
}
// Update tool description to include multi-DB mode and schema-specific permissions
let toolDescription = 'Run SQL queries against PostgreSQL database';
if (isMultiDbMode) {
toolDescription += ' (Multi-DB mode enabled)';
}
if (ALLOW_INSERT_OPERATION || ALLOW_UPDATE_OPERATION || ALLOW_DELETE_OPERATION || ALLOW_DDL_OPERATION) {
// At least one write operation is enabled
toolDescription += ' with support for:';
if (ALLOW_INSERT_OPERATION) {
toolDescription += ' INSERT,';
}
if (ALLOW_UPDATE_OPERATION) {
toolDescription += ' UPDATE,';
}
if (ALLOW_DELETE_OPERATION) {
toolDescription += ' DELETE,';
}
if (ALLOW_DDL_OPERATION) {
toolDescription += ' DDL,';
}
// Remove trailing comma and add READ operations
toolDescription = toolDescription.replace(/,$/, '') + ' and READ operations';
if (Object.keys(SCHEMA_INSERT_PERMISSIONS).length > 0 ||
Object.keys(SCHEMA_UPDATE_PERMISSIONS).length > 0 ||
Object.keys(SCHEMA_DELETE_PERMISSIONS).length > 0 ||
Object.keys(SCHEMA_DDL_PERMISSIONS).length > 0) {
toolDescription += ' (Schema-specific permissions enabled)';
}
}
else {
// Only read operations are allowed
toolDescription += ' (READ-ONLY)';
}
// Update PostgreSQL config to handle blank database name
const config = {
server: {
name: 'mcp-server-postgres',
version: '1.0.0',
connectionTypes: ['stdio'],
},
postgres: {
host: process.env.POSTGRES_HOST || '127.0.0.1',
port: Number(process.env.POSTGRES_PORT || '5432'),
user: process.env.POSTGRES_USER || 'postgres',
password: process.env.POSTGRES_PASS || 'postgres',
database: process.env.POSTGRES_DB || undefined, // Allow undefined database for multi-DB mode
max: 10, // Connection pool max size
...(process.env.POSTGRES_SSL === 'true'
? {
ssl: {
rejectUnauthorized: process.env.POSTGRES_SSL_REJECT_UNAUTHORIZED === 'true',
},
}
: {}),
},
paths: {
schema: 'schema',
},
};
// @INFO: Add debug logging for configuration
log('info', 'PostgreSQL Configuration:', JSON.stringify({
host: config.postgres.host,
port: config.postgres.port,
user: config.postgres.user,
password: config.postgres.password ? '******' : 'not set',
database: config.postgres.database || 'MULTI_DB_MODE',
ssl: process.env.POSTGRES_SSL === 'true' ? 'enabled' : 'disabled',
multiDbMode: isMultiDbMode ? 'enabled' : 'disabled',
}, null, 2));
// @INFO: Lazy load PostgreSQL pool
let poolPromise;
const getPool = () => {
if (!poolPromise) {
poolPromise = new Promise((resolve, reject) => {
try {
const pool = new Pool(config.postgres);
log('info', 'PostgreSQL pool created successfully');
resolve(pool);
}
catch (error) {
log('error', 'Error creating PostgreSQL pool:', error);
reject(error);
}
});
}
return poolPromise;
};
// @INFO: Lazy load server instance
let serverInstance = null;
const getServer = () => {
if (!serverInstance) {
serverInstance = new Promise((resolve) => {
const server = new Server(config.server, {
capabilities: {
resources: {},
tools: {
postgres_query: {
description: toolDescription,
inputSchema: {
type: 'object',
properties: {
sql: {
type: 'string',
description: 'The SQL query to execute',
},
},
required: ['sql'],
},
},
},
},
});
// @INFO: Register request handlers
server.setRequestHandler(ListResourcesRequestSchema, async () => {
try {
log('error', 'Handling ListResourcesRequest');
// If we're in multi-DB mode, list all databases first
if (isMultiDbMode) {
const databases = (await executeQuery('SELECT datname AS "Database" FROM pg_database WHERE datistemplate = false'));
// For each database, list tables
let allResources = [];
for (const db of databases) {
// Skip system databases
if (['postgres', 'template0', 'template1'].includes(db.Database)) {
continue;
}
// Switch connections to use this database
const dbPool = new Pool({
...config.postgres,
database: db.Database,
});
try {
const tables = await dbPool.query(`SELECT table_name FROM information_schema.tables
WHERE table_schema = 'public'`);
allResources.push(...tables.rows.map((row) => ({
uri: new URL(`${db.Database}/${row.table_name}/${config.paths.schema}`, `${config.postgres.host}:${config.postgres.port}`).href,
mimeType: 'application/json',
name: `"${db.Database}.${row.table_name}" database schema`,
})));
}
catch (err) {
log('error', `Error listing tables for database ${db.Database}:`, err);
}
finally {
await dbPool.end();
}
}
return {
resources: allResources,
};
}
else {
// Original behavior for single database mode
const results = (await executeQuery('SELECT table_name FROM information_schema.tables WHERE table_schema = \'public\''));
return {
resources: results.map((row) => ({
uri: new URL(`${row.table_name}/${config.paths.schema}`, `${config.postgres.host}:${config.postgres.port}`).href,
mimeType: 'application/json',
name: `"${row.table_name}" database schema`,
})),
};
}
}
catch (error) {
log('error', 'Error in ListResourcesRequest handler:', error);
throw error;
}
});
server.setRequestHandler(ReadResourceRequestSchema, async (request) => {
try {
log('error', 'Handling ReadResourceRequest');
const resourceUrl = new URL(request.params.uri);
const pathComponents = resourceUrl.pathname.split('/');
const schema = pathComponents.pop();
const tableName = pathComponents.pop();
let dbName = null;
// In multi-DB mode, we expect a database name in the path
if (isMultiDbMode && pathComponents.length > 0) {
dbName = pathComponents.pop() || null;
}
if (schema !== config.paths.schema) {
throw new Error('Invalid resource URI');
}
// If in multi-DB mode and a specific database was identified
if (isMultiDbMode && dbName) {
// Create a new connection to that specific database
const dbPool = new Pool({
...config.postgres,
database: dbName,
});
try {
const columnsResult = await dbPool.query(`SELECT column_name, data_type
FROM information_schema.columns
WHERE table_name = $1
AND table_schema = 'public'`, [tableName]);
return {
contents: [
{
uri: request.params.uri,
mimeType: 'application/json',
text: JSON.stringify(columnsResult.rows, null, 2),
},
],
};
}
finally {
await dbPool.end();
}
}
else {
// Regular mode with current database
const results = (await executeQuery(`SELECT column_name, data_type
FROM information_schema.columns
WHERE table_name = $1
AND table_schema = 'public'`, [tableName]));
return {
contents: [
{
uri: request.params.uri,
mimeType: 'application/json',
text: JSON.stringify(results, null, 2),
},
],
};
}
}
catch (error) {
log('error', 'Error in ReadResourceRequest handler:', error);
throw error;
}
});
server.setRequestHandler(ListToolsRequestSchema, async () => {
log('error', 'Handling ListToolsRequest');
const toolsResponse = {
tools: [
{
name: 'postgres_query',
description: toolDescription,
inputSchema: {
type: 'object',
properties: {
sql: {
type: 'string',
description: 'The SQL query to execute',
},
},
required: ['sql'],
},
},
],
};
log('error', 'ListToolsRequest response:', JSON.stringify(toolsResponse, null, 2));
return toolsResponse;
});
server.setRequestHandler(CallToolRequestSchema, async (request) => {
try {
log('error', 'Handling CallToolRequest:', request.params.name);
if (request.params.name !== 'postgres_query') {
throw new Error(`Unknown tool: ${request.params.name}`);
}
const sql = request.params.arguments?.sql;
return executeReadOnlyQuery(sql);
}
catch (error) {
log('error', 'Error in CallToolRequest handler:', error);
throw error;
}
});
resolve(server);
});
}
return serverInstance;
};
const { Parser } = SqlParser;
const parser = new Parser();
async function getQueryTypes(query) {
try {
log('error', "Parsing SQL query: ", query);
// Parse into AST or array of ASTs - specify the database type as PostgreSQL
const astOrArray = parser.astify(query, { database: 'postgresql' });
const statements = Array.isArray(astOrArray) ? astOrArray : [astOrArray];
log('error', "Parsed SQL AST: ", statements.map(stmt => stmt.type?.toLowerCase() ?? 'unknown'));
// Map each statement to its lowercased type (e.g., 'select', 'update', 'insert', 'delete', etc.)
return statements.map(stmt => stmt.type?.toLowerCase() ?? 'unknown');
}
catch (err) {
log('error', "sqlParser error, query: ", query);
log('error', 'Error parsing SQL query:', err);
throw new Error(`Parsing failed: ${err.message}`);
}
}
async function executeQuery(sql, params = []) {
let client = null;
try {
const pool = await getPool();
client = await pool.connect();
const result = await client.query(sql, params);
return result.rows;
}
catch (error) {
log('error', 'Error executing query:', error);
throw error;
}
finally {
if (client) {
client.release();
log('error', 'Client released');
}
}
}
async function executeReadOnlyQuery(sql) {
let client = null;
try {
// Check the type of query
const queryTypes = await getQueryTypes(sql);
// Get schema for permission checking
const schema = extractSchemaFromQuery(sql);
const isUpdateOperation = queryTypes.some(type => ['update'].includes(type));
const isInsertOperation = queryTypes.some(type => ['insert'].includes(type));
const isDeleteOperation = queryTypes.some(type => ['delete'].includes(type));
const isDDLOperation = queryTypes.some(type => ['create', 'alter', 'drop', 'truncate'].includes(type));
// Check schema-specific permissions
if (isInsertOperation && !isInsertAllowedForSchema(schema)) {
log('error', `INSERT operations are not allowed for schema '${schema || 'default'}'. Configure SCHEMA_INSERT_PERMISSIONS.`);
return {
content: [
{
type: 'text',
text: `Error: INSERT operations are not allowed for schema '${schema || 'default'}'. Ask the administrator to update SCHEMA_INSERT_PERMISSIONS.`,
},
],
isError: true,
};
}
if (isUpdateOperation && !isUpdateAllowedForSchema(schema)) {
log('error', `UPDATE operations are not allowed for schema '${schema || 'default'}'. Configure SCHEMA_UPDATE_PERMISSIONS.`);
return {
content: [
{
type: 'text',
text: `Error: UPDATE operations are not allowed for schema '${schema || 'default'}'. Ask the administrator to update SCHEMA_UPDATE_PERMISSIONS.`,
},
],
isError: true,
};
}
if (isDeleteOperation && !isDeleteAllowedForSchema(schema)) {
log('error', `DELETE operations are not allowed for schema '${schema || 'default'}'. Configure SCHEMA_DELETE_PERMISSIONS.`);
return {
content: [
{
type: 'text',
text: `Error: DELETE operations are not allowed for schema '${schema || 'default'}'. Ask the administrator to update SCHEMA_DELETE_PERMISSIONS.`,
},
],
isError: true,
};
}
if (isDDLOperation && !isDDLAllowedForSchema(schema)) {
log('error', `DDL operations are not allowed for schema '${schema || 'default'}'. Configure SCHEMA_DDL_PERMISSIONS.`);
return {
content: [
{
type: 'text',
text: `Error: DDL operations are not allowed for schema '${schema || 'default'}'. Ask the administrator to update SCHEMA_DDL_PERMISSIONS.`,
},
],
isError: true,
};
}
// For write operations that are allowed, use executeWriteQuery
if ((isInsertOperation && isInsertAllowedForSchema(schema)) ||
(isUpdateOperation && isUpdateAllowedForSchema(schema)) ||
(isDeleteOperation && isDeleteAllowedForSchema(schema)) ||
(isDDLOperation && isDDLAllowedForSchema(schema))) {
return executeWriteQuery(sql);
}
// For read-only operations, continue with the original logic
const pool = await getPool();
client = await pool.connect();
log('error', 'Read-only client acquired');
// Set read-only mode in PostgreSQL
await client.query('SET TRANSACTION READ ONLY');
// Begin transaction
await client.query('BEGIN');
try {
// Execute query
const result = await client.query(sql);
const rows = result.rows;
// Rollback transaction (since it's read-only)
await client.query('ROLLBACK');
// Reset to read-write mode
await client.query('SET TRANSACTION READ WRITE');
return {
content: [
{
type: 'text',
text: JSON.stringify(rows, null, 2),
},
],
isError: false,
};
}
catch (error) {
// Rollback transaction on query error
log('error', 'Error executing read-only query:', error);
await client.query('ROLLBACK');
throw error;
}
}
catch (error) {
// Ensure we rollback and reset transaction mode on any error
log('error', 'Error in read-only query transaction:', error);
try {
if (client) {
await client.query('ROLLBACK');
await client.query('SET TRANSACTION READ WRITE');
}
}
catch (cleanupError) {
// Ignore errors during cleanup
log('error', 'Error during cleanup:', cleanupError);
}
throw error;
}
finally {
if (client) {
client.release();
log('error', 'Read-only client released');
}
}
}
// @INFO: New function to handle write operations
async function executeWriteQuery(sql) {
let client = null;
try {
const pool = await getPool();
client = await pool.connect();
log('error', 'Write client acquired');
// Extract schema for permissions (if needed)
const schema = extractSchemaFromQuery(sql);
// @INFO: Begin transaction for write operation
await client.query('BEGIN');
try {
// @INFO: Execute the write query
const result = await client.query(sql);
// @INFO: Commit the transaction
await client.query('COMMIT');
// @INFO: Format the response based on operation type
let responseText;
// Check the type of query
const queryTypes = await getQueryTypes(sql);
const isUpdateOperation = queryTypes.some(type => ['update'].includes(type));
const isInsertOperation = queryTypes.some(type => ['insert'].includes(type));
const isDeleteOperation = queryTypes.some(type => ['delete'].includes(type));
const isDDLOperation = queryTypes.some(type => ['create', 'alter', 'drop', 'truncate'].includes(type));
// @INFO: Format response based on operation type
if (isInsertOperation) {
responseText = `Insert successful on schema '${schema || 'default'}'. Affected rows: ${result.rowCount}`;
}
else if (isUpdateOperation) {
responseText = `Update successful on schema '${schema || 'default'}'. Affected rows: ${result.rowCount}`;
}
else if (isDeleteOperation) {
responseText = `Delete successful on schema '${schema || 'default'}'. Affected rows: ${result.rowCount}`;
}
else if (isDDLOperation) {
responseText = `DDL operation successful on schema '${schema || 'default'}'.`;
}
else {
responseText = JSON.stringify(result.rows, null, 2);
}
return {
content: [
{
type: 'text',
text: responseText,
},
],
isError: false,
};
}
catch (error) {
// @INFO: Rollback on error
log('error', 'Error executing write query:', error);
await client.query('ROLLBACK');
return {
content: [
{
type: 'text',
text: `Error executing write operation: ${error instanceof Error ? error.message : String(error)}`,
},
],
isError: true,
};
}
}
catch (error) {
log('error', 'Error in write operation transaction:', error);
return {
content: [
{
type: 'text',
text: `Database connection error: ${error instanceof Error ? error.message : String(error)}`,
},
],
isError: true,
};
}
finally {
if (client) {
client.release();
log('error', 'Write client released');
}
}
}
// @INFO: Add exports for the query functions
export { executeQuery, executeReadOnlyQuery, executeWriteQuery, getServer };
// @INFO: Server startup and shutdown
async function runServer() {
try {
log('error', 'Attempting to test database connection...');
// @INFO: Test the connection before fully starting the server
const pool = await getPool();
const client = await pool.connect();
log('error', 'Database connection test successful');
client.release();
const server = await getServer();
const transport = new StdioServerTransport();
log('error', 'Connecting server to transport...');
await server.connect(transport);
log('error', 'Server connected to transport successfully');
}
catch (error) {
log('error', 'Fatal error during server startup:', error);
safeExit(1);
}
}
const shutdown = async (signal) => {
log('error', `Received ${signal}. Shutting down...`);
try {
// @INFO: Only attempt to close the pool if it was created
if (poolPromise) {
const pool = await poolPromise;
await pool.end();
log('error', 'PostgreSQL pool closed successfully');
}
}
catch (err) {
log('error', 'Error closing pool:', err);
throw err;
}
};
process.on('SIGINT', async () => {
try {
await shutdown('SIGINT');
process.exit(0);
}
catch (err) {
log('error', 'Error during SIGINT shutdown:', err);
safeExit(1);
}
});
process.on('SIGTERM', async () => {
try {
await shutdown('SIGTERM');
process.exit(0);
}
catch (err) {
log('error', 'Error during SIGTERM shutdown:', err);
safeExit(1);
}
});
// @INFO: Add unhandled error listeners
process.on('uncaughtException', (error) => {
log('error', 'Uncaught exception:', error);
safeExit(1);
});
process.on('unhandledRejection', (reason, promise) => {
log('error', 'Unhandled rejection at:', promise, 'reason:', reason);
safeExit(1);
});
runServer().catch((error) => {
log('error', 'Server error:', error);
safeExit(1);
});
//# sourceMappingURL=index.js.map