Skip to content

Commit 2b2d615

Browse files
himanshusinghshimanshusinghs
committed
chore: small improvements
1. outputStream.write moved to within the Transform.flush 2. won't send resource updated notification on export-expired event or it might trigger client to fetch expired exports. 3. added ObjectId to the file names to make them unique
1 parent d3d81d6 commit 2b2d615

File tree

4 files changed

+104
-52
lines changed

4 files changed

+104
-52
lines changed

‎src/common/exportsManager.ts‎

Lines changed: 17 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -192,7 +192,6 @@ export class ExportsManager extends EventEmitter<ExportsManagerEvents> {
192192
try {
193193
await fs.mkdir(this.exportsDirectoryPath, { recursive: true });
194194
const outputStream = createWriteStream(inProgressExport.exportPath);
195-
outputStream.write("[");
196195
await pipeline([
197196
input.stream(),
198197
this.docToEJSONStream(this.getEJSONOptionsForFormat(jsonExportFormat)),
@@ -242,20 +241,27 @@ export class ExportsManager extends EventEmitter<ExportsManagerEvents> {
242241
let docsTransformed = 0;
243242
return new Transform({
244243
objectMode: true,
245-
transform: function (chunk: unknown, encoding, callback): void {
246-
++docsTransformed;
244+
transform(chunk: unknown, encoding, callback): void {
247245
try {
248-
const doc: string = EJSON.stringify(chunk, undefined, undefined, ejsonOptions);
249-
const line = `${docsTransformed > 1 ? ",\n" : ""}${doc}`;
250-
251-
callback(null, line);
252-
} catch (err: unknown) {
246+
const doc = EJSON.stringify(chunk, undefined, undefined, ejsonOptions);
247+
if (docsTransformed === 0) {
248+
this.push("[" + doc);
249+
} else {
250+
this.push(",\n" + doc);
251+
}
252+
docsTransformed++;
253+
callback();
254+
} catch (err) {
253255
callback(err as Error);
254256
}
255257
},
256-
final: function (callback): void {
257-
this.push("]");
258-
callback(null);
258+
flush(callback): void {
259+
if (docsTransformed === 0) {
260+
this.push("[]");
261+
} else {
262+
this.push("]");
263+
}
264+
callback();
259265
},
260266
});
261267
}

‎src/resources/common/exportedData.ts‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,13 @@ export class ExportedData {
1515
private server?: Server;
1616

1717
constructor(private readonly session: Session) {
18-
const onExportChanged = (uri: string): void => {
18+
this.session.exportsManager.on("export-available", (uri: string): void => {
1919
this.server?.sendResourceListChanged();
2020
this.server?.sendResourceUpdated(uri);
21-
};
22-
this.session.exportsManager.on("export-available", onExportChanged);
23-
this.session.exportsManager.on("export-expired", onExportChanged);
21+
});
22+
this.session.exportsManager.on("export-expired", (): void => {
23+
this.server?.sendResourceListChanged();
24+
});
2425
}
2526

2627
public register(server: Server): void {
@@ -113,7 +114,7 @@ export class ExportedData {
113114
};
114115

115116
private exportNameToDescription(exportName: string): string {
116-
const match = exportName.match(/^(.+)\.(\d+)\.json$/);
117+
const match = exportName.match(/^(.+)\.(\d+)\.(.+)\.json$/);
117118
if (!match) return "Exported data for an unknown namespace.";
118119

119120
const [, namespace, timestamp] = match;

‎src/tools/mongodb/read/export.ts‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
1+
import z from "zod";
2+
import { ObjectId } from "bson";
13
import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
24
import { OperationType, ToolArgs } from "../../tool.js";
35
import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
46
import { FindArgs } from "./find.js";
57
import { jsonExportFormat } from "../../../common/exportsManager.js";
6-
import z from "zod";
78

89
export class ExportTool extends MongoDBToolBase {
910
public name = "export";
@@ -41,7 +42,11 @@ export class ExportTool extends MongoDBToolBase {
4142
promoteValues: false,
4243
bsonRegExp: true,
4344
});
44-
const exportName = `${database}.${collection}.${Date.now()}.json`;
45+
// The format is namespace.date.objectid.json
46+
// - namespace to identify which namespace the export belongs to
47+
// - date to identify when the export was generated
48+
// - objectid for uniqueness of the names
49+
const exportName = `${database}.${collection}.${Date.now()}.${new ObjectId().toString()}.json`;
4550

4651
const { exportURI, exportPath } = this.session.exportsManager.createJSONExport({
4752
input: findCursor,

‎tests/unit/common/exportsManager.test.ts‎

Lines changed: 74 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -28,14 +28,15 @@ const exportsManagerConfig: ExportsManagerConfig = {
2828

2929
function getExportNameAndPath(
3030
sessionId: string,
31-
timestamp: number
31+
timestamp: number = Date.now(),
32+
objectId: string = new ObjectId().toString()
3233
): {
3334
sessionExportsPath: string;
3435
exportName: string;
3536
exportPath: string;
3637
exportURI: string;
3738
} {
38-
const exportName = `foo.bar.${timestamp}.json`;
39+
const exportName = `foo.bar.${timestamp}.${objectId}.json`;
3940
const sessionExportsPath = path.join(exportsPath, sessionId);
4041
const exportPath = path.join(sessionExportsPath, exportName);
4142
return {
@@ -48,22 +49,21 @@ function getExportNameAndPath(
4849

4950
function createDummyFindCursor(
5051
dataArray: unknown[],
51-
chunkPushTimeoutMs?: number
52+
beforeEachChunk?: (chunkIndex: number) => void | Promise<void>
5253
): { cursor: FindCursor; cursorCloseNotification: Promise<void> } {
5354
let index = 0;
5455
const readable = new Readable({
5556
objectMode: true,
5657
async read(): Promise<void> {
57-
if (index < dataArray.length) {
58-
if (chunkPushTimeoutMs) {
59-
await timeout(chunkPushTimeoutMs);
58+
try {
59+
await beforeEachChunk?.(index);
60+
if (index < dataArray.length) {
61+
this.push(dataArray[index++]);
62+
} else {
63+
this.push(null);
6064
}
61-
this.push(dataArray[index++]);
62-
} else {
63-
if (chunkPushTimeoutMs) {
64-
await timeout(chunkPushTimeoutMs);
65-
}
66-
this.push(null);
65+
} catch (error) {
66+
this.destroy(error as Error);
6767
}
6868
},
6969
});
@@ -90,6 +90,13 @@ function createDummyFindCursor(
9090
};
9191
}
9292

93+
function createDummyFindCursorWithDelay(
94+
dataArray: unknown[],
95+
delayMs: number
96+
): { cursor: FindCursor; cursorCloseNotification: Promise<void> } {
97+
return createDummyFindCursor(dataArray, () => timeout(delayMs));
98+
}
99+
93100
async function fileExists(filePath: string): Promise<boolean> {
94101
try {
95102
await fs.access(filePath);
@@ -125,15 +132,15 @@ describe("ExportsManager unit test", () => {
125132
describe("#availableExport", () => {
126133
it("should list only the exports that are in ready state", async () => {
127134
// This export will finish in at-least 1 second
128-
const { exportName: exportName1 } = getExportNameAndPath(session.sessionId, Date.now());
135+
const { exportName: exportName1 } = getExportNameAndPath(session.sessionId);
129136
manager.createJSONExport({
130-
input: createDummyFindCursor([{ name: "Test1" }], 1000).cursor,
137+
input: createDummyFindCursorWithDelay([{ name: "Test1" }], 1000).cursor,
131138
exportName: exportName1,
132139
jsonExportFormat: "relaxed",
133140
});
134141

135142
// This export will finish way sooner than the first one
136-
const { exportName: exportName2 } = getExportNameAndPath(session.sessionId, Date.now());
143+
const { exportName: exportName2 } = getExportNameAndPath(session.sessionId);
137144
const { cursor, cursorCloseNotification } = createDummyFindCursor([{ name: "Test1" }]);
138145
manager.createJSONExport({
139146
input: cursor,
@@ -154,8 +161,8 @@ describe("ExportsManager unit test", () => {
154161
});
155162

156163
it("should throw if the resource is still being generated", async () => {
157-
const { exportName } = getExportNameAndPath(session.sessionId, Date.now());
158-
const { cursor } = createDummyFindCursor([{ name: "Test1" }], 100);
164+
const { exportName } = getExportNameAndPath(session.sessionId);
165+
const { cursor } = createDummyFindCursorWithDelay([{ name: "Test1" }], 100);
159166
manager.createJSONExport({
160167
input: cursor,
161168
exportName,
@@ -168,7 +175,7 @@ describe("ExportsManager unit test", () => {
168175
});
169176

170177
it("should return the resource content if the resource is ready to be consumed", async () => {
171-
const { exportName } = getExportNameAndPath(session.sessionId, Date.now());
178+
const { exportName } = getExportNameAndPath(session.sessionId);
172179
const { cursor, cursorCloseNotification } = createDummyFindCursor([]);
173180
manager.createJSONExport({
174181
input: cursor,
@@ -198,7 +205,7 @@ describe("ExportsManager unit test", () => {
198205
longNumber: Long.fromNumber(123456),
199206
},
200207
]));
201-
({ exportName, exportPath, exportURI } = getExportNameAndPath(session.sessionId, Date.now()));
208+
({ exportName, exportPath, exportURI } = getExportNameAndPath(session.sessionId));
202209
});
203210

204211
describe("when cursor is empty", () => {
@@ -304,31 +311,37 @@ describe("ExportsManager unit test", () => {
304311
});
305312
});
306313

307-
describe("when there is an error in export generation", () => {
314+
describe("when there is an error during stream transform", () => {
308315
it("should remove the partial export and never make it available", async () => {
309316
const emitSpy = vi.spyOn(manager, "emit");
310317
// eslint-disable-next-line @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-explicit-any
311318
(manager as any).docToEJSONStream = function (ejsonOptions: EJSONOptions | undefined): Transform {
312319
let docsTransformed = 0;
313320
return new Transform({
314321
objectMode: true,
315-
transform: function (chunk: unknown, encoding, callback): void {
316-
++docsTransformed;
322+
transform(chunk: unknown, encoding, callback): void {
317323
try {
318-
if (docsTransformed === 1) {
324+
const doc = EJSON.stringify(chunk, undefined, undefined, ejsonOptions);
325+
if (docsTransformed === 0) {
326+
this.push("[" + doc);
327+
} else if (docsTransformed === 1) {
319328
throw new Error("Could not transform the chunk!");
329+
} else {
330+
this.push(",\n" + doc);
320331
}
321-
const doc: string = EJSON.stringify(chunk, undefined, 2, ejsonOptions);
322-
const line = `${docsTransformed > 1 ? ",\n" : ""}${doc}`;
323-
324-
callback(null, line);
325-
} catch (err: unknown) {
332+
docsTransformed++;
333+
callback();
334+
} catch (err) {
326335
callback(err as Error);
327336
}
328337
},
329-
final: function (callback): void {
330-
this.push("]");
331-
callback(null);
338+
flush(this: Transform, cb): void {
339+
if (docsTransformed === 0) {
340+
this.push("[]");
341+
} else {
342+
this.push("]");
343+
}
344+
cb();
332345
},
333346
});
334347
};
@@ -348,6 +361,33 @@ describe("ExportsManager unit test", () => {
348361
expect(await fileExists(exportPath)).toEqual(false);
349362
});
350363
});
364+
365+
describe("when there is an error on read stream", () => {
366+
it("should remove the partial export and never make it available", async () => {
367+
const emitSpy = vi.spyOn(manager, "emit");
368+
// A cursor that will make the read stream fail after the first chunk
369+
const { cursor, cursorCloseNotification } = createDummyFindCursor([{ name: "Test1" }], (chunkIndex) => {
370+
if (chunkIndex > 0) {
371+
return Promise.reject(new Error("Connection timedout!"));
372+
}
373+
return Promise.resolve();
374+
});
375+
manager.createJSONExport({
376+
input: cursor,
377+
exportName,
378+
jsonExportFormat: "relaxed",
379+
});
380+
await cursorCloseNotification;
381+
382+
// Because the export was never populated in the available exports.
383+
await expect(() => manager.readExport(exportName)).rejects.toThrow(
384+
"Requested export has either expired or does not exist!"
385+
);
386+
expect(emitSpy).not.toHaveBeenCalled();
387+
expect(manager.availableExports).toEqual([]);
388+
expect(await fileExists(exportPath)).toEqual(false);
389+
});
390+
});
351391
});
352392

353393
describe("#cleanupExpiredExports", () => {
@@ -368,7 +408,7 @@ describe("ExportsManager unit test", () => {
368408
});
369409

370410
it("should not clean up in-progress exports", async () => {
371-
const { exportName } = getExportNameAndPath(session.sessionId, Date.now());
411+
const { exportName } = getExportNameAndPath(session.sessionId);
372412
const manager = ExportsManager.init(
373413
session.sessionId,
374414
{
@@ -378,7 +418,7 @@ describe("ExportsManager unit test", () => {
378418
},
379419
new CompositeLogger()
380420
);
381-
const { cursor } = createDummyFindCursor([{ name: "Test" }], 2000);
421+
const { cursor } = createDummyFindCursorWithDelay([{ name: "Test" }], 2000);
382422
manager.createJSONExport({
383423
input: cursor,
384424
exportName,
@@ -395,7 +435,7 @@ describe("ExportsManager unit test", () => {
395435
});
396436

397437
it("should cleanup expired exports", async () => {
398-
const { exportName, exportPath, exportURI } = getExportNameAndPath(session.sessionId, Date.now());
438+
const { exportName, exportPath, exportURI } = getExportNameAndPath(session.sessionId);
399439
const manager = ExportsManager.init(
400440
session.sessionId,
401441
{

0 commit comments

Comments
 (0)