dbQuery.ts•2.29 kB
import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { z } from "zod";
import {
  executeDatabaseQuery,
  type DatabaseQueryResult
} from "../services/databaseService.js";
import { createProgressReporter } from "../utils/progress.js";
const DbQueryInputSchema = z.object({
  profile: z.string().describe("Database profile to use"),
  query: z.string().min(1).describe("SQL query to execute"),
  parameters: z.array(z.unknown()).optional().describe("Positional parameters for the query"),
  timeoutMs: z.number().int().positive().optional().describe("Query timeout in milliseconds"),
  rowLimit: z.number().int().positive().optional().describe("Override maximum number of rows to return")
});
const DbQueryOutputShape = {
  rows: z.array(z.record(z.string(), z.unknown())),
  rowCount: z.number().int().nonnegative(),
  truncated: z.boolean(),
  durationMs: z.number().int().nonnegative()
};
const DbQueryOutputSchema = z.object(DbQueryOutputShape);
export function registerDbTool(server: McpServer): void {
  server.registerTool(
    "dbQuery",
    {
      description: "Execute a SQL query on a PostgreSQL database using a configured profile",
      inputSchema: DbQueryInputSchema.shape,
      outputSchema: DbQueryOutputShape
    },
    async (args, extra) => {
      const progress = createProgressReporter(extra, "dbQuery");
      const total = 1;
      progress?.({ progress: 0, total, message: "Dispatching database query" });
      const result: DatabaseQueryResult = await executeDatabaseQuery(
        args.profile,
        args.query,
        args.parameters,
        {
          timeoutMs: args.timeoutMs,
          rowLimit: args.rowLimit,
          requestId: String(extra.requestId),
          tool: "dbQuery",
          signal: extra.signal,
          onProgress: (update) => {
            progress?.({
              progress: update.progress,
              total,
              message: update.message
            });
          }
        }
      );
      const structuredContent: Record<string, unknown> = {
        rows: result.rows,
        rowCount: result.rowCount,
        truncated: result.truncated,
        durationMs: result.durationMs
      };
      return {
        content: [],
        structuredContent
      };
    }
  );
}