feat: add schema inspection tools

2026-06-12 00:10:55 +05:00
parent aaa2844d46
commit 7514aedb1e
2 changed files with 164 additions and 0 deletions
@@ -0,0 +1,74 @@
+import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
+import type { CallToolResult } from '@modelcontextprotocol/sdk/types.js'
+import * as z from 'zod'
+import type { ManagedConnection, Manager } from '../db/manager.js'
+import { formatDbError } from '../format.js'
+import { errorMessage, fail, ok } from './respond.js'
+
+const withManaged = async (
+    manager: Manager,
+    connection: string,
+    fn: (managed: ManagedConnection) => Promise<unknown>
+): Promise<CallToolResult> => {
+    let managed: ManagedConnection
+    try {
+        managed = await manager.get(connection)
+    } catch (error) {
+        return fail(errorMessage(error))
+    }
+    try {
+        return ok(await fn(managed))
+    } catch (error) {
+        return fail(formatDbError(managed.config.type, error))
+    }
+}
+
+export const registerSchemaTools = (server: McpServer, manager: Manager): void => {
+    server.registerTool(
+        'list_databases',
+        {
+            description:
+                'List databases on a connection with their size in bytes (system databases hidden).',
+            inputSchema: {
+                connection: z.string().describe('connection name, see list_connections')
+            }
+        },
+        async ({ connection }) =>
+            withManaged(manager, connection, (managed) => managed.driver.listDatabases())
+    )
+
+    server.registerTool(
+        'list_tables',
+        {
+            description:
+                'List tables with approximate row counts. For postgres, optionally filter by schema; for mysql, schema is the database.',
+            inputSchema: {
+                connection: z.string(),
+                database: z.string().optional(),
+                schema: z.string().optional()
+            }
+        },
+        async ({ connection, database, schema }) =>
+            withManaged(manager, connection, (managed) =>
+                managed.driver.listTables({ database, schema })
+            )
+    )
+
+    server.registerTool(
+        'describe_table',
+        {
+            description:
+                'Describe a table: columns (type, nullable, default), primary key, indexes and foreign keys.',
+            inputSchema: {
+                connection: z.string(),
+                table: z.string(),
+                database: z.string().optional(),
+                schema: z.string().optional().describe('postgres schema, defaults to public')
+            }
+        },
+        async ({ connection, table, database, schema }) =>
+            withManaged(manager, connection, (managed) =>
+                managed.driver.describeTable({ table, database, schema })
+            )
+    )
+}