feat: add execute_sql tool

2026-06-12 00:08:41 +05:00
parent 96404fe0d3
commit aaa2844d46
2 changed files with 133 additions and 0 deletions
@@ -0,0 +1,85 @@
+import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { registerQueryTools } from '../../../src/tools/query.js'
+import { callTool, connectClient, type FakeManager, fakeManager } from '../helpers.js'
+
+describe('execute_sql', () => {
+    let manager: FakeManager
+    let query: ReturnType<typeof vi.fn>
+    let client: Awaited<ReturnType<typeof connectClient>>
+
+    beforeEach(async () => {
+        query = vi.fn(async () => ({
+            columns: ['n'],
+            rows: [[1]],
+            rowCount: 1,
+            truncated: false
+        }))
+        manager = fakeManager()
+        manager.get.mockImplementation(async () => ({
+            driver: { query },
+            config: { name: 'c', type: 'postgres', host: 'h', user: 'u', readonly: false },
+            source: 'store'
+        }))
+        const server = new McpServer({ name: 'test', version: '0.0.0' })
+        registerQueryTools(server, manager)
+        client = await connectClient(server)
+    })
+
+    it('executes sql and returns the result payload', async () => {
+        const response = await callTool(client, 'execute_sql', {
+            connection: 'c',
+            sql: 'select 1 as n'
+        })
+        expect(response.isError).toBe(false)
+        expect(response.json()).toEqual({
+            columns: ['n'],
+            rows: [[1]],
+            rowCount: 1,
+            truncated: false
+        })
+        expect(query).toHaveBeenCalledWith({
+            sql: 'select 1 as n',
+            params: [],
+            database: undefined,
+            rowLimit: 100
+        })
+    })
+
+    it('passes params, database and clamped rowLimit through', async () => {
+        await callTool(client, 'execute_sql', {
+            connection: 'c',
+            sql: 'select $1',
+            params: [42],
+            database: 'other',
+            rowLimit: 5
+        })
+        expect(query).toHaveBeenCalledWith({
+            sql: 'select $1',
+            params: [42],
+            database: 'other',
+            rowLimit: 5
+        })
+    })
+
+    it('formats driver errors with the engine prefix', async () => {
+        query.mockRejectedValue(
+            Object.assign(new Error('relation "x" does not exist'), { code: '42P01' })
+        )
+        const response = await callTool(client, 'execute_sql', { connection: 'c', sql: 'select' })
+        expect(response.isError).toBe(true)
+        expect(response.text).toBe('[postgres 42P01] relation "x" does not exist')
+    })
+
+    it('reports unknown connections without engine prefix', async () => {
+        manager.get.mockRejectedValue(
+            new Error("connection 'nope' not found. Available connections: c")
+        )
+        const response = await callTool(client, 'execute_sql', {
+            connection: 'nope',
+            sql: 'select 1'
+        })
+        expect(response.isError).toBe(true)
+        expect(response.text).toContain("connection 'nope' not found")
+    })
+})