├── .gitignore ├── README.md ├── package.json ├── LICENSE └── index.mjs /.gitignore: -------------------------------------------------------------------------------- 1 | node_modules 2 | *.tgz 3 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # postgres-context-server 2 | 3 | This module implements a Model Context Protocol server for a PostgreSQL database. 4 | -------------------------------------------------------------------------------- /package.json: -------------------------------------------------------------------------------- 1 | { 2 | "name": "@zeddotdev/postgres-context-server", 3 | "version": "0.1.5", 4 | "description": "a model context protocol server for postgres", 5 | "main": "index.mjs", 6 | "type": "module", 7 | "publishConfig": { 8 | "access": "public" 9 | }, 10 | "files": [ 11 | "vendor", 12 | "index.mjs" 13 | ], 14 | "repository": { 15 | "type": "git", 16 | "url": "git+https://github.com/zed-industries/postgres-context-server.git" 17 | }, 18 | "author": "", 19 | "license": "MIT", 20 | "bugs": { 21 | "url": "https://github.com/zed-industries/postgres-context-server/issues" 22 | }, 23 | "homepage": "https://github.com/zed-industries/postgres-context-server#readme", 24 | "dependencies": { 25 | "@modelcontextprotocol/sdk": "^0.4.0", 26 | "pg": "^8.13.1", 27 | "zod": "^3.23.8" 28 | } 29 | } 30 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024 Zed Industries, Inc. 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /index.mjs: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env node 2 | 3 | import pg from "pg"; 4 | import { Server } from "@modelcontextprotocol/sdk/server/index.js"; 5 | import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js"; 6 | import { 7 | CallToolRequestSchema, 8 | ListResourcesRequestSchema, 9 | ListPromptsRequestSchema, 10 | ListToolsRequestSchema, 11 | ReadResourceRequestSchema, 12 | GetPromptRequestSchema, 13 | CompleteRequestSchema, 14 | } from "@modelcontextprotocol/sdk/types.js"; 15 | 16 | const server = new Server({ 17 | name: "postgres-context-server", 18 | version: "0.1.0", 19 | }); 20 | 21 | const databaseUrl = process.env.DATABASE_URL; 22 | if (typeof databaseUrl == null || databaseUrl.trim().length === 0) { 23 | console.error("Please provide a DATABASE_URL environment variable"); 24 | process.exit(1); 25 | } 26 | 27 | const resourceBaseUrl = new URL(databaseUrl); 28 | resourceBaseUrl.protocol = "postgres:"; 29 | resourceBaseUrl.password = ""; 30 | 31 | process.stderr.write("starting server\n"); 32 | 33 | const pool = new pg.Pool({ 34 | connectionString: databaseUrl, 35 | }); 36 | 37 | const SCHEMA_PATH = "schema"; 38 | const SCHEMA_PROMPT_NAME = "pg-schema"; 39 | const ALL_TABLES = "all-tables"; 40 | 41 | server.setRequestHandler(ListResourcesRequestSchema, async () => { 42 | const client = await pool.connect(); 43 | try { 44 | const result = await client.query( 45 | "SELECT table_name FROM information_schema.tables WHERE table_schema = 'public'", 46 | ); 47 | return { 48 | resources: result.rows.map((row) => ({ 49 | uri: new URL(`${row.table_name}/${SCHEMA_PATH}`, resourceBaseUrl).href, 50 | mimeType: "application/json", 51 | name: `"${row.table_name}" database schema`, 52 | })), 53 | }; 54 | } finally { 55 | client.release(); 56 | } 57 | }); 58 | 59 | server.setRequestHandler(ReadResourceRequestSchema, async (request) => { 60 | const resourceUrl = new URL(request.params.uri); 61 | 62 | const pathComponents = resourceUrl.pathname.split("/"); 63 | const schema = pathComponents.pop(); 64 | const tableName = pathComponents.pop(); 65 | 66 | if (schema !== SCHEMA_PATH) { 67 | throw new Error("Invalid resource URI"); 68 | } 69 | 70 | const client = await pool.connect(); 71 | try { 72 | const result = await client.query( 73 | "SELECT column_name, data_type FROM information_schema.columns WHERE table_name = $1", 74 | [tableName], 75 | ); 76 | 77 | return { 78 | contents: [ 79 | { 80 | uri: request.params.uri, 81 | mimeType: "application/json", 82 | text: JSON.stringify(result.rows, null, 2), 83 | }, 84 | ], 85 | }; 86 | } finally { 87 | client.release(); 88 | } 89 | }); 90 | 91 | server.setRequestHandler(ListToolsRequestSchema, async () => { 92 | return { 93 | tools: [ 94 | { 95 | name: "pg-schema", 96 | description: "Returns the schema for a Postgres database.", 97 | inputSchema: { 98 | type: "object", 99 | properties: { 100 | mode: { 101 | type: "string", 102 | enum: ["all", "specific"], 103 | description: "Mode of schema retrieval", 104 | }, 105 | tableName: { 106 | type: "string", 107 | description: 108 | "Name of the specific table (required if mode is 'specific')", 109 | }, 110 | }, 111 | required: ["mode"], 112 | if: { 113 | properties: { mode: { const: "specific" } }, 114 | }, 115 | then: { 116 | required: ["tableName"], 117 | }, 118 | }, 119 | }, 120 | { 121 | name: "query", 122 | description: "Run a read-only SQL query", 123 | inputSchema: { 124 | type: "object", 125 | properties: { 126 | sql: { type: "string" }, 127 | }, 128 | }, 129 | }, 130 | ], 131 | }; 132 | }); 133 | 134 | server.setRequestHandler(CallToolRequestSchema, async (request) => { 135 | if (request.params.name === "pg-schema") { 136 | const mode = request.params.arguments?.mode; 137 | 138 | const tableName = (() => { 139 | switch (mode) { 140 | case "specific": { 141 | const tableName = request.params.arguments?.tableName; 142 | 143 | if (typeof tableName !== "string" || tableName.length === 0) { 144 | throw new Error(`Invalid tableName: ${tableName}`); 145 | } 146 | 147 | return tableName; 148 | } 149 | case "all": { 150 | return ALL_TABLES; 151 | } 152 | default: 153 | throw new Error(`Invalid mode: ${mode}`); 154 | } 155 | })(); 156 | 157 | const client = await pool.connect(); 158 | 159 | try { 160 | const sql = await getSchema(client, tableName); 161 | 162 | return { 163 | content: [{ type: "text", text: sql }], 164 | }; 165 | } finally { 166 | client.release(); 167 | } 168 | } 169 | 170 | if (request.params.name === "query") { 171 | const sql = request.params.arguments?.sql; 172 | 173 | const client = await pool.connect(); 174 | try { 175 | await client.query("BEGIN TRANSACTION READ ONLY"); 176 | // Force a prepared statement: Prevents multiple statements in the same query. 177 | // Name is unique per session, but we use a single session per query. 178 | const result = await client.query({ 179 | name: "sandboxed-statement", 180 | text: sql, 181 | values: [], 182 | }); 183 | return { 184 | content: [ 185 | { type: "text", text: JSON.stringify(result.rows, undefined, 2) }, 186 | ], 187 | }; 188 | } catch (error) { 189 | throw error; 190 | } finally { 191 | client 192 | .query("ROLLBACK") 193 | .catch((error) => 194 | console.warn("Could not roll back transaction:", error), 195 | ); 196 | 197 | // Destroy session to clean up resources. 198 | client.release(true); 199 | } 200 | } 201 | 202 | throw new Error("Tool not found"); 203 | }); 204 | 205 | server.setRequestHandler(CompleteRequestSchema, async (request) => { 206 | process.stderr.write("Handling completions/complete request\n"); 207 | 208 | if (request.params.ref.name === SCHEMA_PROMPT_NAME) { 209 | const tableNameQuery = request.params.argument.value; 210 | const alreadyHasArg = /\S*\s/.test(tableNameQuery); 211 | 212 | if (alreadyHasArg) { 213 | return { 214 | completion: { 215 | values: [], 216 | }, 217 | }; 218 | } 219 | 220 | const client = await pool.connect(); 221 | try { 222 | const result = await client.query( 223 | "SELECT table_name FROM information_schema.tables WHERE table_schema = 'public'", 224 | ); 225 | const tables = result.rows.map((row) => row.table_name); 226 | return { 227 | completion: { 228 | values: [ALL_TABLES, ...tables], 229 | }, 230 | }; 231 | } finally { 232 | client.release(); 233 | } 234 | } 235 | 236 | throw new Error("unknown prompt"); 237 | }); 238 | 239 | server.setRequestHandler(ListPromptsRequestSchema, async () => { 240 | process.stderr.write("Handling prompts/list request\n"); 241 | 242 | return { 243 | prompts: [ 244 | { 245 | name: SCHEMA_PROMPT_NAME, 246 | description: 247 | "Retrieve the schema for a given table in the postgres database", 248 | arguments: [ 249 | { 250 | name: "tableName", 251 | description: "the table to describe", 252 | required: true, 253 | }, 254 | ], 255 | }, 256 | ], 257 | }; 258 | }); 259 | 260 | server.setRequestHandler(GetPromptRequestSchema, async (request) => { 261 | process.stderr.write("Handling prompts/get request\n"); 262 | 263 | if (request.params.name === SCHEMA_PROMPT_NAME) { 264 | const tableName = request.params.arguments?.tableName; 265 | 266 | if (typeof tableName !== "string" || tableName.length === 0) { 267 | throw new Error(`Invalid tableName: ${tableName}`); 268 | } 269 | 270 | const client = await pool.connect(); 271 | 272 | try { 273 | const sql = await getSchema(client, tableName); 274 | 275 | return { 276 | description: 277 | tableName === ALL_TABLES 278 | ? "all table schemas" 279 | : `${tableName} schema`, 280 | messages: [ 281 | { 282 | role: "user", 283 | content: { 284 | type: "text", 285 | text: sql, 286 | }, 287 | }, 288 | ], 289 | }; 290 | } finally { 291 | client.release(); 292 | } 293 | } 294 | 295 | throw new Error(`Prompt '${request.params.name}' not implemented`); 296 | }); 297 | 298 | /** 299 | * @param tableNameOrAll {string} 300 | */ 301 | async function getSchema(client, tableNameOrAll) { 302 | const select = 303 | "SELECT column_name, data_type, is_nullable, column_default, table_name FROM information_schema.columns"; 304 | 305 | let result; 306 | if (tableNameOrAll === ALL_TABLES) { 307 | result = await client.query( 308 | `${select} WHERE table_schema NOT IN ('pg_catalog', 'information_schema')`, 309 | ); 310 | } else { 311 | result = await client.query(`${select} WHERE table_name = $1`, [ 312 | tableNameOrAll, 313 | ]); 314 | } 315 | 316 | const allTableNames = Array.from( 317 | new Set(result.rows.map((row) => row.table_name).sort()), 318 | ); 319 | 320 | let sql = "```sql\n"; 321 | for (let i = 0, len = allTableNames.length; i < len; i++) { 322 | const tableName = allTableNames[i]; 323 | if (i > 0) { 324 | sql += "\n"; 325 | } 326 | 327 | sql += [ 328 | `create table "${tableName}" (`, 329 | result.rows 330 | .filter((row) => row.table_name === tableName) 331 | .map((row) => { 332 | const notNull = row.is_nullable === "NO" ? "" : " not null"; 333 | const defaultValue = 334 | row.column_default != null ? ` default ${row.column_default}` : ""; 335 | return ` "${row.column_name}" ${row.data_type}${notNull}${defaultValue}`; 336 | }) 337 | .join(",\n"), 338 | ");", 339 | ].join("\n"); 340 | sql += "\n"; 341 | } 342 | sql += "```"; 343 | 344 | return sql; 345 | } 346 | 347 | async function runServer() { 348 | const transport = new StdioServerTransport(); 349 | await server.connect(transport); 350 | } 351 | 352 | runServer().catch((error) => { 353 | console.error(error); 354 | process.exit(1); 355 | }); 356 | --------------------------------------------------------------------------------