fix: use union types and reuse argument definitions

Author: gagik

src/tools/mongodb/metadata/explain.ts

@@ -3,6 +3,9 @@ import { DbOperationArgs, DbOperationType, MongoDBToolBase } from "../mongodbToo
 import { ToolArgs } from "../../tool.js";
 import { z } from "zod";
 import { ExplainVerbosity, Document } from "mongodb";
+import { AggregateArgs } from "../read/aggregate.js";
+import { FindArgs } from "../read/find.js";
+import { CountArgs } from "../read/count.js";
 
 export class ExplainTool extends MongoDBToolBase {
     protected name = "explain";
@@ -11,42 +14,65 @@ export class ExplainTool extends MongoDBToolBase {
 
     protected argsShape = {
         ...DbOperationArgs,
-        method: z.enum(["aggregate", "find"]).describe("The method to run"),
-        methodArguments: z
-            .object({
-                aggregatePipeline: z
-                    .array(z.object({}).passthrough())
-                    .optional()
-                    .describe("aggregate - array of aggregation stages to execute"),
-
-                findQuery: z.object({}).passthrough().optional().describe("find - The query to run"),
-                findProjection: z.object({}).passthrough().optional().describe("find - The projection to apply"),
-            })
-            .describe("The arguments for the method"),
+        method: z
+            .array(
+                z.union([
+                    z.object({
+                        name: z.literal("aggregate"),
+                        arguments: z.object(AggregateArgs),
+                    }),
+                    z.object({
+                        name: z.literal("find"),
+                        arguments: z.object(FindArgs),
+                    }),
+                    z.object({
+                        name: z.literal("count"),
+                        arguments: z.object(CountArgs),
+                    }),
+                ])
+            )
+            .describe("The method and its arguments to run"),
     };
 
     protected operationType: DbOperationType = "metadata";
 
+    static readonly defaultVerbosity = ExplainVerbosity.queryPlanner;
+
     protected async execute({
         database,
         collection,
-        method,
-        methodArguments,
+        method: methods,
     }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
         const provider = this.ensureConnected();
+        const method = methods[0];
+
+        if (!method) {
+            throw new Error("No method provided");
+        }
 
         let result: Document;
-        switch (method) {
+        switch (method.name) {
             case "aggregate": {
-                result = await provider.aggregate(database, collection).explain();
+                const { pipeline, limit } = method.arguments;
+                result = await provider
+                    .aggregate(database, collection, pipeline)
+                    .limit(limit)
+                    .explain(ExplainTool.defaultVerbosity);
                 break;
             }
             case "find": {
-                const query = methodArguments.findQuery ?? {};
-                const projection = methodArguments.findProjection ?? {};
+                const { filter, ...rest } = method.arguments;
                 result = await provider
-                    .find(database, collection, query, { projection })
-                    .explain(ExplainVerbosity.queryPlanner);
+                    .find(database, collection, filter, { ...rest })
+                    .explain(ExplainTool.defaultVerbosity);
+                break;
+            }
+            case "count": {
+                const { query } = method.arguments;
+                // This helper doesn't have explain() command but does have the argument explain
+                result = (await provider.count(database, collection, query, {
+                    explain: ExplainTool.defaultVerbosity,
+                })) as unknown as Document;
                 break;
             }
             default:

src/tools/mongodb/read/aggregate.ts

@@ -1,26 +1,30 @@
 import { z } from "zod";
 import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
-import { DbOperationType, MongoDBToolBase } from "../mongodbTool.js";
+import { DbOperationArgs, DbOperationType, MongoDBToolBase } from "../mongodbTool.js";
 import { ToolArgs } from "../../tool.js";
 
+export const AggregateArgs = {
+    pipeline: z.array(z.object({}).passthrough()).describe("An array of aggregation stages to execute"),
+    limit: z.number().optional().default(10).describe("The maximum number of documents to return"),
+};
+
 export class AggregateTool extends MongoDBToolBase {
     protected name = "aggregate";
     protected description = "Run an aggregation against a MongoDB collection";
     protected argsShape = {
-        collection: z.string().describe("Collection name"),
-        database: z.string().describe("Database name"),
-        pipeline: z.array(z.object({}).passthrough()).describe("An array of aggregation stages to execute"),
-        limit: z.number().optional().default(10).describe("The maximum number of documents to return"),
+        ...DbOperationArgs,
+        ...AggregateArgs,
     };
     protected operationType: DbOperationType = "read";
 
     protected async execute({
         database,
         collection,
         pipeline,
+        limit,
     }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
         const provider = this.ensureConnected();
-        const documents = await provider.aggregate(database, collection, pipeline).toArray();
+        const documents = await provider.aggregate(database, collection, pipeline).limit(limit).toArray();
 
         const content: Array<{ text: string; type: "text" }> = [
             {

src/tools/mongodb/read/count.ts

@@ -3,18 +3,22 @@ import { DbOperationArgs, DbOperationType, MongoDBToolBase } from "../mongodbToo
 import { ToolArgs } from "../../tool.js";
 import { z } from "zod";
 
+export const CountArgs = {
+    query: z
+        .object({})
+        .passthrough()
+        .optional()
+        .describe(
+            "The query filter to count documents. Matches the syntax of the filter argument of db.collection.count()"
+        ),
+};
+
 export class CountTool extends MongoDBToolBase {
     protected name = "count";
     protected description = "Gets the number of documents in a MongoDB collection";
     protected argsShape = {
         ...DbOperationArgs,
-        query: z
-            .object({})
-            .passthrough()
-            .optional()
-            .describe(
-                "The query filter to count documents. Matches the syntax of the filter argument of db.collection.count()"
-            ),
+        ...CountArgs,
     };
 
     protected operationType: DbOperationType = "metadata";

src/tools/mongodb/read/find.ts

@@ -1,32 +1,33 @@
 import { z } from "zod";
 import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
-import { DbOperationType, MongoDBToolBase } from "../mongodbTool.js";
+import { DbOperationArgs, DbOperationType, MongoDBToolBase } from "../mongodbTool.js";
 import { ToolArgs } from "../../tool.js";
 import { SortDirection } from "mongodb";
 
+export const FindArgs = {
+    filter: z
+        .object({})
+        .passthrough()
+        .optional()
+        .describe("The query filter, matching the syntax of the query argument of db.collection.find()"),
+    projection: z
+        .object({})
+        .passthrough()
+        .optional()
+        .describe("The projection, matching the syntax of the projection argument of db.collection.find()"),
+    limit: z.number().optional().default(10).describe("The maximum number of documents to return"),
+    sort: z
+        .record(z.string(), z.custom<SortDirection>())
+        .optional()
+        .describe("A document, describing the sort order, matching the syntax of the sort argument of cursor.sort()"),
+};
+
 export class FindTool extends MongoDBToolBase {
     protected name = "find";
     protected description = "Run a find query against a MongoDB collection";
     protected argsShape = {
-        collection: z.string().describe("Collection name"),
-        database: z.string().describe("Database name"),
-        filter: z
-            .object({})
-            .passthrough()
-            .optional()
-            .describe("The query filter, matching the syntax of the query argument of db.collection.find()"),
-        projection: z
-            .object({})
-            .passthrough()
-            .optional()
-            .describe("The projection, matching the syntax of the projection argument of db.collection.find()"),
-        limit: z.number().optional().default(10).describe("The maximum number of documents to return"),
-        sort: z
-            .record(z.string(), z.custom<SortDirection>())
-            .optional()
-            .describe(
-                "A document, describing the sort order, matching the syntax of the sort argument of cursor.sort()"
-            ),
+        ...DbOperationArgs,
+        ...FindArgs,
     };
     protected operationType: DbOperationType = "read";