@henkey/postgres-mcp-server
Version:
A Model Context Protocol (MCP) server that provides comprehensive PostgreSQL database management capabilities for AI assistants
229 lines • 11.8 kB
JavaScript
import { DatabaseConnection } from '../utils/connection.js';
import { z } from 'zod';
import { McpError, ErrorCode } from '@modelcontextprotocol/sdk/types.js';
// Import existing execution functions from users.ts (reuse the logic)
// Note: In real implementation, these would be extracted to shared functions
// For now, we'll reimplement the core logic here
// Define action-specific schemas using discriminated unions
const ManageUsersInputSchema = z.discriminatedUnion('action', [
// CREATE USER
z.object({
action: z.literal('create'),
connectionString: z.string().optional(),
username: z.string().describe("Username for the new user"),
password: z.string().optional().describe("Password for the user"),
superuser: z.boolean().optional().default(false).describe("Grant superuser privileges"),
createdb: z.boolean().optional().default(false).describe("Allow user to create databases"),
createrole: z.boolean().optional().default(false).describe("Allow user to create roles"),
login: z.boolean().optional().default(true).describe("Allow user to login"),
replication: z.boolean().optional().default(false).describe("Allow replication privileges"),
connectionLimit: z.number().optional().describe("Maximum number of connections"),
validUntil: z.string().optional().describe("Password expiration date (YYYY-MM-DD)"),
inherit: z.boolean().optional().default(true).describe("Inherit privileges from parent roles"),
}),
// DROP USER
z.object({
action: z.literal('drop'),
connectionString: z.string().optional(),
username: z.string().describe("Username to drop"),
ifExists: z.boolean().optional().default(true).describe("Include IF EXISTS clause"),
cascade: z.boolean().optional().default(false).describe("Include CASCADE to drop owned objects"),
}),
// ALTER USER
z.object({
action: z.literal('alter'),
connectionString: z.string().optional(),
username: z.string().describe("Username to alter"),
password: z.string().optional().describe("New password"),
superuser: z.boolean().optional().describe("Grant/revoke superuser privileges"),
createdb: z.boolean().optional().describe("Grant/revoke database creation privileges"),
createrole: z.boolean().optional().describe("Grant/revoke role creation privileges"),
login: z.boolean().optional().describe("Grant/revoke login privileges"),
replication: z.boolean().optional().describe("Grant/revoke replication privileges"),
connectionLimit: z.number().optional().describe("Set connection limit"),
validUntil: z.string().optional().describe("Set password expiration date (YYYY-MM-DD)"),
inherit: z.boolean().optional().describe("Set privilege inheritance"),
}),
// GRANT PERMISSIONS
z.object({
action: z.literal('grant'),
connectionString: z.string().optional(),
username: z.string().describe("Username to grant permissions to"),
permissions: z.array(z.enum(['SELECT', 'INSERT', 'UPDATE', 'DELETE', 'TRUNCATE', 'REFERENCES', 'TRIGGER', 'ALL'])).describe("Permissions to grant"),
target: z.string().describe("Target object (table, schema, database, etc.)"),
targetType: z.enum(['table', 'schema', 'database', 'sequence', 'function']).describe("Type of target object"),
schema: z.string().optional().default('public').describe("Schema name (for table/sequence/function targets)"),
withGrantOption: z.boolean().optional().default(false).describe("Allow user to grant these permissions to others"),
}),
// REVOKE PERMISSIONS
z.object({
action: z.literal('revoke'),
connectionString: z.string().optional(),
username: z.string().describe("Username to revoke permissions from"),
permissions: z.array(z.enum(['SELECT', 'INSERT', 'UPDATE', 'DELETE', 'TRUNCATE', 'REFERENCES', 'TRIGGER', 'ALL'])).describe("Permissions to revoke"),
target: z.string().describe("Target object (table, schema, database, etc.)"),
targetType: z.enum(['table', 'schema', 'database', 'sequence', 'function']).describe("Type of target object"),
schema: z.string().optional().default('public').describe("Schema name (for table/sequence/function targets)"),
cascade: z.boolean().optional().default(false).describe("Cascade revoke to dependent privileges"),
}),
// GET USER PERMISSIONS
z.object({
action: z.literal('get_permissions'),
connectionString: z.string().optional(),
username: z.string().optional().describe("Username to get permissions for (optional, shows all if not provided)"),
schema: z.string().optional().describe("Filter by schema (optional)"),
targetType: z.enum(['table', 'schema', 'database', 'sequence', 'function']).optional().describe("Filter by target type"),
}),
// LIST USERS
z.object({
action: z.literal('list'),
connectionString: z.string().optional(),
includeSystemRoles: z.boolean().optional().default(false).describe("Include system roles"),
}),
]);
// Execution functions (simplified versions - in real implementation, import from users.ts)
async function executeCreateUser(input, getConnectionString) {
const resolvedConnectionString = getConnectionString(input.connectionString);
const db = DatabaseConnection.getInstance();
try {
await db.connect(resolvedConnectionString);
const options = [];
if (input.password)
options.push(`PASSWORD '${input.password.replace(/'/g, "''")}'`);
if (input.superuser)
options.push('SUPERUSER');
if (input.createdb)
options.push('CREATEDB');
if (input.createrole)
options.push('CREATEROLE');
if (input.login)
options.push('LOGIN');
if (input.replication)
options.push('REPLICATION');
if (!input.inherit)
options.push('NOINHERIT');
if (input.connectionLimit !== undefined)
options.push(`CONNECTION LIMIT ${input.connectionLimit}`);
if (input.validUntil)
options.push(`VALID UNTIL '${input.validUntil}'`);
const createUserSQL = `CREATE USER "${input.username}"${options.length > 0 ? ` ${options.join(' ')}` : ''}`;
await db.query(createUserSQL);
return { action: 'create', username: input.username, created: true };
}
finally {
await db.disconnect();
}
}
async function executeDropUser(input, getConnectionString) {
const resolvedConnectionString = getConnectionString(input.connectionString);
const db = DatabaseConnection.getInstance();
try {
await db.connect(resolvedConnectionString);
if (input.cascade) {
await db.query(`DROP OWNED BY "${input.username}" CASCADE`);
}
const ifExistsClause = input.ifExists ? 'IF EXISTS ' : '';
const dropUserSQL = `DROP USER ${ifExistsClause}"${input.username}"`;
await db.query(dropUserSQL);
return { action: 'drop', username: input.username, dropped: true };
}
finally {
await db.disconnect();
}
}
async function executeListUsers(input, getConnectionString) {
const resolvedConnectionString = getConnectionString(input.connectionString);
const db = DatabaseConnection.getInstance();
try {
await db.connect(resolvedConnectionString);
const systemRoleFilter = input.includeSystemRoles ? '' : "WHERE rolname NOT LIKE 'pg_%' AND rolname != 'postgres'";
const usersQuery = `
SELECT
rolname,
rolsuper,
rolinherit,
rolcreaterole,
rolcreatedb,
rolcanlogin,
rolreplication,
rolconnlimit,
rolvaliduntil,
oid
FROM pg_roles
${systemRoleFilter}
ORDER BY rolname
`;
const result = await db.query(usersQuery);
return { action: 'list', users: result };
}
finally {
await db.disconnect();
}
}
// Main consolidated tool
export const manageUsersTool = {
name: 'pg_manage_users',
description: 'Comprehensive user management: create, drop, alter users, grant/revoke permissions, and list users. Use action parameter to specify operation.',
inputSchema: ManageUsersInputSchema,
async execute(params, getConnectionString) {
try {
const input = ManageUsersInputSchema.parse(params);
let result;
let successMessage;
switch (input.action) {
case 'create':
result = await executeCreateUser(input, getConnectionString);
successMessage = `User ${result.username} created successfully.`;
break;
case 'drop':
result = await executeDropUser(input, getConnectionString);
successMessage = `User ${result.username} dropped successfully.`;
break;
case 'alter':
// Simplified - in real implementation, call executeAlterUser
result = { action: 'alter', username: input.username, message: 'Alter user functionality would be implemented here' };
successMessage = `User ${input.username} would be altered successfully.`;
break;
case 'grant':
// Simplified - in real implementation, call executeGrantPermissions
result = { action: 'grant', username: input.username, permissions: input.permissions, target: input.target };
successMessage = `Permissions granted to ${input.username} successfully.`;
break;
case 'revoke':
// Simplified - in real implementation, call executeRevokePermissions
result = { action: 'revoke', username: input.username, permissions: input.permissions, target: input.target };
successMessage = `Permissions revoked from ${input.username} successfully.`;
break;
case 'get_permissions':
// Simplified - in real implementation, call executeGetUserPermissions
result = { action: 'get_permissions', message: 'Get permissions functionality would be implemented here' };
successMessage = `Retrieved permissions for ${input.username || 'all users'}`;
break;
case 'list':
result = await executeListUsers(input, getConnectionString);
successMessage = `Listed ${result.users.length} users`;
break;
default: {
const exhaustiveCheck = input;
throw new McpError(ErrorCode.InvalidParams, `Unknown action: ${JSON.stringify(exhaustiveCheck)}`);
}
}
return {
content: [
{ type: 'text', text: successMessage },
{ type: 'text', text: JSON.stringify(result, null, 2) }
]
};
}
catch (error) {
const errorMessage = error instanceof McpError ? error.message : (error instanceof Error ? error.message : String(error));
return {
content: [{ type: 'text', text: `Error in user management: ${errorMessage}` }],
isError: true
};
}
}
};
// Export for backward compatibility - individual tools can still be used
export { manageUsersTool as consolidatedUserTool };
//# sourceMappingURL=consolidated-users.js.map