Skip to content

Commit 2bd00cf

Browse files
blvablva
committed
Merge remote-tracking branch 'origin/main' into telemetry
2 parents a61d9b4 + 5d378cc commit 2bd00cf

17 files changed

+638
-578
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 & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { z } from "zod";
2-
import { ToolBase, ToolCategory } from "../tool.js";
2+
import { ToolArgs, ToolBase, ToolCategory } from "../tool.js";
33
import { Session } from "../../session.js";
44
import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
55
import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
@@ -30,7 +30,10 @@ export abstract class MongoDBToolBase extends ToolBase {
3030
return this.session.serviceProvider;
3131
}
3232

33-
protected handleError(error: unknown): Promise<CallToolResult> | CallToolResult {
33+
protected handleError(
34+
error: unknown,
35+
args: ToolArgs<typeof this.argsShape>
36+
): Promise<CallToolResult> | CallToolResult {
3437
if (error instanceof MongoDBError && error.code === ErrorCodes.NotConnectedToMongoDB) {
3538
return {
3639
content: [
@@ -47,7 +50,7 @@ export abstract class MongoDBToolBase extends ToolBase {
4750
};
4851
}
4952

50-
return super.handleError(error);
53+
return super.handleError(error, args);
5154
}
5255

5356
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
@@ -85,7 +85,7 @@ export abstract class ToolBase {
8585
error instanceof Error ? error : new Error(String(error))
8686
);
8787

88-
return await this.handleError(error);
88+
return await this.handleError(error, args[0] as ToolArgs<typeof this.argsShape>);
8989
}
9090
};
9191

@@ -117,7 +117,11 @@ export abstract class ToolBase {
117117
}
118118

119119
// This method is intended to be overridden by subclasses to handle errors
120-
protected handleError(error: unknown): Promise<CallToolResult> | CallToolResult {
120+
protected handleError(
121+
error: unknown,
122+
// eslint-disable-next-line @typescript-eslint/no-unused-vars
123+
args: ToolArgs<typeof this.argsShape>
124+
): Promise<CallToolResult> | CallToolResult {
121125
return {
122126
content: [
123127
{

‎tests/integration/helpers.ts‎

Lines changed: 88 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
99
import { MongoClient, ObjectId } from "mongodb";
1010
import { toIncludeAllMembers } from "jest-extended";
1111
import config from "../../src/config.js";
12+
import { McpError } from "@modelcontextprotocol/sdk/types.js";
1213

1314
interface ParameterInfo {
1415
name: string;
@@ -227,10 +228,93 @@ export const dbOperationParameters: ParameterInfo[] = [
227228
{ name: "collection", type: "string", description: "Collection name", required: true },
228229
];
229230

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

236320
export function describeAtlas(name: number | string | Function | jest.FunctionLike, fn: jest.EmptyFunction) {

‎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)