Skip to content

Commit 6775bfb

Browse files
fmenezesfmenezes
committed
Merge branch 'main' into fmenezes/refactor_config
2 parents ed59914 + 5d378cc commit 6775bfb

17 files changed

+638
-577
lines changed

‎src/tools/mongodb/metadata/collectionSchema.ts‎

Lines changed: 16 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
22
import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
33
import { ToolArgs, OperationType } from "../../tool.js";
4-
import { parseSchema, SchemaField } from "mongodb-schema";
4+
import { getSimplifiedSchema } from "mongodb-schema";
55

66
export class CollectionSchemaTool extends MongoDBToolBase {
77
protected name = "collection-schema";
@@ -13,29 +13,31 @@ export class CollectionSchemaTool extends MongoDBToolBase {
1313
protected async execute({ database, collection }: ToolArgs<typeof DbOperationArgs>): Promise<CallToolResult> {
1414
const provider = await this.ensureConnected();
1515
const documents = await provider.find(database, collection, {}, { limit: 5 }).toArray();
16-
const schema = await parseSchema(documents);
16+
const schema = await getSimplifiedSchema(documents);
17+
18+
const fieldsCount = Object.entries(schema).length;
19+
if (fieldsCount === 0) {
20+
return {
21+
content: [
22+
{
23+
text: `Could not deduce the schema for "${database}.${collection}". This may be because it doesn't exist or is empty.`,
24+
type: "text",
25+
},
26+
],
27+
};
28+
}
1729

1830
return {
1931
content: [
2032
{
21-
text: `Found ${schema.fields.length} fields in the schema for \`${database}.${collection}\``,
33+
text: `Found ${fieldsCount} fields in the schema for "${database}.${collection}"`,
2234
type: "text",
2335
},
2436
{
25-
text: this.formatFieldOutput(schema.fields),
37+
text: JSON.stringify(schema),
2638
type: "text",
2739
},
2840
],
2941
};
3042
}
31-
32-
private formatFieldOutput(fields: SchemaField[]): string {
33-
let result = "| Field | Type | Confidence |\n";
34-
result += "|-------|------|-------------|\n";
35-
for (const field of fields) {
36-
const fieldType = Array.isArray(field.type) ? field.type.join(", ") : field.type;
37-
result += `| ${field.name} | \`${fieldType}\` | ${(field.probability * 100).toFixed(0)}% |\n`;
38-
}
39-
return result;
40-
}
4143
}

‎src/tools/mongodb/metadata/collectionStorageSize.ts‎

Lines changed: 41 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ import { ToolArgs, OperationType } from "../../tool.js";
44

55
export class CollectionStorageSizeTool extends MongoDBToolBase {
66
protected name = "collection-storage-size";
7-
protected description = "Gets the size of the collection in MB";
7+
protected description = "Gets the size of the collection";
88
protected argsShape = DbOperationArgs;
99

1010
protected operationType: OperationType = "metadata";
@@ -14,17 +14,55 @@ export class CollectionStorageSizeTool extends MongoDBToolBase {
1414
const [{ value }] = (await provider
1515
.aggregate(database, collection, [
1616
{ $collStats: { storageStats: {} } },
17-
{ $group: { _id: null, value: { $sum: "$storageStats.storageSize" } } },
17+
{ $group: { _id: null, value: { $sum: "$storageStats.size" } } },
1818
])
1919
.toArray()) as [{ value: number }];
2020

21+
const { units, value: scaledValue } = CollectionStorageSizeTool.getStats(value);
22+
2123
return {
2224
content: [
2325
{
24-
text: `The size of \`${database}.${collection}\` is \`${(value / 1024 / 1024).toFixed(2)} MB\``,
26+
text: `The size of "${database}.${collection}" is \`${scaledValue.toFixed(2)} ${units}\``,
2527
type: "text",
2628
},
2729
],
2830
};
2931
}
32+
33+
protected handleError(
34+
error: unknown,
35+
args: ToolArgs<typeof this.argsShape>
36+
): Promise<CallToolResult> | CallToolResult {
37+
if (error instanceof Error && "codeName" in error && error.codeName === "NamespaceNotFound") {
38+
return {
39+
content: [
40+
{
41+
text: `The size of "${args.database}.${args.collection}" cannot be determined because the collection does not exist.`,
42+
type: "text",
43+
},
44+
],
45+
};
46+
}
47+
48+
return super.handleError(error, args);
49+
}
50+
51+
private static getStats(value: number): { value: number; units: string } {
52+
const kb = 1024;
53+
const mb = kb * 1024;
54+
const gb = mb * 1024;
55+
56+
if (value > gb) {
57+
return { value: value / gb, units: "GB" };
58+
}
59+
60+
if (value > mb) {
61+
return { value: value / mb, units: "MB" };
62+
}
63+
if (value > kb) {
64+
return { value: value / kb, units: "KB" };
65+
}
66+
return { value, units: "bytes" };
67+
}
3068
}

‎src/tools/mongodb/mongodbTool.ts‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ import { ToolBase, ToolCategory } from "../tool.js";
33
import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
44
import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
55
import { ErrorCodes, MongoDBError } from "../../errors.js";
6+
import { ToolArgs } from "../tool.js";
67

78
export const DbOperationArgs = {
89
database: z.string().describe("Database name"),
@@ -24,7 +25,10 @@ export abstract class MongoDBToolBase extends ToolBase {
2425
return this.session.serviceProvider;
2526
}
2627

27-
protected handleError(error: unknown): Promise<CallToolResult> | CallToolResult {
28+
protected handleError(
29+
error: unknown,
30+
args: ToolArgs<typeof this.argsShape>
31+
): Promise<CallToolResult> | CallToolResult {
2832
if (error instanceof MongoDBError && error.code === ErrorCodes.NotConnectedToMongoDB) {
2933
return {
3034
content: [
@@ -41,7 +45,7 @@ export abstract class MongoDBToolBase extends ToolBase {
4145
};
4246
}
4347

44-
return super.handleError(error);
48+
return super.handleError(error, args);
4549
}
4650

4751
protected async connectToMongoDB(connectionString: string): Promise<void> {

‎src/tools/tool.ts‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ export abstract class ToolBase {
4747
} catch (error: unknown) {
4848
logger.error(mongoLogId(1_000_000), "tool", `Error executing ${this.name}: ${error as string}`);
4949

50-
return await this.handleError(error);
50+
return await this.handleError(error, args[0] as ToolArgs<typeof this.argsShape>);
5151
}
5252
};
5353

@@ -79,7 +79,11 @@ export abstract class ToolBase {
7979
}
8080

8181
// This method is intended to be overridden by subclasses to handle errors
82-
protected handleError(error: unknown): Promise<CallToolResult> | CallToolResult {
82+
protected handleError(
83+
error: unknown,
84+
// eslint-disable-next-line @typescript-eslint/no-unused-vars
85+
args: ToolArgs<typeof this.argsShape>
86+
): Promise<CallToolResult> | CallToolResult {
8387
return {
8488
content: [
8589
{

‎tests/integration/helpers.ts‎

Lines changed: 88 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import fs from "fs/promises";
77
import { MongoClient, ObjectId } from "mongodb";
88
import { toIncludeAllMembers } from "jest-extended";
99
import config, { UserConfig } from "../../src/config.js";
10+
import { McpError } from "@modelcontextprotocol/sdk/types.js";
1011

1112
interface ParameterInfo {
1213
name: string;
@@ -218,8 +219,91 @@ export const dbOperationParameters: ParameterInfo[] = [
218219
{ name: "collection", type: "string", description: "Collection name", required: true },
219220
];
220221

221-
export function validateParameters(tool: ToolInfo, parameters: ParameterInfo[]): void {
222-
const toolParameters = getParameters(tool);
223-
expect(toolParameters).toHaveLength(parameters.length);
224-
expect(toolParameters).toIncludeAllMembers(parameters);
222+
export const dbOperationInvalidArgTests = [{}, { database: 123 }, { foo: "bar", database: "test" }, { database: [] }];
223+
224+
export function validateToolMetadata(
225+
integration: IntegrationTest,
226+
name: string,
227+
description: string,
228+
parameters: ParameterInfo[]
229+
): void {
230+
it("should have correct metadata", async () => {
231+
const { tools } = await integration.mcpClient().listTools();
232+
const tool = tools.find((tool) => tool.name === name)!;
233+
expect(tool).toBeDefined();
234+
expect(tool.description).toBe(description);
235+
236+
const toolParameters = getParameters(tool);
237+
expect(toolParameters).toHaveLength(parameters.length);
238+
expect(toolParameters).toIncludeAllMembers(parameters);
239+
});
240+
}
241+
242+
export function validateAutoConnectBehavior(
243+
integration: IntegrationTest,
244+
name: string,
245+
validation: () => {
246+
args: { [x: string]: unknown };
247+
expectedResponse?: string;
248+
validate?: (content: unknown) => void;
249+
},
250+
beforeEachImpl?: () => Promise<void>
251+
): void {
252+
describe("when not connected", () => {
253+
if (beforeEachImpl) {
254+
beforeEach(() => beforeEachImpl());
255+
}
256+
257+
it("connects automatically if connection string is configured", async () => {
258+
config.connectionString = integration.connectionString();
259+
260+
const validationInfo = validation();
261+
262+
const response = await integration.mcpClient().callTool({
263+
name,
264+
arguments: validationInfo.args,
265+
});
266+
267+
if (validationInfo.expectedResponse) {
268+
const content = getResponseContent(response.content);
269+
expect(content).toContain(validationInfo.expectedResponse);
270+
}
271+
272+
if (validationInfo.validate) {
273+
validationInfo.validate(response.content);
274+
}
275+
});
276+
277+
it("throws an error if connection string is not configured", async () => {
278+
const response = await integration.mcpClient().callTool({
279+
name,
280+
arguments: validation().args,
281+
});
282+
const content = getResponseContent(response.content);
283+
expect(content).toContain("You need to connect to a MongoDB instance before you can access its data.");
284+
});
285+
});
286+
}
287+
288+
export function validateThrowsForInvalidArguments(
289+
integration: IntegrationTest,
290+
name: string,
291+
args: { [x: string]: unknown }[]
292+
): void {
293+
describe("with invalid arguments", () => {
294+
for (const arg of args) {
295+
it(`throws a schema error for: ${JSON.stringify(arg)}`, async () => {
296+
await integration.connectMcpClient();
297+
try {
298+
await integration.mcpClient().callTool({ name, arguments: arg });
299+
expect.fail("Expected an error to be thrown");
300+
} catch (error) {
301+
expect(error).toBeInstanceOf(McpError);
302+
const mcpError = error as McpError;
303+
expect(mcpError.code).toEqual(-32602);
304+
expect(mcpError.message).toContain(`Invalid arguments for tool ${name}`);
305+
}
306+
});
307+
}
308+
});
225309
}

‎tests/integration/tools/mongodb/create/createCollection.test.ts‎

Lines changed: 16 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -1,50 +1,24 @@
11
import {
22
getResponseContent,
3-
validateParameters,
43
dbOperationParameters,
54
setupIntegrationTest,
5+
validateToolMetadata,
6+
validateAutoConnectBehavior,
7+
validateThrowsForInvalidArguments,
8+
dbOperationInvalidArgTests,
69
} from "../../../helpers.js";
7-
import { toIncludeSameMembers } from "jest-extended";
8-
import { McpError } from "@modelcontextprotocol/sdk/types.js";
9-
import { ObjectId } from "bson";
10-
import config from "../../../../../src/config.js";
1110

1211
describe("createCollection tool", () => {
1312
const integration = setupIntegrationTest();
1413

15-
it("should have correct metadata", async () => {
16-
const { tools } = await integration.mcpClient().listTools();
17-
const listCollections = tools.find((tool) => tool.name === "create-collection")!;
18-
expect(listCollections).toBeDefined();
19-
expect(listCollections.description).toBe(
20-
"Creates a new collection in a database. If the database doesn't exist, it will be created automatically."
21-
);
14+
validateToolMetadata(
15+
integration,
16+
"create-collection",
17+
"Creates a new collection in a database. If the database doesn't exist, it will be created automatically.",
18+
dbOperationParameters
19+
);
2220

23-
validateParameters(listCollections, dbOperationParameters);
24-
});
25-
26-
describe("with invalid arguments", () => {
27-
const args = [
28-
{},
29-
{ database: 123, collection: "bar" },
30-
{ foo: "bar", database: "test", collection: "bar" },
31-
{ collection: [], database: "test" },
32-
];
33-
for (const arg of args) {
34-
it(`throws a schema error for: ${JSON.stringify(arg)}`, async () => {
35-
await integration.connectMcpClient();
36-
try {
37-
await integration.mcpClient().callTool({ name: "create-collection", arguments: arg });
38-
expect.fail("Expected an error to be thrown");
39-
} catch (error) {
40-
expect(error).toBeInstanceOf(McpError);
41-
const mcpError = error as McpError;
42-
expect(mcpError.code).toEqual(-32602);
43-
expect(mcpError.message).toContain("Invalid arguments for tool create-collection");
44-
}
45-
});
46-
}
47-
});
21+
validateThrowsForInvalidArguments(integration, "create-collection", dbOperationInvalidArgTests);
4822

4923
describe("with non-existent database", () => {
5024
it("creates a new collection", async () => {
@@ -114,25 +88,10 @@ describe("createCollection tool", () => {
11488
});
11589
});
11690

117-
describe("when not connected", () => {
118-
it("connects automatically if connection string is configured", async () => {
119-
config.connectionString = integration.connectionString();
120-
121-
const response = await integration.mcpClient().callTool({
122-
name: "create-collection",
123-
arguments: { database: integration.randomDbName(), collection: "new-collection" },
124-
});
125-
const content = getResponseContent(response.content);
126-
expect(content).toEqual(`Collection "new-collection" created in database "${integration.randomDbName()}".`);
127-
});
128-
129-
it("throws an error if connection string is not configured", async () => {
130-
const response = await integration.mcpClient().callTool({
131-
name: "create-collection",
132-
arguments: { database: integration.randomDbName(), collection: "new-collection" },
133-
});
134-
const content = getResponseContent(response.content);
135-
expect(content).toContain("You need to connect to a MongoDB instance before you can access its data.");
136-
});
91+
validateAutoConnectBehavior(integration, "create-collection", () => {
92+
return {
93+
args: { database: integration.randomDbName(), collection: "new-collection" },
94+
expectedResponse: `Collection "new-collection" created in database "${integration.randomDbName()}".`,
95+
};
13796
});
13897
});

0 commit comments

Comments
 (0)