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,137 @@
+import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
+import * as z from 'zod'
+import type { Registry } from '../config/registry.js'
+import {
+    connectionConfigSchema,
+    defaultPort,
+    type ResolvedConnection,
+    sshConfigSchema
+} from '../config/types.js'
+import type { Manager } from '../db/manager.js'
+import { formatDbError } from '../format.js'
+import { errorMessage, fail, ok } from './respond.js'
+
+const patchSchema = z
+    .object({
+        name: z
+            .string()
+            .regex(/^[a-zA-Z0-9_-]+$/)
+            .optional(),
+        type: z.enum(['postgres', 'mysql']).optional(),
+        host: z.string().min(1).optional(),
+        port: z.number().int().positive().nullable().optional(),
+        user: z.string().min(1).optional(),
+        password: z.string().nullable().optional(),
+        database: z.string().nullable().optional(),
+        readonly: z.boolean().optional(),
+        ssh: sshConfigSchema.nullable().optional()
+    })
+    .strict()
+
+const publicView = (resolved: ResolvedConnection) => ({
+    name: resolved.config.name,
+    type: resolved.config.type,
+    host: resolved.config.host,
+    port: resolved.config.port ?? defaultPort(resolved.config.type),
+    database: resolved.config.database ?? null,
+    readonly: resolved.config.readonly,
+    source: resolved.source,
+    ssh: Boolean(resolved.config.ssh)
+})
+
+export const registerConnectionTools = (
+    server: McpServer,
+    registry: Registry,
+    manager: Manager
+): void => {
+    server.registerTool(
+        'list_connections',
+        {
+            description:
+                'List configured database connections with their source layer (env/config/store). Secrets are omitted.'
+        },
+        async () => ok(registry.list().map(publicView))
+    )
+
+    server.registerTool(
+        'add_connection',
+        {
+            description:
+                'Add a named postgres/mysql connection to the persistent store. Supports optional ssh tunnel config and a readonly flag.',
+            inputSchema: connectionConfigSchema.shape
+        },
+        async (config) => {
+            try {
+                return ok({ added: publicView(registry.add(connectionConfigSchema.parse(config))) })
+            } catch (error) {
+                return fail(errorMessage(error))
+            }
+        }
+    )
+
+    server.registerTool(
+        'update_connection',
+        {
+            description:
+                'Patch a stored connection by name. Set a field to null to remove it (e.g. "ssh": null). Only store-sourced connections can be edited.',
+            inputSchema: {
+                name: z.string(),
+                patch: patchSchema
+            }
+        },
+        async ({ name, patch }) => {
+            try {
+                const updated = registry.update(name, patch)
+                await manager.invalidate(name)
+                return ok({ updated: publicView(updated) })
+            } catch (error) {
+                return fail(errorMessage(error))
+            }
+        }
+    )
+
+    server.registerTool(
+        'remove_connection',
+        {
+            description: 'Remove a stored connection by name and drop its cached pools/tunnel.',
+            inputSchema: { name: z.string() }
+        },
+        async ({ name }) => {
+            try {
+                registry.remove(name)
+                await manager.invalidate(name)
+                return ok({ removed: name })
+            } catch (error) {
+                return fail(errorMessage(error))
+            }
+        }
+    )
+
+    server.registerTool(
+        'test_connection',
+        {
+            description:
+                'Verify a connection works end to end (including ssh tunnel if configured). Reports server version and latency.',
+            inputSchema: { name: z.string() }
+        },
+        async ({ name }) => {
+            let managed: Awaited<ReturnType<Manager['get']>>
+            try {
+                managed = await manager.get(name)
+            } catch (error) {
+                return fail(errorMessage(error))
+            }
+            const started = performance.now()
+            try {
+                const version = await managed.driver.serverVersion()
+                return ok({
+                    ok: true,
+                    version,
+                    latencyMs: Math.round(performance.now() - started)
+                })
+            } catch (error) {
+                return fail(formatDbError(managed.config.type, error))
+            }
+        }
+    )
+}