Skip to content

Commit 09a1eec

Browse files
sohankshirsagarsohankshirsagar
committed
misc fixes
1 parent 51960ac commit 09a1eec

10 files changed

Lines changed: 187 additions & 17 deletions

File tree

‎src/instrumentation/libraries/date/Instrumentation.ts‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,14 @@ export class DateInstrumentation extends TdInstrumentationBase {
147147

148148
return self._handleDateCall(args, isConstructorCall);
149149
}
150+
// Preserve the original constructor name so that libraries which validate
151+
// types by checking constructor.name (e.g. Mongoose schema type validation)
152+
// still see "Date" instead of "_TdDate". Without this, any Mongoose schema
153+
// using { type: Date } throws:
154+
// TypeError: Invalid schema configuration: `_TdDate` is not a valid type
155+
// This only changes the name metadata — all instrumentation logic is unaffected.
156+
Object.defineProperty(_TdDate, "name", { value: "Date" });
157+
150158
return _TdDate;
151159
};
152160
}

‎src/instrumentation/libraries/mongodb/Instrumentation.ts‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -458,6 +458,24 @@ export class MongodbInstrumentation extends TdInstrumentationBase {
458458
reconstructBsonValue(mockData.result, this.moduleExports),
459459
);
460460

461+
// Synchronize client-generated ObjectIds with recorded ones.
462+
// Libraries like Mongoose generate _id client-side before calling
463+
// insertOne/insertMany. During replay, the document's _id must match
464+
// the recorded value. We modify the ObjectId buffer IN-PLACE because
465+
// BSON's ObjectId shares the underlying Uint8Array by reference when
466+
// cloned (e.g., Mongoose's toObject() does new ObjectId(obj.id)),
467+
// so in-place mutation propagates back to the caller's model instance.
468+
if (methodName === "insertOne" && result?.insertedId != null && args[0]?._id?.id) {
469+
args[0]._id.id.set(result.insertedId.id);
470+
} else if (methodName === "insertMany" && result?.insertedIds && Array.isArray(args[0])) {
471+
for (const [index, id] of Object.entries(result.insertedIds)) {
472+
const idx = Number(index);
473+
if (args[0]?.[idx]?._id?.id && (id as any)?.id) {
474+
args[0][idx]._id.id.set((id as any).id);
475+
}
476+
}
477+
}
478+
461479
SpanUtils.endSpan(spanInfo.span, { code: SpanStatusCode.OK });
462480
return result;
463481
} catch (error: any) {

‎src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/package.json‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,8 @@
1111
},
1212
"dependencies": {
1313
"@use-tusk/drift-node-sdk": "file:/sdk",
14-
"mongodb": "^6.0.0"
14+
"mongodb": "^6.0.0",
15+
"mongoose": "^8.0.0"
1516
},
1617
"devDependencies": {
1718
"@types/node": "^20.0.0",

‎src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/src/index.ts‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import { TuskDrift } from "./tdInit";
22
import http from "http";
33
import { MongoClient, Db, Collection } from "mongodb";
4+
import mongoose from "mongoose";
45

56
const PORT = process.env.PORT || 3000;
67

@@ -43,6 +44,48 @@ async function initializeDatabase() {
4344
console.log("Database initialized successfully");
4445
}
4546

47+
// --- Mongoose Setup ---
48+
const authorSchema = new mongoose.Schema({
49+
name: { type: String, required: true },
50+
email: { type: String, required: true },
51+
bio: String,
52+
createdAt: { type: Date, default: Date.now },
53+
});
54+
55+
const postSchema = new mongoose.Schema({
56+
title: { type: String, required: true },
57+
body: { type: String, required: true },
58+
author: { type: mongoose.Schema.Types.ObjectId, ref: "Author" },
59+
tags: [String],
60+
views: { type: Number, default: 0 },
61+
published: { type: Boolean, default: false },
62+
createdAt: { type: Number, default: () => Date.now() },
63+
});
64+
65+
const Author = mongoose.model("Author", authorSchema);
66+
const Post = mongoose.model("Post", postSchema);
67+
68+
async function initializeMongoose() {
69+
await mongoose.connect(`mongodb://${mongoHost}:${mongoPort}/${mongoDb}`);
70+
console.log("Mongoose connected successfully");
71+
72+
// Seed mongoose data
73+
await Author.deleteMany({});
74+
await Post.deleteMany({});
75+
76+
const author1 = await Author.create({ name: "Alice Writer", email: "alice@example.com", bio: "Tech blogger" });
77+
const author2 = await Author.create({ name: "Bob Author", email: "bob_author@example.com", bio: "Fiction writer" });
78+
79+
await Post.create([
80+
{ title: "First Post", body: "Hello world content", author: author1._id, tags: ["intro", "tech"], views: 100, published: true },
81+
{ title: "Second Post", body: "Advanced topics", author: author1._id, tags: ["tech", "advanced"], views: 250, published: true },
82+
{ title: "Draft Post", body: "Work in progress", author: author2._id, tags: ["draft"], views: 0, published: false },
83+
{ title: "Fiction Story", body: "Once upon a time", author: author2._id, tags: ["fiction", "story"], views: 50, published: true },
84+
]);
85+
86+
console.log("Mongoose data seeded");
87+
}
88+
4689
function sendJson(res: http.ServerResponse, statusCode: number, data: any) {
4790
res.writeHead(statusCode, { "Content-Type": "application/json" });
4891
res.end(JSON.stringify(data, (_, v) => (typeof v === "bigint" ? v.toString() : v)));
@@ -429,6 +472,34 @@ const server = http.createServer(async (req, res) => {
429472
return;
430473
}
431474

475+
// --- Mongoose: create ---
476+
if (url === "/test/mongoose-create" && method === "GET") {
477+
const tmpAuthor = await Author.create({ name: "Temp Author", email: "temp@example.com", bio: "Temporary" });
478+
const result = { id: tmpAuthor._id, name: tmpAuthor.name, email: tmpAuthor.email };
479+
await Author.deleteOne({ _id: tmpAuthor._id });
480+
sendJson(res, 200, { success: true, data: result });
481+
return;
482+
}
483+
484+
// --- Mongoose: create multiple docs (insertMany path) ---
485+
if (url === "/test/mongoose-create-many" && method === "GET") {
486+
const created = await Post.create([
487+
{ title: "TempMulti1", body: "Multi body 1", tags: ["temp-multi"], views: 0, published: false },
488+
{ title: "TempMulti2", body: "Multi body 2", tags: ["temp-multi"], views: 0, published: false },
489+
]);
490+
const result = created.map(d => ({ title: d.title, id: d._id }));
491+
await Post.deleteMany({ tags: "temp-multi" });
492+
sendJson(res, 200, { success: true, data: result, count: result.length });
493+
return;
494+
}
495+
496+
// --- Raw MongoDB: cursor map transform ---
497+
if (url === "/test/cursor-map" && method === "GET") {
498+
const mapped = await largeData.find({ category: "even" }).map(doc => ({ value: doc.value, doubled: doc.index * 2 })).toArray();
499+
sendJson(res, 200, { success: true, data: mapped });
500+
return;
501+
}
502+
432503
// 404 for unknown routes
433504
sendJson(res, 404, { error: "Not found" });
434505
} catch (error) {
@@ -443,6 +514,7 @@ const server = http.createServer(async (req, res) => {
443514
// Initialize database first, then start server
444515
async function main() {
445516
await initializeDatabase();
517+
await initializeMongoose();
446518
server.listen(PORT, () => {
447519
TuskDrift.markAppAsReady();
448520
console.log(`MongoDB integration test server running on port ${PORT}`);
@@ -459,6 +531,7 @@ main().catch((error) => {
459531
async function shutdown() {
460532
console.log("Shutting down gracefully...");
461533
try {
534+
await mongoose.disconnect();
462535
await client.close();
463536
} catch (error) {
464537
console.error("Error during shutdown:", error);

‎src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/src/test_requests.mjs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,5 +28,8 @@ await makeRequest('GET', '/test/list-collections');
2828
await makeRequest('GET', '/test/transaction');
2929
await makeRequest('GET', '/test/ordered-bulk');
3030
await makeRequest('GET', '/test/unordered-bulk');
31+
await makeRequest('GET', '/test/mongoose-create');
32+
await makeRequest('GET', '/test/mongoose-create-many');
33+
await makeRequest('GET', '/test/cursor-map');
3134

3235
printRequestSummary();

‎src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/package.json‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,8 @@
1111
},
1212
"dependencies": {
1313
"@use-tusk/drift-node-sdk": "file:/sdk",
14-
"mongodb": "^6.0.0"
14+
"mongodb": "^6.0.0",
15+
"mongoose": "^8.0.0"
1516
},
1617
"devDependencies": {
1718
"@types/node": "^20.0.0",

‎src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/src/index.ts‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import { TuskDrift } from "./tdInit.js";
22
import http from "http";
33
import { MongoClient, Db, Collection } from "mongodb";
4+
import mongoose from "mongoose";
45

56
const PORT = process.env.PORT || 3000;
67

@@ -43,6 +44,48 @@ async function initializeDatabase() {
4344
console.log("Database initialized successfully");
4445
}
4546

47+
// --- Mongoose Setup ---
48+
const authorSchema = new mongoose.Schema({
49+
name: { type: String, required: true },
50+
email: { type: String, required: true },
51+
bio: String,
52+
createdAt: { type: Date, default: Date.now },
53+
});
54+
55+
const postSchema = new mongoose.Schema({
56+
title: { type: String, required: true },
57+
body: { type: String, required: true },
58+
author: { type: mongoose.Schema.Types.ObjectId, ref: "Author" },
59+
tags: [String],
60+
views: { type: Number, default: 0 },
61+
published: { type: Boolean, default: false },
62+
createdAt: { type: Number, default: () => Date.now() },
63+
});
64+
65+
const Author = mongoose.model("Author", authorSchema);
66+
const Post = mongoose.model("Post", postSchema);
67+
68+
async function initializeMongoose() {
69+
await mongoose.connect(`mongodb://${mongoHost}:${mongoPort}/${mongoDb}`);
70+
console.log("Mongoose connected successfully");
71+
72+
// Seed mongoose data
73+
await Author.deleteMany({});
74+
await Post.deleteMany({});
75+
76+
const author1 = await Author.create({ name: "Alice Writer", email: "alice@example.com", bio: "Tech blogger" });
77+
const author2 = await Author.create({ name: "Bob Author", email: "bob_author@example.com", bio: "Fiction writer" });
78+
79+
await Post.create([
80+
{ title: "First Post", body: "Hello world content", author: author1._id, tags: ["intro", "tech"], views: 100, published: true },
81+
{ title: "Second Post", body: "Advanced topics", author: author1._id, tags: ["tech", "advanced"], views: 250, published: true },
82+
{ title: "Draft Post", body: "Work in progress", author: author2._id, tags: ["draft"], views: 0, published: false },
83+
{ title: "Fiction Story", body: "Once upon a time", author: author2._id, tags: ["fiction", "story"], views: 50, published: true },
84+
]);
85+
86+
console.log("Mongoose data seeded");
87+
}
88+
4689
function sendJson(res: http.ServerResponse, statusCode: number, data: any) {
4790
res.writeHead(statusCode, { "Content-Type": "application/json" });
4891
res.end(JSON.stringify(data, (_, v) => (typeof v === "bigint" ? v.toString() : v)));
@@ -429,6 +472,34 @@ const server = http.createServer(async (req, res) => {
429472
return;
430473
}
431474

475+
// --- Mongoose: create ---
476+
if (url === "/test/mongoose-create" && method === "GET") {
477+
const tmpAuthor = await Author.create({ name: "Temp Author", email: "temp@example.com", bio: "Temporary" });
478+
const result = { id: tmpAuthor._id, name: tmpAuthor.name, email: tmpAuthor.email };
479+
await Author.deleteOne({ _id: tmpAuthor._id });
480+
sendJson(res, 200, { success: true, data: result });
481+
return;
482+
}
483+
484+
// --- Mongoose: create multiple docs (insertMany path) ---
485+
if (url === "/test/mongoose-create-many" && method === "GET") {
486+
const created = await Post.create([
487+
{ title: "TempMulti1", body: "Multi body 1", tags: ["temp-multi"], views: 0, published: false },
488+
{ title: "TempMulti2", body: "Multi body 2", tags: ["temp-multi"], views: 0, published: false },
489+
]);
490+
const result = created.map(d => ({ title: d.title, id: d._id }));
491+
await Post.deleteMany({ tags: "temp-multi" });
492+
sendJson(res, 200, { success: true, data: result, count: result.length });
493+
return;
494+
}
495+
496+
// --- Raw MongoDB: cursor map transform ---
497+
if (url === "/test/cursor-map" && method === "GET") {
498+
const mapped = await largeData.find({ category: "even" }).map(doc => ({ value: doc.value, doubled: doc.index * 2 })).toArray();
499+
sendJson(res, 200, { success: true, data: mapped });
500+
return;
501+
}
502+
432503
// 404 for unknown routes
433504
sendJson(res, 404, { error: "Not found" });
434505
} catch (error) {
@@ -443,6 +514,7 @@ const server = http.createServer(async (req, res) => {
443514
// Initialize database first, then start server
444515
async function main() {
445516
await initializeDatabase();
517+
await initializeMongoose();
446518
server.listen(PORT, () => {
447519
TuskDrift.markAppAsReady();
448520
console.log(`MongoDB integration test server running on port ${PORT}`);
@@ -459,6 +531,7 @@ main().catch((error) => {
459531
async function shutdown() {
460532
console.log("Shutting down gracefully...");
461533
try {
534+
await mongoose.disconnect();
462535
await client.close();
463536
} catch (error) {
464537
console.error("Error during shutdown:", error);

‎src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/src/test_requests.mjs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,5 +28,8 @@ await makeRequest('GET', '/test/list-collections');
2828
await makeRequest('GET', '/test/transaction');
2929
await makeRequest('GET', '/test/ordered-bulk');
3030
await makeRequest('GET', '/test/unordered-bulk');
31+
await makeRequest('GET', '/test/mongoose-create');
32+
await makeRequest('GET', '/test/mongoose-create-many');
33+
await makeRequest('GET', '/test/cursor-map');
3134

3235
printRequestSummary();

‎src/instrumentation/libraries/mongodb/e2e-tests/run-all.sh‎

100644100755
File mode changed.

‎src/instrumentation/libraries/mongodb/mocks/FakeCursor.ts‎

Lines changed: 5 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -139,21 +139,11 @@ export class TdFakeFindCursor {
139139
return this;
140140
}
141141

142-
map(transform: (doc: any) => any): this {
143-
if (this._mockLoaded) {
144-
this.documents = this.documents.map(transform);
145-
} else {
146-
// Chain the transform for application after mock data loads
147-
const originalLoader = this._mockDataLoader;
148-
if (originalLoader) {
149-
this._mockDataLoader = async () => {
150-
const docs = await originalLoader();
151-
return docs.map(transform);
152-
};
153-
// Reset cached promise since loader changed
154-
this._mockLoadPromise = null;
155-
}
156-
}
142+
map(): this {
143+
// No-op: recorded mock data already contains post-transform results.
144+
// The real cursor's terminal methods (toArray, next, etc.) apply the
145+
// map transform before the instrumentation captures the output, so
146+
// re-applying the transform here would corrupt the data.
157147
return this;
158148
}
159149

0 commit comments

Comments
 (0)