feat: add connection management tools

2026-06-12 00:06:30 +05:00
parent 0c92dd6755
commit 96404fe0d3
2 changed files with 294 additions and 0 deletions
@@ -0,0 +1,157 @@
+import { mkdtempSync, rmSync } from 'node:fs'
+import { tmpdir } from 'node:os'
+import { join } from 'node:path'
+import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+import { createRegistry } from '../../../src/config/registry.js'
+import { writeStore } from '../../../src/config/store.js'
+import type { ConnectionConfig } from '../../../src/config/types.js'
+import { registerConnectionTools } from '../../../src/tools/connections.js'
+import { callTool, connectClient, type FakeManager, fakeManager } from '../helpers.js'
+
+const conn = (name: string, extra: Partial<ConnectionConfig> = {}): ConnectionConfig => ({
+    name,
+    type: 'postgres',
+    host: 'localhost',
+    user: 'postgres',
+    readonly: false,
+    ...extra
+})
+
+describe('connection tools', () => {
+    let dir: string
+    let storePath: string
+    let manager: FakeManager
+    let client: Awaited<ReturnType<typeof connectClient>>
+
+    beforeEach(async () => {
+        dir = mkdtempSync(join(tmpdir(), 'dbmole-tools-'))
+        storePath = join(dir, 'connections.json')
+        writeStore(storePath, [conn('existing', { password: 'secret', database: 'app' })])
+        const registry = createRegistry({ storePath, env: {} })
+        manager = fakeManager()
+        const server = new McpServer({ name: 'test', version: '0.0.0' })
+        registerConnectionTools(server, registry, manager)
+        client = await connectClient(server)
+    })
+
+    afterEach(() => {
+        rmSync(dir, { recursive: true, force: true })
+    })
+
+    it('list_connections returns public view without secrets', async () => {
+        const response = await callTool(client, 'list_connections')
+        expect(response.isError).toBe(false)
+        const list = response.json() as Array<Record<string, unknown>>
+        expect(list).toEqual([
+            {
+                name: 'existing',
+                type: 'postgres',
+                host: 'localhost',
+                port: 5432,
+                database: 'app',
+                readonly: false,
+                source: 'store',
+                ssh: false
+            }
+        ])
+        expect(response.text).not.toContain('secret')
+    })
+
+    it('add_connection persists and reports the new connection', async () => {
+        const response = await callTool(client, 'add_connection', {
+            name: 'fresh',
+            type: 'mysql',
+            host: 'db',
+            user: 'root'
+        })
+        expect(response.isError).toBe(false)
+        const list = await callTool(client, 'list_connections')
+        expect((list.json() as Array<{ name: string }>).map((c) => c.name)).toEqual([
+            'existing',
+            'fresh'
+        ])
+    })
+
+    it('add_connection rejects duplicates with isError', async () => {
+        const response = await callTool(client, 'add_connection', {
+            name: 'existing',
+            type: 'postgres',
+            host: 'x',
+            user: 'u'
+        })
+        expect(response.isError).toBe(true)
+        expect(response.text).toContain('already exists')
+    })
+
+    it('update_connection patches and invalidates the manager cache', async () => {
+        const response = await callTool(client, 'update_connection', {
+            name: 'existing',
+            patch: { host: 'new-host', database: null }
+        })
+        expect(response.isError).toBe(false)
+        expect(manager.invalidate).toHaveBeenCalledWith('existing')
+        const list = await callTool(client, 'list_connections')
+        const updated = (list.json() as Array<Record<string, unknown>>)[0]
+        expect(updated.host).toBe('new-host')
+        expect(updated.database).toBeNull()
+    })
+
+    it('update_connection surfaces validation errors', async () => {
+        const response = await callTool(client, 'update_connection', {
+            name: 'existing',
+            patch: { type: 'oracle' }
+        })
+        expect(response.isError).toBe(true)
+    })
+
+    it('remove_connection deletes and invalidates', async () => {
+        const response = await callTool(client, 'remove_connection', { name: 'existing' })
+        expect(response.isError).toBe(false)
+        expect(manager.invalidate).toHaveBeenCalledWith('existing')
+        const list = await callTool(client, 'list_connections')
+        expect(list.json()).toEqual([])
+    })
+
+    it('remove_connection of unknown name reports available connections', async () => {
+        const response = await callTool(client, 'remove_connection', { name: 'ghost' })
+        expect(response.isError).toBe(true)
+        expect(response.text).toContain('existing')
+    })
+
+    it('test_connection reports version and latency', async () => {
+        manager.get.mockResolvedValue({
+            driver: { serverVersion: vi.fn(async () => '17.2') },
+            config: conn('existing'),
+            source: 'store'
+        })
+        const response = await callTool(client, 'test_connection', { name: 'existing' })
+        expect(response.isError).toBe(false)
+        const payload = response.json() as Record<string, unknown>
+        expect(payload.ok).toBe(true)
+        expect(payload.version).toBe('17.2')
+        expect(typeof payload.latencyMs).toBe('number')
+    })
+
+    it('test_connection formats driver failures', async () => {
+        manager.get.mockResolvedValue({
+            driver: {
+                serverVersion: vi.fn(async () => {
+                    throw Object.assign(new Error('connect ECONNREFUSED'), { code: 'ECONNREFUSED' })
+                })
+            },
+            config: conn('existing'),
+            source: 'store'
+        })
+        const response = await callTool(client, 'test_connection', { name: 'existing' })
+        expect(response.isError).toBe(true)
+        expect(response.text).toContain('ECONNREFUSED')
+    })
+
+    it('test_connection reports manager failures (e.g. tunnel)', async () => {
+        manager.get.mockRejectedValue(new Error('ssh auth failed'))
+        const response = await callTool(client, 'test_connection', { name: 'existing' })
+        expect(response.isError).toBe(true)
+        expect(response.text).toContain('ssh auth failed')
+    })
+})