diff --git a/.github/workflows/e2e.yml b/.github/workflows/e2e.yml index d1854ac0..d8c528d5 100644 --- a/.github/workflows/e2e.yml +++ b/.github/workflows/e2e.yml @@ -21,7 +21,20 @@ jobs: matrix: library: # Exclude upstash-redis-js for now because takes ~15 min. To re-enable when someone needs it. - [fetch, firestore, grpc, http, ioredis, mysql, mysql2, nextjs, pg, postgres, prisma] + [ + fetch, + firestore, + grpc, + http, + ioredis, + mysql, + mysql2, + nextjs, + pg, + postgres, + prisma, + mongodb, + ] steps: - name: Checkout uses: actions/checkout@v4 diff --git a/README.md b/README.md index 96e0fa39..7f48d9fc 100644 --- a/README.md +++ b/README.md @@ -48,6 +48,7 @@ Tusk Drift currently supports the following packages and versions: - **Firestore**: `@google-cloud/firestore@7.x-8.x` - **Postgres**: `postgres@3.x` - **MySQL**: `mysql2@3.x`, `mysql@2.x` +- **MongoDB**: `mongodb@5.x-7.x` - **IORedis**: `ioredis@4.x-5.x` - **Upstash Redis**: `@upstash/redis@1.x` - **GraphQL**: `graphql@15.x-16.x` diff --git a/src/core/TuskDrift.ts b/src/core/TuskDrift.ts index a9084810..85272463 100644 --- a/src/core/TuskDrift.ts +++ b/src/core/TuskDrift.ts @@ -20,6 +20,7 @@ import { NextjsInstrumentation, PrismaInstrumentation, MysqlInstrumentation, + MongodbInstrumentation, } from "../instrumentation/libraries"; import { TdSpanExporter } from "./tracing/TdSpanExporter"; import { trace, Tracer, SpanKind, SpanStatusCode } from "@opentelemetry/api"; @@ -325,6 +326,11 @@ export class TuskDriftCore { enabled: true, mode: this.mode, }); + + new MongodbInstrumentation({ + enabled: true, + mode: this.mode, + }); } private initializeTracing({ baseDirectory }: { baseDirectory: string }): void { diff --git a/src/instrumentation/libraries/date/Instrumentation.ts b/src/instrumentation/libraries/date/Instrumentation.ts index 3d5bc3d0..8ac03458 100644 --- a/src/instrumentation/libraries/date/Instrumentation.ts +++ b/src/instrumentation/libraries/date/Instrumentation.ts @@ -147,6 +147,14 @@ export class DateInstrumentation extends TdInstrumentationBase { return self._handleDateCall(args, isConstructorCall); } + // Preserve the original constructor name so that libraries which validate + // types by checking constructor.name (e.g. Mongoose schema type validation) + // still see "Date" instead of "_TdDate". Without this, any Mongoose schema + // using { type: Date } throws: + // TypeError: Invalid schema configuration: `_TdDate` is not a valid type + // This only changes the name metadata — all instrumentation logic is unaffected. + Object.defineProperty(_TdDate, "name", { value: "Date" }); + return _TdDate; }; } diff --git a/src/instrumentation/libraries/index.ts b/src/instrumentation/libraries/index.ts index 3c5dc152..b0af1a31 100644 --- a/src/instrumentation/libraries/index.ts +++ b/src/instrumentation/libraries/index.ts @@ -15,3 +15,4 @@ export * from "./grpc"; export * from "./firestore"; export * from "./nextjs"; export * from "./prisma"; +export * from "./mongodb"; diff --git a/src/instrumentation/libraries/mongodb/Instrumentation.ts b/src/instrumentation/libraries/mongodb/Instrumentation.ts new file mode 100644 index 00000000..fa3d1a3a --- /dev/null +++ b/src/instrumentation/libraries/mongodb/Instrumentation.ts @@ -0,0 +1,2818 @@ +import { TdInstrumentationBase } from "../../core/baseClasses/TdInstrumentationBase"; +import { TdInstrumentationNodeModule } from "../../core/baseClasses/TdInstrumentationNodeModule"; +import { TdInstrumentationNodeModuleFile } from "../../core/baseClasses/TdInstrumentationNodeModuleFile"; +import { SpanUtils, SpanInfo } from "../../../core/tracing/SpanUtils"; +import { SpanKind, SpanStatusCode, context, Context } from "@opentelemetry/api"; +import { TuskDriftCore, TuskDriftMode } from "../../../core/TuskDrift"; +import { captureStackTrace, wrap } from "../../core/utils"; +import { findMockResponseAsync } from "../../core/utils/mockResponseUtils"; +import { handleRecordMode, handleReplayMode } from "../../core/utils/modeUtils"; +import { PackageType } from "@use-tusk/drift-schemas/core/span"; +import { + MongodbModuleExports, + MongodbInstrumentationConfig, + MongodbCommandInputValue, +} from "./types"; +import { logger, isEsm } from "../../../core/utils"; +import { + createMockInputValue, + createSpanInputValue, +} from "../../../core/utils/dataNormalizationUtils"; +import { TdSpanAttributes } from "../../../core/types"; +import { ConnectionHandler } from "./handlers/ConnectionHandler"; +import { TdFakeFindCursor, TdFakeAggregationCursor, TdFakeChangeStream } from "./mocks/FakeCursor"; +import { TdFakeTopology } from "./mocks/FakeTopology"; +import { + sanitizeBsonValue, + reconstructBsonValue, + addOutputAttributesToSpan, + sanitizeOptions, + wrapCursorOutput, + unwrapCursorOutput, + wrapDirectOutput, + unwrapDirectOutput, +} from "./utils/bsonConversion"; + +/** + * Collection methods that return Promises directly (not cursors). + * Each is wrapped with record/replay instrumentation. + */ +const COLLECTION_METHODS_TO_WRAP = [ + "findOne", + "insertOne", + "insertMany", + "updateOne", + "updateMany", + "deleteOne", + "deleteMany", + "replaceOne", + "findOneAndUpdate", + "findOneAndDelete", + "findOneAndReplace", + "countDocuments", + "estimatedDocumentCount", + "distinct", + "bulkWrite", + // Collection index operations + "createIndex", + "createIndexes", + "dropIndex", + "dropIndexes", + "indexes", +]; + +/** + * Cursor-returning methods on Collection.prototype. + * These require special handling because they return cursors synchronously + * rather than Promises. + */ +const CURSOR_METHODS_TO_WRAP = ["find", "aggregate", "listIndexes"] as const; + +/** + * Db.prototype methods that return Promises directly (not cursors). + */ +const DB_METHODS_TO_WRAP = [ + "command", + "createCollection", + "dropCollection", + "dropDatabase", +] as const; + +/** + * Db.prototype methods that return cursors. + */ +const DB_CURSOR_METHODS_TO_WRAP = ["listCollections", "aggregate"] as const; + +export class MongodbInstrumentation extends TdInstrumentationBase { + private readonly INSTRUMENTATION_NAME = "MongodbInstrumentation"; + private mode: TuskDriftMode; + private tuskDrift: TuskDriftCore; + private connectionHandler: ConnectionHandler; + private moduleExports: any; + + constructor(config: MongodbInstrumentationConfig = {}) { + super("mongodb", config); + this.mode = config.mode || TuskDriftMode.DISABLED; + this.tuskDrift = TuskDriftCore.getInstance(); + this.connectionHandler = new ConnectionHandler(this.mode, this.INSTRUMENTATION_NAME, () => + this.tuskDrift.isAppReady(), + ); + } + + init(): TdInstrumentationNodeModule[] { + return [ + new TdInstrumentationNodeModule({ + name: "mongodb", + supportedVersions: ["5.*", "6.*", "7.*"], + patch: (moduleExports: MongodbModuleExports) => this._patchMongodbModule(moduleExports), + files: [ + new TdInstrumentationNodeModuleFile({ + name: "mongodb/lib/sessions.js", + supportedVersions: ["5.*", "6.*", "7.*"], + patch: (moduleExports: any) => this._patchSessionModule(moduleExports), + }), + // Ordered bulk operations + new TdInstrumentationNodeModuleFile({ + name: "mongodb/lib/bulk/ordered.js", + supportedVersions: ["5.*", "6.*", "7.*"], + patch: (moduleExports: any) => this._patchOrderedBulkModule(moduleExports), + }), + // Unordered bulk operations + new TdInstrumentationNodeModuleFile({ + name: "mongodb/lib/bulk/unordered.js", + supportedVersions: ["5.*", "6.*", "7.*"], + patch: (moduleExports: any) => this._patchUnorderedBulkModule(moduleExports), + }), + ], + }), + ]; + } + + /** + * Patch the mongodb module exports. + * Wraps MongoClient, Collection, Db, and cursor prototypes to intercept + * all database operations for record/replay. + */ + private _patchMongodbModule(mongodbModule: MongodbModuleExports): MongodbModuleExports { + logger.debug(`[${this.INSTRUMENTATION_NAME}] Patching MongoDB module in ${this.mode} mode`); + + if (this.isModulePatched(mongodbModule)) { + logger.debug(`[${this.INSTRUMENTATION_NAME}] MongoDB module already patched, skipping`); + return mongodbModule; + } + + // Resolve actual exports (handle ESM vs CJS) + const actualExports = isEsm(mongodbModule) ? mongodbModule.default : mongodbModule; + + // Store module exports for BSON reconstruction during replay + this.moduleExports = actualExports; + + if (!actualExports || !actualExports.MongoClient) { + logger.error( + `[${this.INSTRUMENTATION_NAME}] MongoClient not found in module exports, cannot patch`, + ); + return mongodbModule; + } + + // Patch MongoClient connection lifecycle methods + this._wrap(actualExports.MongoClient.prototype, "connect", (original: any) => { + const self = this; + return function (this: any, ...args: any[]) { + return self.connectionHandler.handleConnect(original, this, args); + }; + }); + + this._wrap(actualExports.MongoClient.prototype, "close", (original: any) => { + const self = this; + return function (this: any, ...args: any[]) { + return self.connectionHandler.handleClose(original, this, args); + }; + }); + + this._wrap(actualExports.MongoClient.prototype, "db", (original: any) => { + const self = this; + return function (this: any, ...args: any[]) { + return self.connectionHandler.handleDb(original, this, args); + }; + }); + + // Patch Collection.prototype CRUD methods + try { + this._patchCollectionMethods(actualExports); + } catch (error) { + logger.error( + `[${this.INSTRUMENTATION_NAME}] Error patching Collection methods, skipping:`, + error, + ); + } + + // Patch cursor-returning Collection methods (find, aggregate, listIndexes) + try { + this._patchCursorReturningMethods(actualExports); + } catch (error) { + logger.error( + `[${this.INSTRUMENTATION_NAME}] Error patching cursor-returning methods, skipping:`, + error, + ); + } + + // Patch Collection.prototype.initializeOrderedBulkOp / initializeUnorderedBulkOp + // In replay mode, inject FakeTopology before calling original to prevent + // "MongoClient must be connected" error from BulkOperationBase constructor. + try { + this._patchBulkOpInitMethods(actualExports); + } catch (error) { + logger.error( + `[${this.INSTRUMENTATION_NAME}] Error patching bulk op init methods, skipping:`, + error, + ); + } + + // Patch Db.prototype methods + try { + this._patchDbMethods(actualExports); + this._patchDbCursorMethods(actualExports); + } catch (error) { + logger.error(`[${this.INSTRUMENTATION_NAME}] Error patching Db methods, skipping:`, error); + } + + // Patch MongoClient.prototype.startSession + try { + if (typeof actualExports.MongoClient.prototype.startSession === "function") { + this._wrap( + actualExports.MongoClient.prototype, + "startSession", + this._getStartSessionWrapper(), + ); + logger.debug(`[${this.INSTRUMENTATION_NAME}] Wrapped MongoClient.prototype.startSession`); + } + } catch (error) { + logger.error(`[${this.INSTRUMENTATION_NAME}] Error patching startSession, skipping:`, error); + } + + // Wrap watch() methods as passthrough (RECORD) or no-op (REPLAY) + try { + this._patchWatchMethods(actualExports); + } catch (error) { + logger.error(`[${this.INSTRUMENTATION_NAME}] Error patching watch methods, skipping:`, error); + } + + this.markModuleAsPatched(mongodbModule); + logger.debug(`[${this.INSTRUMENTATION_NAME}] MongoDB module patching complete`); + + return mongodbModule; + } + + // --------------------------------------------------------------------------- + // Collection-level CRUD method patching + // --------------------------------------------------------------------------- + + /** + * Patch all Collection.prototype CRUD methods that return Promises. + * Guards each method with typeof check for version compatibility. + */ + private _patchCollectionMethods(actualExports: any): void { + const Collection = actualExports.Collection; + if (!Collection || !Collection.prototype) { + logger.warn( + `[${this.INSTRUMENTATION_NAME}] Collection not found in module exports, skipping collection method patching`, + ); + return; + } + + for (const methodName of COLLECTION_METHODS_TO_WRAP) { + if (typeof Collection.prototype[methodName] === "function") { + this._wrap(Collection.prototype, methodName, this._getCollectionMethodWrapper(methodName)); + logger.debug(`[${this.INSTRUMENTATION_NAME}] Wrapped Collection.prototype.${methodName}`); + } else { + logger.debug( + `[${this.INSTRUMENTATION_NAME}] Collection.prototype.${methodName} not found (may not exist in this version), skipping`, + ); + } + } + } + + /** + * Returns a wrapper function for a given Collection method. + * Handles RECORD, REPLAY, and DISABLED modes. + */ + private _getCollectionMethodWrapper(methodName: string): (original: Function) => Function { + const self = this; + const spanName = `mongodb.${methodName}`; + const submodule = `collection-${methodName}`; + + return (original: Function) => { + return function (this: any, ...args: any[]) { + const collectionName = this?.s?.namespace?.collection; + const databaseName = this?.s?.namespace?.db; + const inputValue = self._extractCollectionInput( + methodName, + collectionName, + databaseName, + args, + ); + + if (self.mode === TuskDriftMode.DISABLED) { + return original.apply(this, args); + } + + if (self.mode === TuskDriftMode.RECORD) { + return self._handleRecordCollectionMethod( + original, + this, + args, + inputValue, + spanName, + submodule, + ); + } + + if (self.mode === TuskDriftMode.REPLAY) { + return self._handleReplayCollectionMethod( + original, + this, + args, + inputValue, + spanName, + submodule, + methodName, + ); + } + + return original.apply(this, args); + }; + }; + } + + /** + * Handle RECORD mode for a collection method. + * Calls the original method, wraps the promise to capture output, creates a span. + */ + private _handleRecordCollectionMethod( + original: Function, + thisArg: any, + args: any[], + inputValue: MongodbCommandInputValue, + spanName: string, + submodule: string, + ): any { + return handleRecordMode({ + originalFunctionCall: () => original.apply(thisArg, args), + recordModeHandler: ({ isPreAppStart }) => { + return SpanUtils.createAndExecuteSpan( + this.mode, + () => original.apply(thisArg, args), + { + name: spanName, + kind: SpanKind.CLIENT, + submodule, + packageType: PackageType.MONGODB, + packageName: "mongodb", + instrumentationName: this.INSTRUMENTATION_NAME, + inputValue, + isPreAppStart, + stopRecordingChildSpans: true, + }, + (spanInfo: SpanInfo) => { + const resultPromise = original.apply(thisArg, args) as Promise; + + return resultPromise + .then((result: any) => { + try { + addOutputAttributesToSpan(spanInfo, wrapDirectOutput(result)); + SpanUtils.endSpan(spanInfo.span, { + code: SpanStatusCode.OK, + }); + } catch (error) { + logger.error( + `[${this.INSTRUMENTATION_NAME}] Error adding span attributes for ${spanName}:`, + error, + ); + } + return result; + }) + .catch((error: any) => { + try { + SpanUtils.addSpanAttributes(spanInfo.span, { + outputValue: { + error: error?.message || "Unknown error", + }, + }); + SpanUtils.endSpan(spanInfo.span, { + code: SpanStatusCode.ERROR, + message: error?.message || "Operation failed", + }); + } catch (spanError) { + logger.error( + `[${this.INSTRUMENTATION_NAME}] Error recording span for ${spanName} error:`, + spanError, + ); + } + throw error; + }); + }, + ); + }, + spanKind: SpanKind.CLIENT, + }); + } + + /** + * Handle REPLAY mode for a collection method. + * Looks up mock data, reconstructs BSON types, returns mocked result. + */ + private _handleReplayCollectionMethod( + original: Function, + thisArg: any, + args: any[], + inputValue: MongodbCommandInputValue, + spanName: string, + submodule: string, + methodName: string, + ): any { + const stackTrace = captureStackTrace(["MongodbInstrumentation"]); + return handleReplayMode({ + noOpRequestHandler: () => { + return Promise.resolve(this._getNoOpResult(methodName)); + }, + isServerRequest: false, + replayModeHandler: () => { + return SpanUtils.createAndExecuteSpan( + this.mode, + () => original.apply(thisArg, args), + { + name: spanName, + kind: SpanKind.CLIENT, + submodule, + packageType: PackageType.MONGODB, + packageName: "mongodb", + instrumentationName: this.INSTRUMENTATION_NAME, + inputValue, + isPreAppStart: !this.tuskDrift.isAppReady(), + stopRecordingChildSpans: true, + }, + async (spanInfo: SpanInfo) => { + try { + const mockData = await this._findMockData({ + spanInfo, + name: spanName, + inputValue, + submoduleName: submodule, + stackTrace, + }); + + if (!mockData) { + const errorMsg = `[${this.INSTRUMENTATION_NAME}] No matching mock found for ${spanName} (collection: ${inputValue.collection})`; + logger.warn(errorMsg); + throw new Error(errorMsg); + } + + const result = unwrapDirectOutput( + reconstructBsonValue(mockData.result, this.moduleExports), + ); + + // Synchronize client-generated ObjectIds with recorded ones. + // Libraries like Mongoose generate _id client-side before calling + // insertOne/insertMany. During replay, the document's _id must match + // the recorded value. We modify the ObjectId buffer IN-PLACE because + // BSON's ObjectId shares the underlying Uint8Array by reference when + // cloned (e.g., Mongoose's toObject() does new ObjectId(obj.id)), + // so in-place mutation propagates back to the caller's model instance. + if (methodName === "insertOne" && result?.insertedId != null && args[0]?._id?.id) { + args[0]._id.id.set(result.insertedId.id); + } else if ( + methodName === "insertMany" && + result?.insertedIds && + Array.isArray(args[0]) + ) { + for (const [index, id] of Object.entries(result.insertedIds)) { + const idx = Number(index); + if (args[0]?.[idx]?._id?.id && (id as any)?.id) { + args[0][idx]._id.id.set((id as any).id); + } + } + } + + SpanUtils.endSpan(spanInfo.span, { code: SpanStatusCode.OK }); + return result; + } catch (error: any) { + SpanUtils.endSpan(spanInfo.span, { + code: SpanStatusCode.ERROR, + message: error?.message || "Replay failed", + }); + throw error; + } + }, + ); + }, + }); + } + + // --------------------------------------------------------------------------- + // Cursor-returning method patching (find, aggregate) + // --------------------------------------------------------------------------- + + /** + * Patch Collection.prototype methods that return cursors (find, aggregate). + * Unlike CRUD methods, these are synchronous and return cursor objects. + * The actual query doesn't execute until a terminal method is called. + */ + private _patchCursorReturningMethods(actualExports: any): void { + const Collection = actualExports.Collection; + if (!Collection || !Collection.prototype) { + logger.warn( + `[${this.INSTRUMENTATION_NAME}] Collection not found, skipping cursor method patching`, + ); + return; + } + + for (const methodName of CURSOR_METHODS_TO_WRAP) { + if (typeof Collection.prototype[methodName] === "function") { + this._wrap(Collection.prototype, methodName, this._getCursorMethodWrapper(methodName)); + logger.debug(`[${this.INSTRUMENTATION_NAME}] Wrapped Collection.prototype.${methodName}`); + } else { + logger.debug( + `[${this.INSTRUMENTATION_NAME}] Collection.prototype.${methodName} not found, skipping`, + ); + } + } + } + + /** + * Returns a wrapper function for a cursor-returning Collection method. + * + * - DISABLED: passthrough + * - RECORD: call original (get real cursor), wrap terminal methods on the instance + * - REPLAY: return a lazy fake cursor that loads mock data on first terminal call + */ + private _getCursorMethodWrapper(methodName: string): (original: Function) => Function { + const self = this; + const spanName = `mongodb.${methodName}`; + const submodule = `collection-${methodName}`; + + return (original: Function) => { + return function (this: any, ...args: any[]) { + const collectionName = this?.s?.namespace?.collection; + const databaseName = this?.s?.namespace?.db; + const inputValue = self._extractCursorInput(methodName, collectionName, databaseName, args); + + if (self.mode === TuskDriftMode.DISABLED) { + return original.apply(this, args); + } + + // Capture context at cursor creation time (before builder chaining) + const creationContext = context.active(); + + if (self.mode === TuskDriftMode.RECORD) { + return self._handleRecordCursorMethod( + original, + this, + args, + inputValue, + spanName, + submodule, + creationContext, + ); + } + + if (self.mode === TuskDriftMode.REPLAY) { + return self._handleReplayCursorMethod( + inputValue, + spanName, + submodule, + methodName, + creationContext, + ); + } + + return original.apply(this, args); + }; + }; + } + + /** + * Handle RECORD mode for a cursor-returning method. + * Calls the original to get a real cursor, then wraps its terminal methods + * to capture documents and create spans. + */ + private _handleRecordCursorMethod( + original: Function, + thisArg: any, + args: any[], + inputValue: MongodbCommandInputValue, + spanName: string, + submodule: string, + creationContext: Context, + ): any { + // Let MongoDB create the real cursor (no query executes yet) + const cursor = original.apply(thisArg, args); + + // Shared state for this cursor's lifetime + const cursorState = { + collectedDocuments: [] as any[], + spanInfo: null as SpanInfo | null, + recorded: false, + spanCreated: false, + }; + + this._wrapCursorTerminalMethods( + cursor, + inputValue, + spanName, + submodule, + creationContext, + cursorState, + ); + + return cursor; + } + + /** + * Wrap terminal methods on a real cursor instance for RECORD mode. + * All terminal methods share state (collected documents, span, recorded flag). + */ + private _wrapCursorTerminalMethods( + cursor: any, + inputValue: MongodbCommandInputValue, + spanName: string, + submodule: string, + creationContext: Context, + cursorState: { + collectedDocuments: any[]; + spanInfo: SpanInfo | null; + recorded: boolean; + spanCreated: boolean; + }, + ): void { + const self = this; + + // Helper: finalize the cursor span (called once when cursor is exhausted) + const finalizeCursorSpan = (): void => { + if (cursorState.recorded || !cursorState.spanInfo) return; + cursorState.recorded = true; + try { + addOutputAttributesToSpan( + cursorState.spanInfo, + wrapCursorOutput(cursorState.collectedDocuments), + ); + SpanUtils.endSpan(cursorState.spanInfo.span, { + code: SpanStatusCode.OK, + }); + } catch (error) { + logger.error( + `[${self.INSTRUMENTATION_NAME}] Error finalizing cursor span for ${spanName}:`, + error, + ); + } + }; + + // Helper: handle errors on the cursor span + const handleCursorError = (error: any): void => { + if (cursorState.recorded || !cursorState.spanInfo) return; + cursorState.recorded = true; + try { + SpanUtils.addSpanAttributes(cursorState.spanInfo.span, { + outputValue: { error: error?.message || "Unknown error" }, + }); + SpanUtils.endSpan(cursorState.spanInfo.span, { + code: SpanStatusCode.ERROR, + message: error?.message || "Cursor operation failed", + }); + } catch (spanError) { + logger.error( + `[${self.INSTRUMENTATION_NAME}] Error recording cursor error span for ${spanName}:`, + spanError, + ); + } + }; + + // Helper: create span on first terminal call for next/hasNext/asyncIterator + // Returns false if recording should be skipped (background request). + const ensureSpanCreated = (): boolean => { + if (cursorState.spanCreated) return cursorState.spanInfo !== null; + cursorState.spanCreated = true; + + const isAppReady = self.tuskDrift.isAppReady(); + const currentSpanInfo = SpanUtils.getCurrentSpanInfo(); + + // Background request: app ready, no parent span, not a server request + if (isAppReady && !currentSpanInfo) { + return false; + } + + const isPreAppStart = !isAppReady; + + const spanInfo = SpanUtils.createSpan({ + name: spanName, + kind: SpanKind.CLIENT, + isPreAppStart, + parentContext: creationContext, + attributes: { + [TdSpanAttributes.NAME]: spanName, + [TdSpanAttributes.PACKAGE_NAME]: "mongodb", + [TdSpanAttributes.SUBMODULE_NAME]: submodule, + [TdSpanAttributes.INSTRUMENTATION_NAME]: self.INSTRUMENTATION_NAME, + [TdSpanAttributes.PACKAGE_TYPE]: PackageType.MONGODB, + [TdSpanAttributes.INPUT_VALUE]: createSpanInputValue(inputValue), + [TdSpanAttributes.IS_PRE_APP_START]: isPreAppStart, + }, + }); + + if (!spanInfo) { + return false; + } + + cursorState.spanInfo = spanInfo; + return true; + }; + + // --- Wrap toArray --- + if (typeof cursor.toArray === "function") { + const originalToArray = cursor.toArray.bind(cursor); + cursor.toArray = (): Promise => { + if (cursorState.recorded) return originalToArray(); + cursorState.recorded = true; + + return context.with(creationContext, () => { + return handleRecordMode({ + originalFunctionCall: () => originalToArray(), + recordModeHandler: ({ isPreAppStart }) => { + return SpanUtils.createAndExecuteSpan( + self.mode, + () => originalToArray(), + { + name: spanName, + kind: SpanKind.CLIENT, + submodule, + packageType: PackageType.MONGODB, + packageName: "mongodb", + instrumentationName: self.INSTRUMENTATION_NAME, + inputValue, + isPreAppStart, + stopRecordingChildSpans: true, + }, + (spanInfo: SpanInfo) => { + return originalToArray() + .then((result: any[]) => { + try { + addOutputAttributesToSpan(spanInfo, wrapCursorOutput(result)); + SpanUtils.endSpan(spanInfo.span, { + code: SpanStatusCode.OK, + }); + } catch (error) { + logger.error( + `[${self.INSTRUMENTATION_NAME}] Error adding span attributes for ${spanName}.toArray:`, + error, + ); + } + return result; + }) + .catch((error: any) => { + try { + SpanUtils.addSpanAttributes(spanInfo.span, { + outputValue: { + error: error?.message || "Unknown error", + }, + }); + SpanUtils.endSpan(spanInfo.span, { + code: SpanStatusCode.ERROR, + message: error?.message || "toArray failed", + }); + } catch (spanError) { + logger.error( + `[${self.INSTRUMENTATION_NAME}] Error recording cursor error span for ${spanName}.toArray:`, + spanError, + ); + } + throw error; + }); + }, + ); + }, + spanKind: SpanKind.CLIENT, + }); + }); + }; + } + + // --- Wrap next --- + if (typeof cursor.next === "function") { + const originalNext = cursor.next.bind(cursor); + cursor.next = async (): Promise => { + if (cursorState.recorded) return originalNext(); + + const shouldRecord = context.with(creationContext, () => ensureSpanCreated()); + if (!shouldRecord) return originalNext(); + + try { + const doc = await (cursorState.spanInfo + ? SpanUtils.withSpan(cursorState.spanInfo, () => originalNext()) + : originalNext()); + + if (doc !== null) { + cursorState.collectedDocuments.push(doc); + } else { + finalizeCursorSpan(); + } + return doc; + } catch (error) { + handleCursorError(error); + throw error; + } + }; + } + + // --- Wrap hasNext --- + if (typeof cursor.hasNext === "function") { + const originalHasNext = cursor.hasNext.bind(cursor); + cursor.hasNext = async (): Promise => { + if (cursorState.recorded) return originalHasNext(); + + const shouldRecord = context.with(creationContext, () => ensureSpanCreated()); + if (!shouldRecord) return originalHasNext(); + + try { + const result = await (cursorState.spanInfo + ? SpanUtils.withSpan(cursorState.spanInfo, () => originalHasNext()) + : originalHasNext()); + + if (!result) { + finalizeCursorSpan(); + } + return result; + } catch (error) { + handleCursorError(error); + throw error; + } + }; + } + + // --- Wrap forEach --- + if (typeof cursor.forEach === "function") { + const originalForEach = cursor.forEach.bind(cursor); + cursor.forEach = (iterator: (doc: any) => boolean | void): Promise => { + if (cursorState.recorded) return originalForEach(iterator); + cursorState.recorded = true; + + return context.with(creationContext, () => { + return handleRecordMode({ + originalFunctionCall: () => originalForEach(iterator), + recordModeHandler: ({ isPreAppStart }) => { + return SpanUtils.createAndExecuteSpan( + self.mode, + () => originalForEach(iterator), + { + name: spanName, + kind: SpanKind.CLIENT, + submodule, + packageType: PackageType.MONGODB, + packageName: "mongodb", + instrumentationName: self.INSTRUMENTATION_NAME, + inputValue, + isPreAppStart, + stopRecordingChildSpans: true, + }, + (spanInfo: SpanInfo) => { + const collectedDocs: any[] = []; + + const wrappedIterator = (doc: any) => { + collectedDocs.push(doc); + return iterator(doc); + }; + + return originalForEach(wrappedIterator) + .then(() => { + try { + addOutputAttributesToSpan(spanInfo, wrapCursorOutput(collectedDocs)); + SpanUtils.endSpan(spanInfo.span, { + code: SpanStatusCode.OK, + }); + } catch (error) { + logger.error( + `[${self.INSTRUMENTATION_NAME}] Error adding span attributes for ${spanName}.forEach:`, + error, + ); + } + }) + .catch((error: any) => { + try { + SpanUtils.addSpanAttributes(spanInfo.span, { + outputValue: { + error: error?.message || "Unknown error", + }, + }); + SpanUtils.endSpan(spanInfo.span, { + code: SpanStatusCode.ERROR, + message: error?.message || "forEach failed", + }); + } catch (spanError) { + logger.error( + `[${self.INSTRUMENTATION_NAME}] Error recording cursor error span for ${spanName}.forEach:`, + spanError, + ); + } + throw error; + }); + }, + ); + }, + spanKind: SpanKind.CLIENT, + }); + }); + }; + } + + // --- Wrap [Symbol.asyncIterator] --- + if (typeof cursor[Symbol.asyncIterator] === "function") { + const originalAsyncIterator = cursor[Symbol.asyncIterator].bind(cursor); + cursor[Symbol.asyncIterator] = async function* () { + const shouldRecord = context.with(creationContext, () => ensureSpanCreated()); + if (!shouldRecord || cursorState.recorded) { + yield* originalAsyncIterator(); + return; + } + cursorState.recorded = true; + let spanFinalized = false; + + try { + for await (const doc of { + [Symbol.asyncIterator]: originalAsyncIterator, + }) { + cursorState.collectedDocuments.push(doc); + yield doc; + } + // Directly finalize span — can't use finalizeCursorSpan() since recorded=true + // is needed early to prevent the wrapped next() from double-recording documents + if (cursorState.spanInfo) { + addOutputAttributesToSpan( + cursorState.spanInfo, + wrapCursorOutput(cursorState.collectedDocuments), + ); + SpanUtils.endSpan(cursorState.spanInfo.span, { + code: SpanStatusCode.OK, + }); + spanFinalized = true; + } + } catch (error: any) { + // Directly handle error on span (same reason as above) + if (cursorState.spanInfo) { + SpanUtils.addSpanAttributes(cursorState.spanInfo.span, { + outputValue: { error: error?.message || "Unknown error" }, + }); + SpanUtils.endSpan(cursorState.spanInfo.span, { + code: SpanStatusCode.ERROR, + message: error?.message || "Cursor operation failed", + }); + spanFinalized = true; + } + throw error; + } finally { + // Handle early break from for-await loop + if (!spanFinalized && cursorState.spanInfo) { + addOutputAttributesToSpan( + cursorState.spanInfo, + wrapCursorOutput(cursorState.collectedDocuments), + ); + SpanUtils.endSpan(cursorState.spanInfo.span, { + code: SpanStatusCode.OK, + }); + } + } + }; + } + } + + /** + * Handle REPLAY mode for a cursor-returning method. + * Returns a lazy fake cursor that defers mock data loading until + * the first terminal method is called. + */ + private _handleReplayCursorMethod( + inputValue: MongodbCommandInputValue, + spanName: string, + submodule: string, + methodName: string, + creationContext: Context, + ): any { + const self = this; + const stackTrace = captureStackTrace(["MongodbInstrumentation"]); + + // Create a mock data loader function (called lazily on first terminal method) + const loadMockData = (): Promise => { + return context.with(creationContext, () => { + return handleReplayMode({ + noOpRequestHandler: () => Promise.resolve([]), + isServerRequest: false, + replayModeHandler: () => { + return SpanUtils.createAndExecuteSpan( + self.mode, + () => Promise.resolve([]), + { + name: spanName, + kind: SpanKind.CLIENT, + submodule, + packageType: PackageType.MONGODB, + packageName: "mongodb", + instrumentationName: self.INSTRUMENTATION_NAME, + inputValue, + isPreAppStart: !self.tuskDrift.isAppReady(), + stopRecordingChildSpans: true, + }, + async (spanInfo: SpanInfo) => { + try { + const mockData = await self._findMockData({ + spanInfo, + name: spanName, + inputValue, + submoduleName: submodule, + stackTrace, + }); + + if (!mockData) { + const errorMsg = `[${self.INSTRUMENTATION_NAME}] No matching mock found for ${spanName} (collection: ${inputValue.collection})`; + logger.warn(errorMsg); + throw new Error(errorMsg); + } + + const reconstructed = reconstructBsonValue(mockData.result, self.moduleExports); + const documents = unwrapCursorOutput(reconstructed); + + SpanUtils.endSpan(spanInfo.span, { + code: SpanStatusCode.OK, + }); + return documents; + } catch (error: any) { + SpanUtils.endSpan(spanInfo.span, { + code: SpanStatusCode.ERROR, + message: error?.message || "Replay failed", + }); + throw error; + } + }, + ); + }, + }); + }); + }; + + if (methodName === "aggregate") { + return new TdFakeAggregationCursor([], loadMockData); + } + return new TdFakeFindCursor([], loadMockData); + } + + /** + * Build the input value for a cursor-returning method call. + * For find: captures filter and options. + * For aggregate: captures pipeline and options. + */ + private _extractCursorInput( + methodName: string, + collectionName: string | undefined, + databaseName: string | undefined, + args: any[], + ): MongodbCommandInputValue { + let commandArgs: Record; + + if (methodName === "find") { + commandArgs = { + filter: args[0] || {}, + options: sanitizeOptions(args[1]), + }; + } else if (methodName === "aggregate") { + commandArgs = { + pipeline: args[0] || [], + options: sanitizeOptions(args[1]), + }; + } else if (methodName === "listIndexes") { + commandArgs = { + options: sanitizeOptions(args[0]), + }; + } else { + commandArgs = { args: args.map((a: any, i: number) => ({ [i]: a })) }; + } + + return { + command: methodName, + collection: collectionName, + database: databaseName, + commandArgs: sanitizeBsonValue(commandArgs), + }; + } + + // --------------------------------------------------------------------------- + // Input extraction helpers + // --------------------------------------------------------------------------- + + /** + * Build the input value for a collection method call. + * Extracts relevant arguments per method and sanitizes session/BSON. + */ + private _extractCollectionInput( + methodName: string, + collectionName: string | undefined, + databaseName: string | undefined, + args: any[], + ): MongodbCommandInputValue { + const commandArgs = this._extractCommandArgs(methodName, args); + + return { + command: methodName, + collection: collectionName, + database: databaseName, + commandArgs: sanitizeBsonValue(commandArgs), + }; + } + + /** + * Extract command-specific arguments from the method call args. + * Strips session from options for all methods. + * Strips non-deterministic metadata from all operations so that mock + * matching uses only the stable, user-provided content. + */ + private _extractCommandArgs(methodName: string, args: any[]): Record { + // Strip non-deterministic metadata fields (_id, createdAt, updatedAt, __v) + // from objects so the value hash is stable across recording and replay. + // These fields change every run (client-generated ObjectIds, timestamps, + // Mongoose version keys) and would cause hash mismatches. + const stripMetadata = (obj: any) => { + if (!obj || typeof obj !== "object") return obj; + const { _id, createdAt, updatedAt, __v, ...rest } = obj; + return rest; + }; + + switch (methodName) { + case "findOne": + case "countDocuments": + // (filter?, options?) + return { + filter: stripMetadata(args[0]), + options: sanitizeOptions(args[1]), + }; + + case "estimatedDocumentCount": + // (options?) + return { + options: sanitizeOptions(args[0]), + }; + + case "insertOne": + // (doc, options?) + return { + document: stripMetadata(args[0]), + options: sanitizeOptions(args[1]), + }; + + case "insertMany": + // (docs, options?) + return { + documents: Array.isArray(args[0]) ? args[0].map(stripMetadata) : args[0], + options: sanitizeOptions(args[1]), + }; + + case "updateOne": + case "updateMany": + case "findOneAndUpdate": + // (filter, update, options?) + return { + filter: stripMetadata(args[0]), + update: args[1], + options: sanitizeOptions(args[2]), + }; + + case "deleteOne": + case "deleteMany": + case "findOneAndDelete": + // (filter, options?) + return { + filter: stripMetadata(args[0]), + options: sanitizeOptions(args[1]), + }; + + case "replaceOne": + case "findOneAndReplace": + // (filter, replacement, options?) + return { + filter: stripMetadata(args[0]), + replacement: args[1], + options: sanitizeOptions(args[2]), + }; + + case "distinct": + // (key, filter?, options?) + return { + key: args[0], + filter: stripMetadata(args[1]), + options: sanitizeOptions(args[2]), + }; + + case "bulkWrite": + // (operations, options?) + return { + operations: args[0], + options: sanitizeOptions(args[1]), + }; + + // Collection index operations + case "createIndex": + // (indexSpec, options?) + return { + indexSpec: args[0], + options: sanitizeOptions(args[1]), + }; + + case "createIndexes": + // (indexSpecs, options?) + return { + indexSpecs: args[0], + options: sanitizeOptions(args[1]), + }; + + case "dropIndex": + // (indexName, options?) + return { + indexName: args[0], + options: sanitizeOptions(args[1]), + }; + + case "dropIndexes": + // (options?) + return { + options: sanitizeOptions(args[0]), + }; + + case "indexes": + // (options?) + return { + options: sanitizeOptions(args[0]), + }; + + default: + return { args: args.map((a: any, i: number) => ({ [i]: a })) }; + } + } + + // --------------------------------------------------------------------------- + // Mock data lookup + // --------------------------------------------------------------------------- + + /** + * Find mock response data for replay mode. + * Wraps the core findMockResponseAsync utility. + */ + private async _findMockData({ + spanInfo, + name, + inputValue, + submoduleName, + stackTrace, + }: { + spanInfo: SpanInfo; + name: string; + inputValue: any; + submoduleName: string; + stackTrace?: string; + }): Promise { + return findMockResponseAsync({ + mockRequestData: { + traceId: spanInfo.traceId, + spanId: spanInfo.spanId, + name, + inputValue: createMockInputValue(inputValue), + packageName: "mongodb", + instrumentationName: this.INSTRUMENTATION_NAME, + submoduleName, + kind: SpanKind.CLIENT, + stackTrace, + }, + tuskDrift: this.tuskDrift, + }); + } + + // --------------------------------------------------------------------------- + // No-op results for replay background requests + // --------------------------------------------------------------------------- + + /** + * Returns an appropriate empty result for a given collection method. + * Used for background requests in replay mode (app ready, no parent span). + */ + private _getNoOpResult(methodName: string): any { + switch (methodName) { + case "findOne": + case "findOneAndUpdate": + case "findOneAndDelete": + case "findOneAndReplace": + return null; + + case "insertOne": + return { acknowledged: false, insertedId: null }; + + case "insertMany": + return { acknowledged: false, insertedIds: {}, insertedCount: 0 }; + + case "updateOne": + case "updateMany": + case "replaceOne": + return { + acknowledged: false, + matchedCount: 0, + modifiedCount: 0, + upsertedCount: 0, + upsertedId: null, + }; + + case "deleteOne": + case "deleteMany": + return { acknowledged: false, deletedCount: 0 }; + + case "countDocuments": + case "estimatedDocumentCount": + return 0; + + case "distinct": + return []; + + case "bulkWrite": + return { + acknowledged: false, + insertedCount: 0, + matchedCount: 0, + modifiedCount: 0, + deletedCount: 0, + upsertedCount: 0, + insertedIds: {}, + upsertedIds: {}, + }; + + // Collection index operations + case "createIndex": + return ""; + + case "createIndexes": + return []; + + case "dropIndex": + return {}; + + case "dropIndexes": + return { ok: 1 }; + + case "indexes": + return []; + + default: + return null; + } + } + + // --------------------------------------------------------------------------- + // Db-level promise method patching + // --------------------------------------------------------------------------- + + /** + * Patch Db.prototype methods that return Promises (command, createCollection, + * dropCollection, dropDatabase). + */ + private _patchDbMethods(actualExports: any): void { + const Db = actualExports.Db; + if (!Db || !Db.prototype) { + logger.warn( + `[${this.INSTRUMENTATION_NAME}] Db not found in module exports, skipping Db method patching`, + ); + return; + } + + for (const methodName of DB_METHODS_TO_WRAP) { + if (typeof Db.prototype[methodName] === "function") { + this._wrap(Db.prototype, methodName, this._getDbMethodWrapper(methodName)); + logger.debug(`[${this.INSTRUMENTATION_NAME}] Wrapped Db.prototype.${methodName}`); + } else { + logger.debug( + `[${this.INSTRUMENTATION_NAME}] Db.prototype.${methodName} not found, skipping`, + ); + } + } + } + + /** + * Returns a wrapper function for a given Db method. + * Handles RECORD, REPLAY, and DISABLED modes. + */ + private _getDbMethodWrapper(methodName: string): (original: Function) => Function { + const self = this; + const spanName = `mongodb.db.${methodName}`; + const submodule = `db-${methodName}`; + + return (original: Function) => { + return function (this: any, ...args: any[]) { + const databaseName = this?.databaseName || this?.s?.namespace?.db; + const inputValue = self._extractDbInput(methodName, databaseName, args); + + if (self.mode === TuskDriftMode.DISABLED) { + return original.apply(this, args); + } + + if (self.mode === TuskDriftMode.RECORD) { + if (methodName === "createCollection") { + return self._handleRecordCreateCollection( + original, + this, + args, + inputValue, + spanName, + submodule, + ); + } + return self._handleRecordDbMethod(original, this, args, inputValue, spanName, submodule); + } + + if (self.mode === TuskDriftMode.REPLAY) { + if (methodName === "createCollection") { + return self._handleReplayCreateCollection(this, args, inputValue, spanName, submodule); + } + return self._handleReplayDbMethod( + original, + this, + args, + inputValue, + spanName, + submodule, + methodName, + ); + } + + return original.apply(this, args); + }; + }; + } + + /** + * Handle RECORD mode for a Db method. + * Same structure as _handleRecordCollectionMethod. + */ + private _handleRecordDbMethod( + original: Function, + thisArg: any, + args: any[], + inputValue: MongodbCommandInputValue, + spanName: string, + submodule: string, + ): any { + return handleRecordMode({ + originalFunctionCall: () => original.apply(thisArg, args), + recordModeHandler: ({ isPreAppStart }) => { + return SpanUtils.createAndExecuteSpan( + this.mode, + () => original.apply(thisArg, args), + { + name: spanName, + kind: SpanKind.CLIENT, + submodule, + packageType: PackageType.MONGODB, + packageName: "mongodb", + instrumentationName: this.INSTRUMENTATION_NAME, + inputValue, + isPreAppStart, + stopRecordingChildSpans: true, + }, + (spanInfo: SpanInfo) => { + const resultPromise = original.apply(thisArg, args) as Promise; + + return resultPromise + .then((result: any) => { + try { + addOutputAttributesToSpan(spanInfo, wrapDirectOutput(result)); + SpanUtils.endSpan(spanInfo.span, { + code: SpanStatusCode.OK, + }); + } catch (error) { + logger.error( + `[${this.INSTRUMENTATION_NAME}] Error adding span attributes for ${spanName}:`, + error, + ); + } + return result; + }) + .catch((error: any) => { + try { + SpanUtils.addSpanAttributes(spanInfo.span, { + outputValue: { + error: error?.message || "Unknown error", + }, + }); + SpanUtils.endSpan(spanInfo.span, { + code: SpanStatusCode.ERROR, + message: error?.message || "Operation failed", + }); + } catch (spanError) { + logger.error( + `[${this.INSTRUMENTATION_NAME}] Error recording span for ${spanName} error:`, + spanError, + ); + } + throw error; + }); + }, + ); + }, + spanKind: SpanKind.CLIENT, + }); + } + + /** + * Handle REPLAY mode for a Db method. + * Same structure as _handleReplayCollectionMethod. + */ + private _handleReplayDbMethod( + original: Function, + thisArg: any, + args: any[], + inputValue: MongodbCommandInputValue, + spanName: string, + submodule: string, + methodName: string, + ): any { + const stackTrace = captureStackTrace(["MongodbInstrumentation"]); + + return handleReplayMode({ + noOpRequestHandler: () => { + return Promise.resolve(this._getDbNoOpResult(methodName)); + }, + isServerRequest: false, + replayModeHandler: () => { + return SpanUtils.createAndExecuteSpan( + this.mode, + () => original.apply(thisArg, args), + { + name: spanName, + kind: SpanKind.CLIENT, + submodule, + packageType: PackageType.MONGODB, + packageName: "mongodb", + instrumentationName: this.INSTRUMENTATION_NAME, + inputValue, + isPreAppStart: !this.tuskDrift.isAppReady(), + stopRecordingChildSpans: true, + }, + async (spanInfo: SpanInfo) => { + try { + const mockData = await this._findMockData({ + spanInfo, + name: spanName, + inputValue, + submoduleName: submodule, + stackTrace, + }); + + if (!mockData) { + const errorMsg = `[${this.INSTRUMENTATION_NAME}] No matching mock found for ${spanName} (database: ${inputValue.database})`; + logger.warn(errorMsg); + throw new Error(errorMsg); + } + + const result = unwrapDirectOutput( + reconstructBsonValue(mockData.result, this.moduleExports), + ); + + SpanUtils.endSpan(spanInfo.span, { code: SpanStatusCode.OK }); + return result; + } catch (error: any) { + SpanUtils.endSpan(spanInfo.span, { + code: SpanStatusCode.ERROR, + message: error?.message || "Replay failed", + }); + throw error; + } + }, + ); + }, + }); + } + + // --------------------------------------------------------------------------- + // createCollection special handling + // --------------------------------------------------------------------------- + + /** + * Handle RECORD mode for createCollection. + * Collection objects are not serializable, so we record + * { collectionName } as the output instead of the actual Collection. + */ + private _handleRecordCreateCollection( + original: Function, + thisArg: any, + args: any[], + inputValue: MongodbCommandInputValue, + spanName: string, + submodule: string, + ): any { + return handleRecordMode({ + originalFunctionCall: () => original.apply(thisArg, args), + recordModeHandler: ({ isPreAppStart }) => { + return SpanUtils.createAndExecuteSpan( + this.mode, + () => original.apply(thisArg, args), + { + name: spanName, + kind: SpanKind.CLIENT, + submodule, + packageType: PackageType.MONGODB, + packageName: "mongodb", + instrumentationName: this.INSTRUMENTATION_NAME, + inputValue, + isPreAppStart, + stopRecordingChildSpans: true, + }, + (spanInfo: SpanInfo) => { + const resultPromise = original.apply(thisArg, args) as Promise; + + return resultPromise + .then((result: any) => { + try { + // Collection objects are not serializable — record the name only + const collectionName = args[0]; + SpanUtils.addSpanAttributes(spanInfo.span, { + outputValue: sanitizeBsonValue({ collectionName }), + }); + SpanUtils.endSpan(spanInfo.span, { + code: SpanStatusCode.OK, + }); + } catch (error) { + logger.error( + `[${this.INSTRUMENTATION_NAME}] Error adding span attributes for ${spanName}:`, + error, + ); + } + return result; + }) + .catch((error: any) => { + try { + SpanUtils.addSpanAttributes(spanInfo.span, { + outputValue: { + error: error?.message || "Unknown error", + }, + }); + SpanUtils.endSpan(spanInfo.span, { + code: SpanStatusCode.ERROR, + message: error?.message || "Operation failed", + }); + } catch (spanError) { + logger.error( + `[${this.INSTRUMENTATION_NAME}] Error recording span for ${spanName} error:`, + spanError, + ); + } + throw error; + }); + }, + ); + }, + spanKind: SpanKind.CLIENT, + }); + } + + /** + * Handle REPLAY mode for createCollection. + * Returns a Collection reference via thisArg.collection(name) which + * is synchronous and doesn't require a server connection. + */ + private _handleReplayCreateCollection( + thisArg: any, + args: any[], + inputValue: MongodbCommandInputValue, + spanName: string, + submodule: string, + ): any { + const stackTrace = captureStackTrace(["MongodbInstrumentation"]); + const collectionName = args[0]; + + return handleReplayMode({ + noOpRequestHandler: () => { + return Promise.resolve(thisArg.collection(collectionName)); + }, + isServerRequest: false, + replayModeHandler: () => { + return SpanUtils.createAndExecuteSpan( + this.mode, + () => Promise.resolve(thisArg.collection(collectionName)), + { + name: spanName, + kind: SpanKind.CLIENT, + submodule, + packageType: PackageType.MONGODB, + packageName: "mongodb", + instrumentationName: this.INSTRUMENTATION_NAME, + inputValue, + isPreAppStart: !this.tuskDrift.isAppReady(), + stopRecordingChildSpans: true, + }, + async (spanInfo: SpanInfo) => { + try { + // Consume mock data to advance the mock counter + await this._findMockData({ + spanInfo, + name: spanName, + inputValue, + submoduleName: submodule, + stackTrace, + }); + + SpanUtils.endSpan(spanInfo.span, { code: SpanStatusCode.OK }); + return thisArg.collection(collectionName); + } catch (error: any) { + SpanUtils.endSpan(spanInfo.span, { + code: SpanStatusCode.ERROR, + message: error?.message || "Replay failed", + }); + throw error; + } + }, + ); + }, + }); + } + + // --------------------------------------------------------------------------- + // Db-level cursor method patching (listCollections, Db.aggregate) + // --------------------------------------------------------------------------- + + /** + * Patch Db.prototype cursor-returning methods (listCollections, aggregate). + */ + private _patchDbCursorMethods(actualExports: any): void { + const Db = actualExports.Db; + if (!Db || !Db.prototype) { + logger.warn( + `[${this.INSTRUMENTATION_NAME}] Db not found, skipping Db cursor method patching`, + ); + return; + } + + for (const methodName of DB_CURSOR_METHODS_TO_WRAP) { + if (typeof Db.prototype[methodName] === "function") { + this._wrap(Db.prototype, methodName, this._getDbCursorMethodWrapper(methodName)); + logger.debug(`[${this.INSTRUMENTATION_NAME}] Wrapped Db.prototype.${methodName}`); + } else { + logger.debug( + `[${this.INSTRUMENTATION_NAME}] Db.prototype.${methodName} not found, skipping`, + ); + } + } + } + + /** + * Returns a wrapper function for a cursor-returning Db method. + * Similar to _getCursorMethodWrapper but for Db-level operations. + */ + private _getDbCursorMethodWrapper(methodName: string): (original: Function) => Function { + const self = this; + const spanName = `mongodb.db.${methodName}`; + const submodule = `db-${methodName}`; + + return (original: Function) => { + return function (this: any, ...args: any[]) { + const databaseName = this?.databaseName || this?.s?.namespace?.db; + const inputValue = self._extractDbCursorInput(methodName, databaseName, args); + + if (self.mode === TuskDriftMode.DISABLED) { + return original.apply(this, args); + } + + const creationContext = context.active(); + + if (self.mode === TuskDriftMode.RECORD) { + return self._handleRecordCursorMethod( + original, + this, + args, + inputValue, + spanName, + submodule, + creationContext, + ); + } + + if (self.mode === TuskDriftMode.REPLAY) { + return self._handleReplayDbCursorMethod( + inputValue, + spanName, + submodule, + methodName, + creationContext, + ); + } + + return original.apply(this, args); + }; + }; + } + + /** + * Handle REPLAY mode for a Db cursor-returning method. + * Returns a fake cursor with lazy mock data loading. + */ + private _handleReplayDbCursorMethod( + inputValue: MongodbCommandInputValue, + spanName: string, + submodule: string, + methodName: string, + creationContext: Context, + ): any { + const self = this; + const stackTrace = captureStackTrace(["MongodbInstrumentation"]); + + const loadMockData = (): Promise => { + return context.with(creationContext, () => { + return handleReplayMode({ + noOpRequestHandler: () => Promise.resolve([]), + isServerRequest: false, + replayModeHandler: () => { + return SpanUtils.createAndExecuteSpan( + self.mode, + () => Promise.resolve([]), + { + name: spanName, + kind: SpanKind.CLIENT, + submodule, + packageType: PackageType.MONGODB, + packageName: "mongodb", + instrumentationName: self.INSTRUMENTATION_NAME, + inputValue, + isPreAppStart: !self.tuskDrift.isAppReady(), + stopRecordingChildSpans: true, + }, + async (spanInfo: SpanInfo) => { + try { + const mockData = await self._findMockData({ + spanInfo, + name: spanName, + inputValue, + submoduleName: submodule, + stackTrace, + }); + + if (!mockData) { + const errorMsg = `[${self.INSTRUMENTATION_NAME}] No matching mock found for ${spanName} (database: ${inputValue.database})`; + logger.warn(errorMsg); + throw new Error(errorMsg); + } + + const reconstructed = reconstructBsonValue(mockData.result, self.moduleExports); + const documents = unwrapCursorOutput(reconstructed); + + SpanUtils.endSpan(spanInfo.span, { + code: SpanStatusCode.OK, + }); + return documents; + } catch (error: any) { + SpanUtils.endSpan(spanInfo.span, { + code: SpanStatusCode.ERROR, + message: error?.message || "Replay failed", + }); + throw error; + } + }, + ); + }, + }); + }); + }; + + if (methodName === "aggregate") { + return new TdFakeAggregationCursor([], loadMockData); + } + return new TdFakeFindCursor([], loadMockData); + } + + // --------------------------------------------------------------------------- + // Db-level input extraction + // --------------------------------------------------------------------------- + + /** + * Build the input value for a Db method call. + */ + private _extractDbInput( + methodName: string, + databaseName: string | undefined, + args: any[], + ): MongodbCommandInputValue { + const commandArgs = this._extractDbCommandArgs(methodName, args); + + return { + command: methodName, + database: databaseName, + commandArgs: sanitizeBsonValue(commandArgs), + }; + } + + /** + * Extract command-specific arguments from a Db method call. + */ + private _extractDbCommandArgs(methodName: string, args: any[]): Record { + switch (methodName) { + case "command": + // (command, options?) + return { + command: args[0], + options: sanitizeOptions(args[1]), + }; + + case "createCollection": + // (name, options?) + return { + collectionName: args[0], + options: sanitizeOptions(args[1]), + }; + + case "dropCollection": + // (name, options?) + return { + collectionName: args[0], + options: sanitizeOptions(args[1]), + }; + + case "dropDatabase": + // (options?) + return { + options: sanitizeOptions(args[0]), + }; + + default: + return { args: args.map((a: any, i: number) => ({ [i]: a })) }; + } + } + + /** + * Build the input value for a cursor-returning Db method call. + */ + private _extractDbCursorInput( + methodName: string, + databaseName: string | undefined, + args: any[], + ): MongodbCommandInputValue { + let commandArgs: Record; + + if (methodName === "listCollections") { + commandArgs = { + filter: args[0] || {}, + options: sanitizeOptions(args[1]), + }; + } else if (methodName === "aggregate") { + commandArgs = { + pipeline: args[0] || [], + options: sanitizeOptions(args[1]), + }; + } else { + commandArgs = { args: args.map((a: any, i: number) => ({ [i]: a })) }; + } + + return { + command: methodName, + database: databaseName, + commandArgs: sanitizeBsonValue(commandArgs), + }; + } + + // --------------------------------------------------------------------------- + // Db no-op results + // --------------------------------------------------------------------------- + + /** + * Returns an appropriate empty result for a given Db method. + * Used for background requests in replay mode. + */ + private _getDbNoOpResult(methodName: string): any { + switch (methodName) { + case "command": + return {}; + + case "dropCollection": + return true; + + case "dropDatabase": + return true; + + default: + return null; + } + } + + // --------------------------------------------------------------------------- + // Session & Transaction patching + // --------------------------------------------------------------------------- + + /** + * Patch ClientSession.prototype methods from the mongodb/lib/sessions.js + * internal module file. Wraps commitTransaction, abortTransaction, and + * endSession for record/replay. + */ + private _patchSessionModule(moduleExports: any): any { + logger.debug(`[${this.INSTRUMENTATION_NAME}] Patching session module`); + + if (this.isModulePatched(moduleExports)) { + logger.debug(`[${this.INSTRUMENTATION_NAME}] Session module already patched, skipping`); + return moduleExports; + } + + if (!moduleExports?.ClientSession?.prototype) { + logger.warn( + `[${this.INSTRUMENTATION_NAME}] ClientSession not found in sessions module, skipping`, + ); + return moduleExports; + } + + if (typeof moduleExports.ClientSession.prototype.commitTransaction === "function") { + this._wrap( + moduleExports.ClientSession.prototype, + "commitTransaction", + this._getSessionMethodWrapper("commitTransaction"), + ); + logger.debug( + `[${this.INSTRUMENTATION_NAME}] Wrapped ClientSession.prototype.commitTransaction`, + ); + } + + if (typeof moduleExports.ClientSession.prototype.abortTransaction === "function") { + this._wrap( + moduleExports.ClientSession.prototype, + "abortTransaction", + this._getSessionMethodWrapper("abortTransaction"), + ); + logger.debug( + `[${this.INSTRUMENTATION_NAME}] Wrapped ClientSession.prototype.abortTransaction`, + ); + } + + if (typeof moduleExports.ClientSession.prototype.endSession === "function") { + this._wrap(moduleExports.ClientSession.prototype, "endSession", this._getEndSessionWrapper()); + logger.debug(`[${this.INSTRUMENTATION_NAME}] Wrapped ClientSession.prototype.endSession`); + } + + this.markModuleAsPatched(moduleExports); + logger.debug(`[${this.INSTRUMENTATION_NAME}] Session module patching complete`); + return moduleExports; + } + + /** + * Returns a wrapper for MongoClient.prototype.startSession. + * startSession is synchronous and returns a ClientSession instance. + * In all modes we call the original — the session's prototype methods + * (commitTransaction, abortTransaction, endSession) are already patched + * via file-level patching of mongodb/lib/sessions.js. + */ + private _getStartSessionWrapper(): (original: Function) => Function { + const self = this; + return (original: Function) => { + return function (this: any, ...args: any[]) { + if (self.mode === TuskDriftMode.REPLAY) { + // In replay mode, still create a real session object. + // Its methods (commitTransaction, abortTransaction, endSession) + // are already patched on the prototype and will handle replay. + logger.debug(`[${self.INSTRUMENTATION_NAME}] startSession called in REPLAY mode`); + return original.apply(this, args); + } + return original.apply(this, args); + }; + }; + } + + // --------------------------------------------------------------------------- + // ChangeStream (watch) handling + // --------------------------------------------------------------------------- + + /** + * Patch watch() methods on Collection, Db, and MongoClient prototypes. + * ChangeStreams are long-lived event-based streams not suited for + * span-level record/replay. In RECORD mode: passthrough. + * In REPLAY mode: return a fake ChangeStream that emits no events. + */ + private _patchWatchMethods(actualExports: any): void { + const self = this; + + const targets = [ + { proto: actualExports.Collection?.prototype, name: "Collection" }, + { proto: actualExports.Db?.prototype, name: "Db" }, + { proto: actualExports.MongoClient?.prototype, name: "MongoClient" }, + ]; + + for (const { proto, name } of targets) { + if (proto && typeof proto.watch === "function") { + this._wrap(proto, "watch", (original: Function) => { + return function (this: any, ...args: any[]) { + if (self.mode === TuskDriftMode.REPLAY) { + logger.debug( + `[${self.INSTRUMENTATION_NAME}] ${name}.watch() called in REPLAY mode, returning fake ChangeStream`, + ); + return new TdFakeChangeStream(); + } + // RECORD and DISABLED: passthrough + return original.apply(this, args); + }; + }); + logger.debug(`[${this.INSTRUMENTATION_NAME}] Wrapped ${name}.prototype.watch`); + } + } + } + + /** + * Returns a wrapper function for commitTransaction or abortTransaction. + * Both are async and follow the same record/replay pattern. + */ + private _getSessionMethodWrapper(methodName: string): (original: Function) => Function { + const self = this; + const spanName = `mongodb.session.${methodName}`; + const submodule = `session-${methodName}`; + + return (original: Function) => { + return function (this: any, ...args: any[]) { + const inputValue: MongodbCommandInputValue = { + command: methodName, + }; + + if (self.mode === TuskDriftMode.DISABLED) { + return original.apply(this, args); + } + + if (self.mode === TuskDriftMode.RECORD) { + return self._handleRecordSessionMethod( + original, + this, + args, + inputValue, + spanName, + submodule, + ); + } + + if (self.mode === TuskDriftMode.REPLAY) { + return self._handleReplaySessionMethod( + original, + this, + args, + inputValue, + spanName, + submodule, + ); + } + + return original.apply(this, args); + }; + }; + } + + /** + * Handle RECORD mode for a session method (commitTransaction / abortTransaction). + * Calls the original method, wraps the promise to capture output, creates a span. + */ + private _handleRecordSessionMethod( + original: Function, + thisArg: any, + args: any[], + inputValue: MongodbCommandInputValue, + spanName: string, + submodule: string, + ): any { + return handleRecordMode({ + originalFunctionCall: () => original.apply(thisArg, args), + recordModeHandler: ({ isPreAppStart }) => { + return SpanUtils.createAndExecuteSpan( + this.mode, + () => original.apply(thisArg, args), + { + name: spanName, + kind: SpanKind.CLIENT, + submodule, + packageType: PackageType.MONGODB, + packageName: "mongodb", + instrumentationName: this.INSTRUMENTATION_NAME, + inputValue, + isPreAppStart, + stopRecordingChildSpans: true, + }, + (spanInfo: SpanInfo) => { + const resultPromise = original.apply(thisArg, args) as Promise; + + return resultPromise + .then((result: any) => { + try { + addOutputAttributesToSpan(spanInfo, result ?? { success: true }); + SpanUtils.endSpan(spanInfo.span, { + code: SpanStatusCode.OK, + }); + } catch (error) { + logger.error( + `[${this.INSTRUMENTATION_NAME}] Error adding span attributes for ${spanName}:`, + error, + ); + } + return result; + }) + .catch((error: any) => { + try { + SpanUtils.addSpanAttributes(spanInfo.span, { + outputValue: { + error: error?.message || "Unknown error", + }, + }); + SpanUtils.endSpan(spanInfo.span, { + code: SpanStatusCode.ERROR, + message: error?.message || "Operation failed", + }); + } catch (spanError) { + logger.error( + `[${this.INSTRUMENTATION_NAME}] Error recording span for ${spanName} error:`, + spanError, + ); + } + throw error; + }); + }, + ); + }, + spanKind: SpanKind.CLIENT, + }); + } + + /** + * Handle REPLAY mode for a session method (commitTransaction / abortTransaction). + * Looks up mock data, returns mocked result. + */ + private _handleReplaySessionMethod( + original: Function, + thisArg: any, + args: any[], + inputValue: MongodbCommandInputValue, + spanName: string, + submodule: string, + ): any { + const stackTrace = captureStackTrace(["MongodbInstrumentation"]); + + return handleReplayMode({ + noOpRequestHandler: () => Promise.resolve(undefined), + isServerRequest: false, + replayModeHandler: () => { + return SpanUtils.createAndExecuteSpan( + this.mode, + () => original.apply(thisArg, args), + { + name: spanName, + kind: SpanKind.CLIENT, + submodule, + packageType: PackageType.MONGODB, + packageName: "mongodb", + instrumentationName: this.INSTRUMENTATION_NAME, + inputValue, + isPreAppStart: !this.tuskDrift.isAppReady(), + stopRecordingChildSpans: true, + }, + async (spanInfo: SpanInfo) => { + try { + const mockData = await this._findMockData({ + spanInfo, + name: spanName, + inputValue, + submoduleName: submodule, + stackTrace, + }); + + // Session operations return void — treat missing mock result as success + const result = + mockData?.result != null + ? reconstructBsonValue(mockData.result, this.moduleExports) + : undefined; + + SpanUtils.endSpan(spanInfo.span, { + code: SpanStatusCode.OK, + }); + return result; + } catch (error: any) { + SpanUtils.endSpan(spanInfo.span, { + code: SpanStatusCode.ERROR, + message: error?.message || "Replay failed", + }); + throw error; + } + }, + ); + }, + }); + } + + /** + * Returns a wrapper for ClientSession.prototype.endSession. + * - RECORD: call original (let real session clean up) + * - REPLAY: no-op (no real session to end) + * - DISABLED: passthrough + */ + private _getEndSessionWrapper(): (original: Function) => Function { + const self = this; + return (original: Function) => { + return function (this: any, ...args: any[]) { + if (self.mode === TuskDriftMode.RECORD) { + return original.apply(this, args); + } + if (self.mode === TuskDriftMode.REPLAY) { + return Promise.resolve(); + } + return original.apply(this, args); + }; + }; + } + + // --------------------------------------------------------------------------- + // Bulk Operations (Ordered/Unordered) + // --------------------------------------------------------------------------- + + /** + * Patch Collection.prototype.initializeOrderedBulkOp and + * Collection.prototype.initializeUnorderedBulkOp. + * + * In replay mode, BulkOperationBase's constructor calls getTopology(collection) + * which throws if no topology is connected. We inject a FakeTopology onto the + * collection (and its client) BEFORE calling the original, so the constructor + * finds a valid topology object and proceeds with default size limits. + */ + private _patchBulkOpInitMethods(actualExports: any): void { + const Collection = actualExports.Collection; + if (!Collection || !Collection.prototype) { + logger.warn( + `[${this.INSTRUMENTATION_NAME}] Collection not found, skipping bulk op init patching`, + ); + return; + } + + const self = this; + + // Patch initializeOrderedBulkOp + if (typeof Collection.prototype.initializeOrderedBulkOp === "function") { + this._wrap(Collection.prototype, "initializeOrderedBulkOp", (original: Function) => { + return function (this: any, ...args: any[]) { + if (self.mode === TuskDriftMode.REPLAY) { + self._injectFakeTopology(this); + } + return original.apply(this, args); + }; + }); + logger.debug( + `[${this.INSTRUMENTATION_NAME}] Wrapped Collection.prototype.initializeOrderedBulkOp`, + ); + } + + // Patch initializeUnorderedBulkOp + if (typeof Collection.prototype.initializeUnorderedBulkOp === "function") { + this._wrap(Collection.prototype, "initializeUnorderedBulkOp", (original: Function) => { + return function (this: any, ...args: any[]) { + if (self.mode === TuskDriftMode.REPLAY) { + self._injectFakeTopology(this); + } + return original.apply(this, args); + }; + }); + logger.debug( + `[${this.INSTRUMENTATION_NAME}] Wrapped Collection.prototype.initializeUnorderedBulkOp`, + ); + } + } + + /** + * Inject a FakeTopology onto a collection and its client for replay mode. + * + * getTopology() in the MongoDB driver checks: + * 1. provider.topology (direct property on collection) + * 2. provider.client.topology (via the MongoClient) + * We set both to ensure the topology lookup succeeds. + */ + private _injectFakeTopology(collection: any): void { + const fakeTopology = new TdFakeTopology(); + + // Set on the client (satisfies getTopology's client.topology check) + if (collection.client && !collection.client.topology) { + collection.client.topology = fakeTopology; + } + + // Set on collection directly as fallback + if (!collection.topology) { + collection.topology = fakeTopology; + } + } + + /** + * Patch OrderedBulkOperation.prototype.execute from mongodb/lib/bulk/ordered.js. + */ + private _patchOrderedBulkModule(moduleExports: any): any { + logger.debug(`[${this.INSTRUMENTATION_NAME}] Patching ordered bulk operation module`); + + if (this.isModulePatched(moduleExports)) { + logger.debug(`[${this.INSTRUMENTATION_NAME}] Ordered bulk module already patched, skipping`); + return moduleExports; + } + + if (!moduleExports?.OrderedBulkOperation?.prototype) { + logger.warn( + `[${this.INSTRUMENTATION_NAME}] OrderedBulkOperation not found in bulk/ordered module, skipping`, + ); + return moduleExports; + } + + if (typeof moduleExports.OrderedBulkOperation.prototype.execute === "function") { + this._wrap( + moduleExports.OrderedBulkOperation.prototype, + "execute", + this._getBulkOpExecuteWrapper(true), + ); + logger.debug(`[${this.INSTRUMENTATION_NAME}] Wrapped OrderedBulkOperation.prototype.execute`); + } + + this.markModuleAsPatched(moduleExports); + return moduleExports; + } + + /** + * Patch UnorderedBulkOperation.prototype.execute from mongodb/lib/bulk/unordered.js. + */ + private _patchUnorderedBulkModule(moduleExports: any): any { + logger.debug(`[${this.INSTRUMENTATION_NAME}] Patching unordered bulk operation module`); + + if (this.isModulePatched(moduleExports)) { + logger.debug( + `[${this.INSTRUMENTATION_NAME}] Unordered bulk module already patched, skipping`, + ); + return moduleExports; + } + + if (!moduleExports?.UnorderedBulkOperation?.prototype) { + logger.warn( + `[${this.INSTRUMENTATION_NAME}] UnorderedBulkOperation not found in bulk/unordered module, skipping`, + ); + return moduleExports; + } + + if (typeof moduleExports.UnorderedBulkOperation.prototype.execute === "function") { + this._wrap( + moduleExports.UnorderedBulkOperation.prototype, + "execute", + this._getBulkOpExecuteWrapper(false), + ); + logger.debug( + `[${this.INSTRUMENTATION_NAME}] Wrapped UnorderedBulkOperation.prototype.execute`, + ); + } + + this.markModuleAsPatched(moduleExports); + return moduleExports; + } + + /** + * Returns a wrapper function for BulkOperation.prototype.execute. + * @param isOrdered — true for OrderedBulkOperation, false for UnorderedBulkOperation + */ + private _getBulkOpExecuteWrapper(isOrdered: boolean): (original: Function) => Function { + const self = this; + const opType = isOrdered ? "ordered" : "unordered"; + const spanName = "mongodb.bulkOp.execute"; + const submodule = `bulkOp-${opType}Execute`; + + return (original: Function) => { + return function (this: any, ...args: any[]) { + const collectionName = this?.s?.namespace?.collection; + const databaseName = this?.s?.namespace?.db; + const inputValue = self._extractBulkOpInput( + this, + isOrdered, + collectionName, + databaseName, + args, + ); + + if (self.mode === TuskDriftMode.DISABLED) { + return original.apply(this, args); + } + + if (self.mode === TuskDriftMode.RECORD) { + return self._handleRecordBulkOpExecute( + original, + this, + args, + inputValue, + spanName, + submodule, + ); + } + + if (self.mode === TuskDriftMode.REPLAY) { + return self._handleReplayBulkOpExecute( + original, + this, + args, + inputValue, + spanName, + submodule, + ); + } + + return original.apply(this, args); + }; + }; + } + + /** + * Extract input from a bulk operation's internal state. + * Reads queued operations from batches before execute() moves them. + */ + private _extractBulkOpInput( + bulkOp: any, + isOrdered: boolean, + collectionName: string | undefined, + databaseName: string | undefined, + args: any[], + ): MongodbCommandInputValue { + const operations = this._extractBulkOpOperations(bulkOp, isOrdered); + + return { + command: "bulkOp.execute", + collection: collectionName, + database: databaseName, + commandArgs: sanitizeBsonValue({ + isOrdered, + operations, + options: sanitizeOptions(args[0]), + }), + }; + } + + /** + * Extract readable operations from a bulk operation's internal state. + * Must be called BEFORE original execute() since execute() moves + * currentBatch into batches and clears state. + * + * Internal state layout: + * - Ordered: s.batches[] + s.currentBatch + * - Unordered: s.batches[] + s.currentInsertBatch + s.currentUpdateBatch + s.currentRemoveBatch + * + * BatchType constants: INSERT=1, UPDATE=2, DELETE=3 + */ + private _extractBulkOpOperations(bulkOp: any, isOrdered: boolean): any[] { + try { + const internalState = bulkOp?.s; + if (!internalState) { + return []; + } + + let batches: any[] = []; + + if (isOrdered) { + const { batches: internalBatches, currentBatch } = internalState; + batches = [...(internalBatches || [])]; + if (currentBatch) { + batches.push(currentBatch); + } + } else { + const { + batches: internalBatches, + currentInsertBatch, + currentUpdateBatch, + currentRemoveBatch, + } = internalState; + batches = [...(internalBatches || [])]; + if (currentInsertBatch) { + batches.push(currentInsertBatch); + } + if (currentUpdateBatch) { + batches.push(currentUpdateBatch); + } + if (currentRemoveBatch) { + batches.push(currentRemoveBatch); + } + } + + const readableOperations: any[] = []; + for (const batch of batches) { + if (!batch?.operations) { + continue; + } + const { batchType, operations } = batch; + for (const operation of operations) { + readableOperations.push(this._makeReadableOperation(operation, batchType)); + } + } + + return readableOperations; + } catch (error) { + logger.error(`[${this.INSTRUMENTATION_NAME}] Error extracting bulk op operations:`, error); + return []; + } + } + + /** + * Convert an internal bulk operation to a readable format. + * BatchType: 1=INSERT, 2=UPDATE, 3=DELETE + */ + private _makeReadableOperation(operation: any, batchType: number): Record { + const readableOp: Record = {}; + + // Strip non-deterministic metadata from documents and filters + const stripMetadata = (obj: any) => { + if (!obj || typeof obj !== "object") return obj; + const { _id, createdAt, updatedAt, __v, ...rest } = obj; + return rest; + }; + + if (batchType === 1) { + // INSERT + readableOp.operation = "Insert"; + readableOp.document = stripMetadata(operation); + } else if (batchType === 2) { + // UPDATE / REPLACE + if (!this._hasAtomicOperators(operation.u)) { + readableOp.operation = "ReplaceOne"; + } else if (operation.multi === true) { + readableOp.operation = "Update"; + } else { + readableOp.operation = "UpdateOne"; + } + readableOp.query = stripMetadata(operation.q); + readableOp.document = { ...operation.u }; + } else if (batchType === 3) { + // DELETE + readableOp.operation = operation.limit === 1 ? "DeleteOne" : "Delete"; + readableOp.query = stripMetadata(operation.q); + } + + if (operation.upsert) { + readableOp.upsert = operation.upsert; + } + if (operation.hint) { + readableOp.hint = operation.hint; + } + if (operation.collation) { + readableOp.collation = operation.collation; + } + if (operation.arrayFilters) { + readableOp.arrayFilters = operation.arrayFilters; + } + + return readableOp; + } + + /** + * Check if a document contains atomic operators (keys starting with '$'). + */ + private _hasAtomicOperators(doc: any): boolean { + if (Array.isArray(doc)) { + for (const item of doc) { + if (this._hasAtomicOperators(item)) { + return true; + } + } + return false; + } + if (!doc || typeof doc !== "object") { + return false; + } + const keys = Object.keys(doc); + return keys.length > 0 && keys[0][0] === "$"; + } + + /** + * Handle RECORD mode for BulkOperation.prototype.execute. + * Same pattern as _handleRecordCollectionMethod. + */ + private _handleRecordBulkOpExecute( + original: Function, + thisArg: any, + args: any[], + inputValue: MongodbCommandInputValue, + spanName: string, + submodule: string, + ): any { + return handleRecordMode({ + originalFunctionCall: () => original.apply(thisArg, args), + recordModeHandler: ({ isPreAppStart }) => { + return SpanUtils.createAndExecuteSpan( + this.mode, + () => original.apply(thisArg, args), + { + name: spanName, + kind: SpanKind.CLIENT, + submodule, + packageType: PackageType.MONGODB, + packageName: "mongodb", + instrumentationName: this.INSTRUMENTATION_NAME, + inputValue, + isPreAppStart, + stopRecordingChildSpans: true, + }, + (spanInfo: SpanInfo) => { + const resultPromise = original.apply(thisArg, args) as Promise; + + return resultPromise + .then((result: any) => { + try { + addOutputAttributesToSpan(spanInfo, wrapDirectOutput(result)); + SpanUtils.endSpan(spanInfo.span, { + code: SpanStatusCode.OK, + }); + } catch (error) { + logger.error( + `[${this.INSTRUMENTATION_NAME}] Error adding span attributes for ${spanName}:`, + error, + ); + } + return result; + }) + .catch((error: any) => { + try { + SpanUtils.addSpanAttributes(spanInfo.span, { + outputValue: { + error: error?.message || "Unknown error", + }, + }); + SpanUtils.endSpan(spanInfo.span, { + code: SpanStatusCode.ERROR, + message: error?.message || "Operation failed", + }); + } catch (spanError) { + logger.error( + `[${this.INSTRUMENTATION_NAME}] Error recording span for ${spanName} error:`, + spanError, + ); + } + throw error; + }); + }, + ); + }, + spanKind: SpanKind.CLIENT, + }); + } + + /** + * Handle REPLAY mode for BulkOperation.prototype.execute. + * Same pattern as _handleReplayCollectionMethod. + */ + private _handleReplayBulkOpExecute( + original: Function, + thisArg: any, + args: any[], + inputValue: MongodbCommandInputValue, + spanName: string, + submodule: string, + ): any { + const stackTrace = captureStackTrace(["MongodbInstrumentation"]); + + return handleReplayMode({ + noOpRequestHandler: () => { + return Promise.resolve({ + acknowledged: false, + insertedCount: 0, + matchedCount: 0, + modifiedCount: 0, + deletedCount: 0, + upsertedCount: 0, + insertedIds: {}, + upsertedIds: {}, + }); + }, + isServerRequest: false, + replayModeHandler: () => { + return SpanUtils.createAndExecuteSpan( + this.mode, + () => original.apply(thisArg, args), + { + name: spanName, + kind: SpanKind.CLIENT, + submodule, + packageType: PackageType.MONGODB, + packageName: "mongodb", + instrumentationName: this.INSTRUMENTATION_NAME, + inputValue, + isPreAppStart: !this.tuskDrift.isAppReady(), + stopRecordingChildSpans: true, + }, + async (spanInfo: SpanInfo) => { + try { + const mockData = await this._findMockData({ + spanInfo, + name: spanName, + inputValue, + submoduleName: submodule, + stackTrace, + }); + + if (!mockData) { + const errorMsg = `[${this.INSTRUMENTATION_NAME}] No matching mock found for ${spanName} (collection: ${inputValue.collection})`; + logger.warn(errorMsg); + throw new Error(errorMsg); + } + + const result = unwrapDirectOutput( + reconstructBsonValue(mockData.result, this.moduleExports), + ); + + SpanUtils.endSpan(spanInfo.span, { code: SpanStatusCode.OK }); + return result; + } catch (error: any) { + SpanUtils.endSpan(spanInfo.span, { + code: SpanStatusCode.ERROR, + message: error?.message || "Replay failed", + }); + throw error; + } + }, + ); + }, + }); + } + + // --------------------------------------------------------------------------- + // Utilities + // --------------------------------------------------------------------------- + + /** + * Convenience wrapper around the shimmer wrap utility. + */ + private _wrap(target: any, propertyName: string, wrapper: (original: any) => any): void { + wrap(target, propertyName, wrapper); + } +} diff --git a/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/.tusk/config.yaml b/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/.tusk/config.yaml new file mode 100644 index 00000000..a2cc7489 --- /dev/null +++ b/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/.tusk/config.yaml @@ -0,0 +1,31 @@ +version: 1 # version of the config file format + +service: + id: "cjs-mongodb-e2e-test-id" + name: "cjs-mongodb-e2e-test" + port: 3000 + start: + command: "npm run dev" + readiness_check: + command: "curl http://localhost:3000/health" + timeout: 45s + interval: 5s + +tusk_api: + url: "http://localhost:3000" + +test_execution: + concurrent_limit: 10 + batch_size: 10 + timeout: 30s + +comparison: + ignore_fields: + - created_at + +recording: + sampling_rate: 1.0 # 100% + export_spans: false + +replay: + enable_telemetry: false diff --git a/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/Dockerfile b/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/Dockerfile new file mode 100644 index 00000000..4600b342 --- /dev/null +++ b/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/Dockerfile @@ -0,0 +1,29 @@ +FROM node:18 + +WORKDIR /app + +# Copy package files from e2e test directory +COPY src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/package*.json ./ +COPY src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/tsconfig.json ./ + +# Add cache-busting argument to force fresh CLI download +ARG CACHEBUST=1 +ARG TUSK_CLI_VERSION=latest + +# Install Tusk Drift CLI +RUN if [ "$TUSK_CLI_VERSION" = "latest" ]; then \ + curl -fsSL https://raw.githubusercontent.com/Use-Tusk/tusk-drift-cli/main/install.sh | sh; \ + else \ + curl -fsSL https://raw.githubusercontent.com/Use-Tusk/tusk-drift-cli/main/install.sh | sh -s -- ${TUSK_CLI_VERSION}; \ + fi + +# Expose the server port +EXPOSE 3000 + +COPY src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/entrypoint.sh ./ +RUN chmod +x entrypoint.sh +RUN mkdir -p /app/.tusk/traces /app/.tusk/logs +COPY src/instrumentation/libraries/e2e-common/base-entrypoint.sh /app/base-entrypoint.sh +COPY src/instrumentation/libraries/e2e-common/test-utils.mjs /app/test-utils.mjs +RUN chmod +x /app/base-entrypoint.sh +ENTRYPOINT ["./entrypoint.sh"] diff --git a/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/docker-compose.yml b/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/docker-compose.yml new file mode 100644 index 00000000..239c2021 --- /dev/null +++ b/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/docker-compose.yml @@ -0,0 +1,49 @@ +services: + mongo: + image: mongo:7 + command: ["mongod", "--replSet", "rs0", "--bind_ip_all"] + healthcheck: + test: > + mongosh --quiet --eval " + try { + rs.status().ok && 1 + } catch(e) { + rs.initiate({_id: 'rs0', members: [{_id: 0, host: 'mongo:27017'}]}); + sleep(2000); + rs.status().ok && 1 + } + " || exit 1 + interval: 5s + timeout: 10s + retries: 10 + + app: + build: + context: ../../../../../.. + dockerfile: src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/Dockerfile + args: + - CACHEBUST=${CACHEBUST:-1} + - TUSK_CLI_VERSION=${TUSK_CLI_VERSION:-latest} + environment: + - BENCHMARKS=${BENCHMARKS:-} + - BENCHMARK_DURATION=${BENCHMARK_DURATION:-5} + - BENCHMARK_WARMUP=${BENCHMARK_WARMUP:-3} + - PORT=3000 + - MONGO_HOST=mongo + - MONGO_PORT=27017 + - MONGO_DB=testdb + - TUSK_ANALYTICS_DISABLED=1 + volumes: + # Mount SDK source for hot reload (this is what package.json expects) + - ../../../../../..:/sdk:ro + # Mount .tusk config to persist configuration + - ./.tusk/config.yaml:/app/.tusk/config.yaml:ro + # Persist traces and logs on host + - ./.tusk/traces:/app/.tusk/traces + - ./.tusk/logs:/app/.tusk/logs + # Mount app source for development + - ./src:/app/src + working_dir: /app + depends_on: + mongo: + condition: service_healthy diff --git a/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/entrypoint.sh b/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/entrypoint.sh new file mode 100644 index 00000000..ed1dce09 --- /dev/null +++ b/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/entrypoint.sh @@ -0,0 +1,3 @@ +#!/bin/bash +SERVER_WAIT_TIME=8 +source /app/base-entrypoint.sh diff --git a/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/package.json b/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/package.json new file mode 100644 index 00000000..5cb72730 --- /dev/null +++ b/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/package.json @@ -0,0 +1,21 @@ +{ + "name": "cjs-mongodb-e2e-test", + "version": "1.0.0", + "description": "E2E tests for MongoDB instrumentation (CommonJS)", + "main": "dist/index.js", + "type": "commonjs", + "scripts": { + "build": "tsc", + "start": "node dist/index.js", + "dev": "node dist/index.js" + }, + "dependencies": { + "@use-tusk/drift-node-sdk": "file:/sdk", + "mongodb": "^6.0.0", + "mongoose": "^8.0.0" + }, + "devDependencies": { + "@types/node": "^20.0.0", + "typescript": "^5.0.0" + } +} diff --git a/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/run.sh b/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/run.sh new file mode 100755 index 00000000..8d3bcaff --- /dev/null +++ b/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/run.sh @@ -0,0 +1,51 @@ +#!/bin/bash +set -e + +APP_PORT=${1:-3000} +export APP_PORT + +PROJECT_NAME="mongodb-cjs-${APP_PORT}" + +GREEN='\033[0;32m' +RED='\033[0;31m' +YELLOW='\033[1;33m' +BLUE='\033[0;34m' +NC='\033[0m' + +echo -e "${BLUE}========================================${NC}" +echo -e "${BLUE}Running Node E2E Test: MongoDB (CJS)${NC}" +echo -e "${BLUE}Port: ${APP_PORT}${NC}" +echo -e "${BLUE}========================================${NC}" +echo "" + +cleanup() { + echo "" + echo -e "${YELLOW}Cleaning up containers...${NC}" + docker compose -p "$PROJECT_NAME" down -v 2>/dev/null || true +} + +trap cleanup EXIT + +echo -e "${BLUE}Building containers...${NC}" +docker compose -p "$PROJECT_NAME" build --no-cache + +echo -e "${BLUE}Starting test...${NC}" +echo "" + +set +e +docker compose -p "$PROJECT_NAME" run --rm -e TUSK_USE_RUST_CORE="${TUSK_USE_RUST_CORE:-}" app +EXIT_CODE=$? +set -e + +echo "" +if [ $EXIT_CODE -eq 0 ]; then + echo -e "${GREEN}========================================${NC}" + echo -e "${GREEN}Test passed!${NC}" + echo -e "${GREEN}========================================${NC}" +else + echo -e "${RED}========================================${NC}" + echo -e "${RED}Test failed with exit code ${EXIT_CODE}${NC}" + echo -e "${RED}========================================${NC}" +fi + +exit $EXIT_CODE diff --git a/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/src/index.ts b/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/src/index.ts new file mode 100644 index 00000000..c85191bc --- /dev/null +++ b/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/src/index.ts @@ -0,0 +1,629 @@ +import { TuskDrift } from "./tdInit"; +import http from "http"; +import { MongoClient, Db, Collection } from "mongodb"; +import mongoose from "mongoose"; + +const PORT = process.env.PORT || 3000; + +// MongoDB configuration +const mongoHost = process.env.MONGO_HOST || "mongo"; +const mongoPort = process.env.MONGO_PORT || "27017"; +const mongoDb = process.env.MONGO_DB || "testdb"; +const mongoUrl = `mongodb://${mongoHost}:${mongoPort}`; + +let client: MongoClient; +let db: Db; +let testUsers: Collection; +let largeData: Collection; + +async function initializeDatabase() { + console.log(`Connecting to MongoDB: ${mongoUrl}`); + + client = new MongoClient(mongoUrl); + await client.connect(); + db = client.db(mongoDb); + + // Create test_users collection with seed data + testUsers = db.collection("test_users"); + await testUsers.deleteMany({}); + await testUsers.insertMany([ + { name: "John Doe", email: "john@example.com", age: 30, tags: ["admin", "user"] }, + { name: "Jane Smith", email: "jane@example.com", age: 25, tags: ["user"] }, + { name: "Bob Johnson", email: "bob@example.com", age: 35, tags: ["user", "moderator"] }, + ]); + + // Create large_data collection for cursor testing + largeData = db.collection("large_data"); + await largeData.deleteMany({}); + const largeDataDocs = []; + for (let i = 1; i <= 10; i++) { + largeDataDocs.push({ + value: `test_data_${i}`, + index: i, + category: i % 2 === 0 ? "even" : "odd", + }); + } + await largeData.insertMany(largeDataDocs); + + console.log("Database initialized successfully"); +} + +// --- Mongoose Setup --- +const authorSchema = new mongoose.Schema({ + name: { type: String, required: true }, + email: { type: String, required: true }, + bio: String, + createdAt: { type: Date, default: Date.now }, +}); + +const postSchema = new mongoose.Schema({ + title: { type: String, required: true }, + body: { type: String, required: true }, + author: { type: mongoose.Schema.Types.ObjectId, ref: "Author" }, + tags: [String], + views: { type: Number, default: 0 }, + published: { type: Boolean, default: false }, + createdAt: { type: Number, default: () => Date.now() }, +}); + +const Author = mongoose.model("Author", authorSchema); +const Post = mongoose.model("Post", postSchema); + +async function initializeMongoose() { + await mongoose.connect(`mongodb://${mongoHost}:${mongoPort}/${mongoDb}`); + console.log("Mongoose connected successfully"); + + // Seed mongoose data + await Author.deleteMany({}); + await Post.deleteMany({}); + + const author1 = await Author.create({ + name: "Alice Writer", + email: "alice@example.com", + bio: "Tech blogger", + }); + const author2 = await Author.create({ + name: "Bob Author", + email: "bob_author@example.com", + bio: "Fiction writer", + }); + + await Post.create([ + { + title: "First Post", + body: "Hello world content", + author: author1._id, + tags: ["intro", "tech"], + views: 100, + published: true, + }, + { + title: "Second Post", + body: "Advanced topics", + author: author1._id, + tags: ["tech", "advanced"], + views: 250, + published: true, + }, + { + title: "Draft Post", + body: "Work in progress", + author: author2._id, + tags: ["draft"], + views: 0, + published: false, + }, + { + title: "Fiction Story", + body: "Once upon a time", + author: author2._id, + tags: ["fiction", "story"], + views: 50, + published: true, + }, + ]); + + console.log("Mongoose data seeded"); +} + +function sendJson(res: http.ServerResponse, statusCode: number, data: any) { + res.writeHead(statusCode, { "Content-Type": "application/json" }); + res.end(JSON.stringify(data, (_, v) => (typeof v === "bigint" ? v.toString() : v))); +} + +// Create HTTP server with test endpoints +const server = http.createServer(async (req, res) => { + const url = req.url || "/"; + const method = req.method || "GET"; + + try { + // Health check + if (url === "/health" && method === "GET") { + sendJson(res, 200, { success: true }); + return; + } + + // --- insertOne --- + if (url === "/test/insert-one" && method === "GET") { + const tempCol = db.collection("temp_insert_one"); + await tempCol.deleteMany({}); + const result = await tempCol.insertOne({ + name: "Test User", + email: "test@example.com", + age: 28, + }); + await tempCol.drop(); + sendJson(res, 200, { + success: true, + acknowledged: result.acknowledged, + insertedId: result.insertedId, + }); + return; + } + + // --- insertMany --- + if (url === "/test/insert-many" && method === "GET") { + const tempCol = db.collection("temp_insert_many"); + await tempCol.deleteMany({}); + const result = await tempCol.insertMany([ + { name: "User A", email: "a@example.com" }, + { name: "User B", email: "b@example.com" }, + { name: "User C", email: "c@example.com" }, + ]); + await tempCol.drop(); + sendJson(res, 200, { + success: true, + acknowledged: result.acknowledged, + insertedCount: result.insertedCount, + insertedIds: result.insertedIds, + }); + return; + } + + // --- findOne --- + if (url === "/test/find-one" && method === "GET") { + const result = await testUsers.findOne({ name: "John Doe" }); + sendJson(res, 200, { success: true, data: result }); + return; + } + + // --- find (toArray) --- + if (url === "/test/find" && method === "GET") { + const result = await testUsers.find({}).toArray(); + sendJson(res, 200, { success: true, data: result, count: result.length }); + return; + } + + // --- find with sort, limit, skip, project --- + if (url === "/test/find-with-options" && method === "GET") { + const result = await testUsers + .find({}) + .sort({ age: -1 }) + .limit(2) + .skip(0) + .project({ name: 1, age: 1, _id: 0 }) + .toArray(); + sendJson(res, 200, { success: true, data: result }); + return; + } + + // --- find cursor with next/hasNext --- + if (url === "/test/find-cursor-next" && method === "GET") { + const cursor = largeData.find({}).sort({ index: 1 }).limit(3); + const documents: any[] = []; + while (await cursor.hasNext()) { + const doc = await cursor.next(); + if (doc) documents.push(doc); + } + await cursor.close(); + sendJson(res, 200, { success: true, data: documents, count: documents.length }); + return; + } + + // --- updateOne --- + if (url === "/test/update-one" && method === "GET") { + const tempCol = db.collection("temp_update_one"); + await tempCol.deleteMany({}); + await tempCol.insertOne({ name: "Update Target", status: "pending" }); + const result = await tempCol.updateOne( + { name: "Update Target" }, + { $set: { status: "completed" } }, + ); + await tempCol.drop(); + sendJson(res, 200, { + success: true, + acknowledged: result.acknowledged, + matchedCount: result.matchedCount, + modifiedCount: result.modifiedCount, + }); + return; + } + + // --- updateMany --- + if (url === "/test/update-many" && method === "GET") { + const tempCol = db.collection("temp_update_many"); + await tempCol.deleteMany({}); + await tempCol.insertMany([ + { group: "A", status: "pending" }, + { group: "A", status: "pending" }, + { group: "B", status: "pending" }, + ]); + const result = await tempCol.updateMany({ group: "A" }, { $set: { status: "done" } }); + await tempCol.drop(); + sendJson(res, 200, { + success: true, + acknowledged: result.acknowledged, + matchedCount: result.matchedCount, + modifiedCount: result.modifiedCount, + }); + return; + } + + // --- deleteOne --- + if (url === "/test/delete-one" && method === "GET") { + const tempCol = db.collection("temp_delete_one"); + await tempCol.deleteMany({}); + await tempCol.insertOne({ name: "Delete Me" }); + const result = await tempCol.deleteOne({ name: "Delete Me" }); + await tempCol.drop(); + sendJson(res, 200, { + success: true, + acknowledged: result.acknowledged, + deletedCount: result.deletedCount, + }); + return; + } + + // --- deleteMany --- + if (url === "/test/delete-many" && method === "GET") { + const tempCol = db.collection("temp_delete_many"); + await tempCol.deleteMany({}); + await tempCol.insertMany([ + { group: "delete", value: 1 }, + { group: "delete", value: 2 }, + { group: "keep", value: 3 }, + ]); + const result = await tempCol.deleteMany({ group: "delete" }); + await tempCol.drop(); + sendJson(res, 200, { + success: true, + acknowledged: result.acknowledged, + deletedCount: result.deletedCount, + }); + return; + } + + // --- replaceOne --- + if (url === "/test/replace-one" && method === "GET") { + const tempCol = db.collection("temp_replace_one"); + await tempCol.deleteMany({}); + await tempCol.insertOne({ name: "Original", version: 1 }); + const result = await tempCol.replaceOne( + { name: "Original" }, + { name: "Replaced", version: 2, replaced: true }, + ); + await tempCol.drop(); + sendJson(res, 200, { + success: true, + acknowledged: result.acknowledged, + matchedCount: result.matchedCount, + modifiedCount: result.modifiedCount, + }); + return; + } + + // --- findOneAndUpdate --- + if (url === "/test/find-one-and-update" && method === "GET") { + const tempCol = db.collection("temp_find_update"); + await tempCol.deleteMany({}); + await tempCol.insertOne({ name: "FindAndUpdate", counter: 0 }); + const result = await tempCol.findOneAndUpdate( + { name: "FindAndUpdate" }, + { $inc: { counter: 1 } }, + { returnDocument: "after" }, + ); + await tempCol.drop(); + sendJson(res, 200, { success: true, data: result }); + return; + } + + // --- findOneAndDelete --- + if (url === "/test/find-one-and-delete" && method === "GET") { + const tempCol = db.collection("temp_find_delete"); + await tempCol.deleteMany({}); + await tempCol.insertOne({ name: "FindAndDelete", value: 42 }); + const result = await tempCol.findOneAndDelete({ name: "FindAndDelete" }); + await tempCol.drop(); + sendJson(res, 200, { success: true, data: result }); + return; + } + + // --- findOneAndReplace --- + if (url === "/test/find-one-and-replace" && method === "GET") { + const tempCol = db.collection("temp_find_replace"); + await tempCol.deleteMany({}); + await tempCol.insertOne({ name: "FindAndReplace", version: 1 }); + const result = await tempCol.findOneAndReplace( + { name: "FindAndReplace" }, + { name: "Replaced", version: 2 }, + { returnDocument: "after" }, + ); + await tempCol.drop(); + sendJson(res, 200, { success: true, data: result }); + return; + } + + // --- countDocuments --- + if (url === "/test/count-documents" && method === "GET") { + const count = await testUsers.countDocuments({ age: { $gte: 25 } }); + sendJson(res, 200, { success: true, count }); + return; + } + + // --- estimatedDocumentCount --- + if (url === "/test/estimated-count" && method === "GET") { + const count = await testUsers.estimatedDocumentCount(); + sendJson(res, 200, { success: true, count }); + return; + } + + // --- distinct --- + if (url === "/test/distinct" && method === "GET") { + const values = await testUsers.distinct("tags"); + sendJson(res, 200, { success: true, data: values }); + return; + } + + // --- aggregate --- + if (url === "/test/aggregate" && method === "GET") { + const result = await largeData + .aggregate([ + { $match: { index: { $lte: 6 } } }, + { $group: { _id: "$category", total: { $sum: 1 }, avgIndex: { $avg: "$index" } } }, + { $sort: { _id: 1 } }, + ]) + .toArray(); + sendJson(res, 200, { success: true, data: result }); + return; + } + + // --- bulkWrite --- + if (url === "/test/bulk-write" && method === "GET") { + const tempCol = db.collection("temp_bulk_write"); + await tempCol.deleteMany({}); + await tempCol.insertOne({ name: "Existing", value: 1 }); + const result = await tempCol.bulkWrite([ + { insertOne: { document: { name: "BulkInsert1", value: 10 } } }, + { insertOne: { document: { name: "BulkInsert2", value: 20 } } }, + { updateOne: { filter: { name: "Existing" }, update: { $set: { value: 100 } } } }, + { deleteOne: { filter: { name: "BulkInsert1" } } }, + ]); + await tempCol.drop(); + sendJson(res, 200, { + success: true, + ok: result.ok, + insertedCount: result.insertedCount, + matchedCount: result.matchedCount, + modifiedCount: result.modifiedCount, + deletedCount: result.deletedCount, + upsertedCount: result.upsertedCount, + }); + return; + } + + // --- createIndex / createIndexes --- + if (url === "/test/create-index" && method === "GET") { + const tempCol = db.collection("temp_create_index"); + await tempCol.deleteMany({}); + await tempCol.insertOne({ field1: "a", field2: "b" }); + const indexName = await tempCol.createIndex({ field1: 1 }); + const indexNames = await tempCol.createIndexes([{ key: { field2: 1 }, name: "field2_idx" }]); + await tempCol.drop(); + sendJson(res, 200, { success: true, indexName, indexNames }); + return; + } + + // --- dropIndex --- + if (url === "/test/drop-index" && method === "GET") { + const tempCol = db.collection("temp_drop_index"); + await tempCol.deleteMany({}); + await tempCol.insertOne({ field1: "a" }); + await tempCol.createIndex({ field1: 1 }, { name: "field1_drop_idx" }); + await tempCol.dropIndex("field1_drop_idx"); + sendJson(res, 200, { success: true }); + return; + } + + // --- listIndexes --- + if (url === "/test/list-indexes" && method === "GET") { + const indexes = await testUsers.listIndexes().toArray(); + sendJson(res, 200, { success: true, data: indexes }); + return; + } + + // --- db.command --- + if (url === "/test/db-command" && method === "GET") { + const result = await db.command({ ping: 1 }); + sendJson(res, 200, { success: true, data: result }); + return; + } + + // --- listCollections --- + if (url === "/test/list-collections" && method === "GET") { + const collections = await db.listCollections().toArray(); + sendJson(res, 200, { success: true, data: collections }); + return; + } + + // --- transaction (session) --- + if (url === "/test/transaction" && method === "GET") { + const session = client.startSession(); + try { + session.startTransaction(); + const txnCol = db.collection("temp_transaction"); + await txnCol.deleteMany({}, { session }); + await txnCol.insertOne({ name: "TxnUser", value: 1 }, { session }); + const found = await txnCol.findOne({ name: "TxnUser" }, { session }); + await txnCol.updateOne({ name: "TxnUser" }, { $set: { value: 2 } }, { session }); + await session.commitTransaction(); + await txnCol.drop(); + sendJson(res, 200, { success: true, data: found }); + } catch (error) { + await session.abortTransaction(); + throw error; + } finally { + await session.endSession(); + } + return; + } + + // --- ordered bulk operation --- + if (url === "/test/ordered-bulk" && method === "GET") { + const tempCol = db.collection("temp_ordered_bulk"); + await tempCol.deleteMany({}); + const bulk = tempCol.initializeOrderedBulkOp(); + bulk.insert({ name: "OrderedBulk1", value: 1 }); + bulk.insert({ name: "OrderedBulk2", value: 2 }); + bulk.insert({ name: "OrderedBulk3", value: 3 }); + bulk.find({ name: "OrderedBulk1" }).updateOne({ $set: { value: 10 } }); + bulk.find({ name: "OrderedBulk3" }).deleteOne(); + const result = await bulk.execute(); + await tempCol.drop(); + sendJson(res, 200, { + success: true, + ok: result.ok, + insertedCount: result.insertedCount, + modifiedCount: result.modifiedCount, + deletedCount: result.deletedCount, + }); + return; + } + + // --- unordered bulk operation --- + if (url === "/test/unordered-bulk" && method === "GET") { + const tempCol = db.collection("temp_unordered_bulk"); + await tempCol.deleteMany({}); + const bulk = tempCol.initializeUnorderedBulkOp(); + bulk.insert({ name: "UnorderedBulk1", value: 1 }); + bulk.insert({ name: "UnorderedBulk2", value: 2 }); + bulk.insert({ name: "UnorderedBulk3", value: 3 }); + const result = await bulk.execute(); + await tempCol.drop(); + sendJson(res, 200, { + success: true, + ok: result.ok, + insertedCount: result.insertedCount, + }); + return; + } + + // --- Mongoose: create --- + if (url === "/test/mongoose-create" && method === "GET") { + const tmpAuthor = await Author.create({ + name: "Temp Author", + email: "temp@example.com", + bio: "Temporary", + }); + const result = { id: tmpAuthor._id, name: tmpAuthor.name, email: tmpAuthor.email }; + await Author.deleteOne({ _id: tmpAuthor._id }); + sendJson(res, 200, { success: true, data: result }); + return; + } + + // --- Mongoose: create multiple docs (insertMany path) --- + if (url === "/test/mongoose-create-many" && method === "GET") { + const created = await Post.create([ + { + title: "TempMulti1", + body: "Multi body 1", + tags: ["temp-multi"], + views: 0, + published: false, + }, + { + title: "TempMulti2", + body: "Multi body 2", + tags: ["temp-multi"], + views: 0, + published: false, + }, + ]); + const result = created.map((d) => ({ title: d.title, id: d._id })); + await Post.deleteMany({ tags: "temp-multi" }); + sendJson(res, 200, { success: true, data: result, count: result.length }); + return; + } + + // --- Raw MongoDB: cursor map transform --- + if (url === "/test/cursor-map" && method === "GET") { + const mapped = await largeData + .find({ category: "even" }) + .map((doc) => ({ value: doc.value, doubled: doc.index * 2 })) + .toArray(); + sendJson(res, 200, { success: true, data: mapped }); + return; + } + + // --- Raw MongoDB: cursor async iterator (for await...of) --- + if (url === "/test/cursor-async-iterator" && method === "GET") { + const items: any[] = []; + const cursor = largeData.find({ category: "even" }).sort({ index: 1 }); + for await (const doc of cursor) { + items.push({ value: doc.value, index: doc.index }); + } + sendJson(res, 200, { success: true, data: items }); + return; + } + + // 404 for unknown routes + sendJson(res, 404, { error: "Not found" }); + } catch (error) { + console.error("Error handling request:", error); + sendJson(res, 500, { + success: false, + error: error instanceof Error ? error.message : String(error), + }); + } +}); + +// Initialize database first, then start server +async function main() { + await initializeDatabase(); + await initializeMongoose(); + server.listen(PORT, () => { + TuskDrift.markAppAsReady(); + console.log(`MongoDB integration test server running on port ${PORT}`); + console.log(`Test mode: ${process.env.TUSK_DRIFT_MODE}`); + }); +} + +main().catch((error) => { + console.error("Failed to start server:", error); + process.exit(1); +}); + +// Graceful shutdown +async function shutdown() { + console.log("Shutting down gracefully..."); + try { + await mongoose.disconnect(); + await client.close(); + } catch (error) { + console.error("Error during shutdown:", error); + } + process.exit(0); +} + +process.on("SIGTERM", shutdown); +process.on("SIGINT", shutdown); + +// Handle uncaught exceptions +process.on("uncaughtException", async (error) => { + console.error("Uncaught exception:", error); + await shutdown(); +}); + +process.on("unhandledRejection", async (reason, promise) => { + console.error("Unhandled rejection at:", promise, "reason:", reason); + await shutdown(); +}); diff --git a/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/src/tdInit.ts b/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/src/tdInit.ts new file mode 100644 index 00000000..9b1acad7 --- /dev/null +++ b/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/src/tdInit.ts @@ -0,0 +1,8 @@ +import { TuskDrift } from '@use-tusk/drift-node-sdk'; + +TuskDrift.initialize({ + apiKey: "api-key", + env: "dev", +}); + +export { TuskDrift }; diff --git a/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/src/test_requests.mjs b/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/src/test_requests.mjs new file mode 100644 index 00000000..b4dac25c --- /dev/null +++ b/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/src/test_requests.mjs @@ -0,0 +1,36 @@ +import { makeRequest, printRequestSummary } from "/app/test-utils.mjs"; + +await makeRequest("GET", "/health"); +await makeRequest("GET", "/test/insert-one"); +await makeRequest("GET", "/test/insert-many"); +await makeRequest("GET", "/test/find-one"); +await makeRequest("GET", "/test/find"); +await makeRequest("GET", "/test/find-with-options"); +await makeRequest("GET", "/test/find-cursor-next"); +await makeRequest("GET", "/test/update-one"); +await makeRequest("GET", "/test/update-many"); +await makeRequest("GET", "/test/delete-one"); +await makeRequest("GET", "/test/delete-many"); +await makeRequest("GET", "/test/replace-one"); +await makeRequest("GET", "/test/find-one-and-update"); +await makeRequest("GET", "/test/find-one-and-delete"); +await makeRequest("GET", "/test/find-one-and-replace"); +await makeRequest("GET", "/test/count-documents"); +await makeRequest("GET", "/test/estimated-count"); +await makeRequest("GET", "/test/distinct"); +await makeRequest("GET", "/test/aggregate"); +await makeRequest("GET", "/test/bulk-write"); +await makeRequest("GET", "/test/create-index"); +await makeRequest("GET", "/test/drop-index"); +await makeRequest("GET", "/test/list-indexes"); +await makeRequest("GET", "/test/db-command"); +await makeRequest("GET", "/test/list-collections"); +await makeRequest("GET", "/test/transaction"); +await makeRequest("GET", "/test/ordered-bulk"); +await makeRequest("GET", "/test/unordered-bulk"); +await makeRequest("GET", "/test/mongoose-create"); +await makeRequest("GET", "/test/mongoose-create-many"); +await makeRequest("GET", "/test/cursor-map"); +await makeRequest("GET", "/test/cursor-async-iterator"); + +printRequestSummary(); diff --git a/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/tsconfig.json b/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/tsconfig.json new file mode 100644 index 00000000..9b415ad8 --- /dev/null +++ b/src/instrumentation/libraries/mongodb/e2e-tests/cjs-mongodb/tsconfig.json @@ -0,0 +1,17 @@ +{ + "compilerOptions": { + "target": "ES2020", + "module": "commonjs", + "lib": ["ES2020"], + "outDir": "./dist", + "rootDir": "./src", + "strict": true, + "esModuleInterop": true, + "skipLibCheck": true, + "forceConsistentCasingInFileNames": true, + "resolveJsonModule": true, + "moduleResolution": "node" + }, + "include": ["src/**/*"], + "exclude": ["node_modules", "dist"] +} diff --git a/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/.tusk/config.yaml b/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/.tusk/config.yaml new file mode 100644 index 00000000..e57870e5 --- /dev/null +++ b/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/.tusk/config.yaml @@ -0,0 +1,31 @@ +version: 1 # version of the config file format + +service: + id: "esm-mongodb-e2e-test-id" + name: "esm-mongodb-e2e-test" + port: 3000 + start: + command: "npm run dev" + readiness_check: + command: "curl http://localhost:3000/health" + timeout: 45s + interval: 5s + +tusk_api: + url: "http://localhost:3000" + +test_execution: + concurrent_limit: 10 + batch_size: 10 + timeout: 30s + +comparison: + ignore_fields: + - created_at + +recording: + sampling_rate: 1.0 # 100% + export_spans: false + +replay: + enable_telemetry: false diff --git a/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/Dockerfile b/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/Dockerfile new file mode 100644 index 00000000..a62ab87a --- /dev/null +++ b/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/Dockerfile @@ -0,0 +1,29 @@ +FROM node:18 + +WORKDIR /app + +# Copy package files from e2e test directory +COPY src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/package*.json ./ +COPY src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/tsconfig.json ./ + +# Add cache-busting argument to force fresh CLI download +ARG CACHEBUST=1 +ARG TUSK_CLI_VERSION=latest + +# Install Tusk Drift CLI +RUN if [ "$TUSK_CLI_VERSION" = "latest" ]; then \ + curl -fsSL https://raw.githubusercontent.com/Use-Tusk/tusk-drift-cli/main/install.sh | sh; \ + else \ + curl -fsSL https://raw.githubusercontent.com/Use-Tusk/tusk-drift-cli/main/install.sh | sh -s -- ${TUSK_CLI_VERSION}; \ + fi + +# Expose the server port +EXPOSE 3000 + +COPY src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/entrypoint.sh ./ +RUN chmod +x entrypoint.sh +RUN mkdir -p /app/.tusk/traces /app/.tusk/logs +COPY src/instrumentation/libraries/e2e-common/base-entrypoint.sh /app/base-entrypoint.sh +COPY src/instrumentation/libraries/e2e-common/test-utils.mjs /app/test-utils.mjs +RUN chmod +x /app/base-entrypoint.sh +ENTRYPOINT ["./entrypoint.sh"] diff --git a/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/docker-compose.yml b/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/docker-compose.yml new file mode 100644 index 00000000..9fdfebf7 --- /dev/null +++ b/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/docker-compose.yml @@ -0,0 +1,49 @@ +services: + mongo: + image: mongo:7 + command: ["mongod", "--replSet", "rs0", "--bind_ip_all"] + healthcheck: + test: > + mongosh --quiet --eval " + try { + rs.status().ok && 1 + } catch(e) { + rs.initiate({_id: 'rs0', members: [{_id: 0, host: 'mongo:27017'}]}); + sleep(2000); + rs.status().ok && 1 + } + " || exit 1 + interval: 5s + timeout: 10s + retries: 10 + + app: + build: + context: ../../../../../.. + dockerfile: src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/Dockerfile + args: + - CACHEBUST=${CACHEBUST:-1} + - TUSK_CLI_VERSION=${TUSK_CLI_VERSION:-latest} + environment: + - BENCHMARKS=${BENCHMARKS:-} + - BENCHMARK_DURATION=${BENCHMARK_DURATION:-5} + - BENCHMARK_WARMUP=${BENCHMARK_WARMUP:-3} + - PORT=3000 + - MONGO_HOST=mongo + - MONGO_PORT=27017 + - MONGO_DB=testdb + - TUSK_ANALYTICS_DISABLED=1 + volumes: + # Mount SDK source for hot reload (this is what package.json expects) + - ../../../../../..:/sdk:ro + # Mount .tusk config to persist configuration + - ./.tusk/config.yaml:/app/.tusk/config.yaml:ro + # Persist traces and logs on host + - ./.tusk/traces:/app/.tusk/traces + - ./.tusk/logs:/app/.tusk/logs + # Mount app source for development + - ./src:/app/src + working_dir: /app + depends_on: + mongo: + condition: service_healthy diff --git a/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/entrypoint.sh b/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/entrypoint.sh new file mode 100755 index 00000000..ed1dce09 --- /dev/null +++ b/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/entrypoint.sh @@ -0,0 +1,3 @@ +#!/bin/bash +SERVER_WAIT_TIME=8 +source /app/base-entrypoint.sh diff --git a/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/package.json b/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/package.json new file mode 100644 index 00000000..ac7da54a --- /dev/null +++ b/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/package.json @@ -0,0 +1,21 @@ +{ + "name": "esm-mongodb-e2e-test", + "version": "1.0.0", + "description": "E2E tests for MongoDB instrumentation (ESM)", + "main": "dist/index.js", + "type": "module", + "scripts": { + "build": "tsc", + "start": "node --import ./dist/tdInit.js dist/index.js", + "dev": "node --import ./dist/tdInit.js dist/index.js" + }, + "dependencies": { + "@use-tusk/drift-node-sdk": "file:/sdk", + "mongodb": "^6.0.0", + "mongoose": "^8.0.0" + }, + "devDependencies": { + "@types/node": "^20.0.0", + "typescript": "^5.0.0" + } +} diff --git a/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/run.sh b/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/run.sh new file mode 100755 index 00000000..9f2cffd6 --- /dev/null +++ b/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/run.sh @@ -0,0 +1,51 @@ +#!/bin/bash +set -e + +APP_PORT=${1:-3000} +export APP_PORT + +PROJECT_NAME="mongodb-esm-${APP_PORT}" + +GREEN='\033[0;32m' +RED='\033[0;31m' +YELLOW='\033[1;33m' +BLUE='\033[0;34m' +NC='\033[0m' + +echo -e "${BLUE}========================================${NC}" +echo -e "${BLUE}Running Node E2E Test: MongoDB (ESM)${NC}" +echo -e "${BLUE}Port: ${APP_PORT}${NC}" +echo -e "${BLUE}========================================${NC}" +echo "" + +cleanup() { + echo "" + echo -e "${YELLOW}Cleaning up containers...${NC}" + docker compose -p "$PROJECT_NAME" down -v 2>/dev/null || true +} + +trap cleanup EXIT + +echo -e "${BLUE}Building containers...${NC}" +docker compose -p "$PROJECT_NAME" build --no-cache + +echo -e "${BLUE}Starting test...${NC}" +echo "" + +set +e +docker compose -p "$PROJECT_NAME" run --rm -e TUSK_USE_RUST_CORE="${TUSK_USE_RUST_CORE:-}" app +EXIT_CODE=$? +set -e + +echo "" +if [ $EXIT_CODE -eq 0 ]; then + echo -e "${GREEN}========================================${NC}" + echo -e "${GREEN}Test passed!${NC}" + echo -e "${GREEN}========================================${NC}" +else + echo -e "${RED}========================================${NC}" + echo -e "${RED}Test failed with exit code ${EXIT_CODE}${NC}" + echo -e "${RED}========================================${NC}" +fi + +exit $EXIT_CODE diff --git a/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/src/index.ts b/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/src/index.ts new file mode 100644 index 00000000..eee30b3b --- /dev/null +++ b/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/src/index.ts @@ -0,0 +1,629 @@ +import { TuskDrift } from "./tdInit.js"; +import http from "http"; +import { MongoClient, Db, Collection } from "mongodb"; +import mongoose from "mongoose"; + +const PORT = process.env.PORT || 3000; + +// MongoDB configuration +const mongoHost = process.env.MONGO_HOST || "mongo"; +const mongoPort = process.env.MONGO_PORT || "27017"; +const mongoDb = process.env.MONGO_DB || "testdb"; +const mongoUrl = `mongodb://${mongoHost}:${mongoPort}`; + +let client: MongoClient; +let db: Db; +let testUsers: Collection; +let largeData: Collection; + +async function initializeDatabase() { + console.log(`Connecting to MongoDB: ${mongoUrl}`); + + client = new MongoClient(mongoUrl); + await client.connect(); + db = client.db(mongoDb); + + // Create test_users collection with seed data + testUsers = db.collection("test_users"); + await testUsers.deleteMany({}); + await testUsers.insertMany([ + { name: "John Doe", email: "john@example.com", age: 30, tags: ["admin", "user"] }, + { name: "Jane Smith", email: "jane@example.com", age: 25, tags: ["user"] }, + { name: "Bob Johnson", email: "bob@example.com", age: 35, tags: ["user", "moderator"] }, + ]); + + // Create large_data collection for cursor testing + largeData = db.collection("large_data"); + await largeData.deleteMany({}); + const largeDataDocs = []; + for (let i = 1; i <= 10; i++) { + largeDataDocs.push({ + value: `test_data_${i}`, + index: i, + category: i % 2 === 0 ? "even" : "odd", + }); + } + await largeData.insertMany(largeDataDocs); + + console.log("Database initialized successfully"); +} + +// --- Mongoose Setup --- +const authorSchema = new mongoose.Schema({ + name: { type: String, required: true }, + email: { type: String, required: true }, + bio: String, + createdAt: { type: Date, default: Date.now }, +}); + +const postSchema = new mongoose.Schema({ + title: { type: String, required: true }, + body: { type: String, required: true }, + author: { type: mongoose.Schema.Types.ObjectId, ref: "Author" }, + tags: [String], + views: { type: Number, default: 0 }, + published: { type: Boolean, default: false }, + createdAt: { type: Number, default: () => Date.now() }, +}); + +const Author = mongoose.model("Author", authorSchema); +const Post = mongoose.model("Post", postSchema); + +async function initializeMongoose() { + await mongoose.connect(`mongodb://${mongoHost}:${mongoPort}/${mongoDb}`); + console.log("Mongoose connected successfully"); + + // Seed mongoose data + await Author.deleteMany({}); + await Post.deleteMany({}); + + const author1 = await Author.create({ + name: "Alice Writer", + email: "alice@example.com", + bio: "Tech blogger", + }); + const author2 = await Author.create({ + name: "Bob Author", + email: "bob_author@example.com", + bio: "Fiction writer", + }); + + await Post.create([ + { + title: "First Post", + body: "Hello world content", + author: author1._id, + tags: ["intro", "tech"], + views: 100, + published: true, + }, + { + title: "Second Post", + body: "Advanced topics", + author: author1._id, + tags: ["tech", "advanced"], + views: 250, + published: true, + }, + { + title: "Draft Post", + body: "Work in progress", + author: author2._id, + tags: ["draft"], + views: 0, + published: false, + }, + { + title: "Fiction Story", + body: "Once upon a time", + author: author2._id, + tags: ["fiction", "story"], + views: 50, + published: true, + }, + ]); + + console.log("Mongoose data seeded"); +} + +function sendJson(res: http.ServerResponse, statusCode: number, data: any) { + res.writeHead(statusCode, { "Content-Type": "application/json" }); + res.end(JSON.stringify(data, (_, v) => (typeof v === "bigint" ? v.toString() : v))); +} + +// Create HTTP server with test endpoints +const server = http.createServer(async (req, res) => { + const url = req.url || "/"; + const method = req.method || "GET"; + + try { + // Health check + if (url === "/health" && method === "GET") { + sendJson(res, 200, { success: true }); + return; + } + + // --- insertOne --- + if (url === "/test/insert-one" && method === "GET") { + const tempCol = db.collection("temp_insert_one"); + await tempCol.deleteMany({}); + const result = await tempCol.insertOne({ + name: "Test User", + email: "test@example.com", + age: 28, + }); + await tempCol.drop(); + sendJson(res, 200, { + success: true, + acknowledged: result.acknowledged, + insertedId: result.insertedId, + }); + return; + } + + // --- insertMany --- + if (url === "/test/insert-many" && method === "GET") { + const tempCol = db.collection("temp_insert_many"); + await tempCol.deleteMany({}); + const result = await tempCol.insertMany([ + { name: "User A", email: "a@example.com" }, + { name: "User B", email: "b@example.com" }, + { name: "User C", email: "c@example.com" }, + ]); + await tempCol.drop(); + sendJson(res, 200, { + success: true, + acknowledged: result.acknowledged, + insertedCount: result.insertedCount, + insertedIds: result.insertedIds, + }); + return; + } + + // --- findOne --- + if (url === "/test/find-one" && method === "GET") { + const result = await testUsers.findOne({ name: "John Doe" }); + sendJson(res, 200, { success: true, data: result }); + return; + } + + // --- find (toArray) --- + if (url === "/test/find" && method === "GET") { + const result = await testUsers.find({}).toArray(); + sendJson(res, 200, { success: true, data: result, count: result.length }); + return; + } + + // --- find with sort, limit, skip, project --- + if (url === "/test/find-with-options" && method === "GET") { + const result = await testUsers + .find({}) + .sort({ age: -1 }) + .limit(2) + .skip(0) + .project({ name: 1, age: 1, _id: 0 }) + .toArray(); + sendJson(res, 200, { success: true, data: result }); + return; + } + + // --- find cursor with next/hasNext --- + if (url === "/test/find-cursor-next" && method === "GET") { + const cursor = largeData.find({}).sort({ index: 1 }).limit(3); + const documents: any[] = []; + while (await cursor.hasNext()) { + const doc = await cursor.next(); + if (doc) documents.push(doc); + } + await cursor.close(); + sendJson(res, 200, { success: true, data: documents, count: documents.length }); + return; + } + + // --- updateOne --- + if (url === "/test/update-one" && method === "GET") { + const tempCol = db.collection("temp_update_one"); + await tempCol.deleteMany({}); + await tempCol.insertOne({ name: "Update Target", status: "pending" }); + const result = await tempCol.updateOne( + { name: "Update Target" }, + { $set: { status: "completed" } }, + ); + await tempCol.drop(); + sendJson(res, 200, { + success: true, + acknowledged: result.acknowledged, + matchedCount: result.matchedCount, + modifiedCount: result.modifiedCount, + }); + return; + } + + // --- updateMany --- + if (url === "/test/update-many" && method === "GET") { + const tempCol = db.collection("temp_update_many"); + await tempCol.deleteMany({}); + await tempCol.insertMany([ + { group: "A", status: "pending" }, + { group: "A", status: "pending" }, + { group: "B", status: "pending" }, + ]); + const result = await tempCol.updateMany({ group: "A" }, { $set: { status: "done" } }); + await tempCol.drop(); + sendJson(res, 200, { + success: true, + acknowledged: result.acknowledged, + matchedCount: result.matchedCount, + modifiedCount: result.modifiedCount, + }); + return; + } + + // --- deleteOne --- + if (url === "/test/delete-one" && method === "GET") { + const tempCol = db.collection("temp_delete_one"); + await tempCol.deleteMany({}); + await tempCol.insertOne({ name: "Delete Me" }); + const result = await tempCol.deleteOne({ name: "Delete Me" }); + await tempCol.drop(); + sendJson(res, 200, { + success: true, + acknowledged: result.acknowledged, + deletedCount: result.deletedCount, + }); + return; + } + + // --- deleteMany --- + if (url === "/test/delete-many" && method === "GET") { + const tempCol = db.collection("temp_delete_many"); + await tempCol.deleteMany({}); + await tempCol.insertMany([ + { group: "delete", value: 1 }, + { group: "delete", value: 2 }, + { group: "keep", value: 3 }, + ]); + const result = await tempCol.deleteMany({ group: "delete" }); + await tempCol.drop(); + sendJson(res, 200, { + success: true, + acknowledged: result.acknowledged, + deletedCount: result.deletedCount, + }); + return; + } + + // --- replaceOne --- + if (url === "/test/replace-one" && method === "GET") { + const tempCol = db.collection("temp_replace_one"); + await tempCol.deleteMany({}); + await tempCol.insertOne({ name: "Original", version: 1 }); + const result = await tempCol.replaceOne( + { name: "Original" }, + { name: "Replaced", version: 2, replaced: true }, + ); + await tempCol.drop(); + sendJson(res, 200, { + success: true, + acknowledged: result.acknowledged, + matchedCount: result.matchedCount, + modifiedCount: result.modifiedCount, + }); + return; + } + + // --- findOneAndUpdate --- + if (url === "/test/find-one-and-update" && method === "GET") { + const tempCol = db.collection("temp_find_update"); + await tempCol.deleteMany({}); + await tempCol.insertOne({ name: "FindAndUpdate", counter: 0 }); + const result = await tempCol.findOneAndUpdate( + { name: "FindAndUpdate" }, + { $inc: { counter: 1 } }, + { returnDocument: "after" }, + ); + await tempCol.drop(); + sendJson(res, 200, { success: true, data: result }); + return; + } + + // --- findOneAndDelete --- + if (url === "/test/find-one-and-delete" && method === "GET") { + const tempCol = db.collection("temp_find_delete"); + await tempCol.deleteMany({}); + await tempCol.insertOne({ name: "FindAndDelete", value: 42 }); + const result = await tempCol.findOneAndDelete({ name: "FindAndDelete" }); + await tempCol.drop(); + sendJson(res, 200, { success: true, data: result }); + return; + } + + // --- findOneAndReplace --- + if (url === "/test/find-one-and-replace" && method === "GET") { + const tempCol = db.collection("temp_find_replace"); + await tempCol.deleteMany({}); + await tempCol.insertOne({ name: "FindAndReplace", version: 1 }); + const result = await tempCol.findOneAndReplace( + { name: "FindAndReplace" }, + { name: "Replaced", version: 2 }, + { returnDocument: "after" }, + ); + await tempCol.drop(); + sendJson(res, 200, { success: true, data: result }); + return; + } + + // --- countDocuments --- + if (url === "/test/count-documents" && method === "GET") { + const count = await testUsers.countDocuments({ age: { $gte: 25 } }); + sendJson(res, 200, { success: true, count }); + return; + } + + // --- estimatedDocumentCount --- + if (url === "/test/estimated-count" && method === "GET") { + const count = await testUsers.estimatedDocumentCount(); + sendJson(res, 200, { success: true, count }); + return; + } + + // --- distinct --- + if (url === "/test/distinct" && method === "GET") { + const values = await testUsers.distinct("tags"); + sendJson(res, 200, { success: true, data: values }); + return; + } + + // --- aggregate --- + if (url === "/test/aggregate" && method === "GET") { + const result = await largeData + .aggregate([ + { $match: { index: { $lte: 6 } } }, + { $group: { _id: "$category", total: { $sum: 1 }, avgIndex: { $avg: "$index" } } }, + { $sort: { _id: 1 } }, + ]) + .toArray(); + sendJson(res, 200, { success: true, data: result }); + return; + } + + // --- bulkWrite --- + if (url === "/test/bulk-write" && method === "GET") { + const tempCol = db.collection("temp_bulk_write"); + await tempCol.deleteMany({}); + await tempCol.insertOne({ name: "Existing", value: 1 }); + const result = await tempCol.bulkWrite([ + { insertOne: { document: { name: "BulkInsert1", value: 10 } } }, + { insertOne: { document: { name: "BulkInsert2", value: 20 } } }, + { updateOne: { filter: { name: "Existing" }, update: { $set: { value: 100 } } } }, + { deleteOne: { filter: { name: "BulkInsert1" } } }, + ]); + await tempCol.drop(); + sendJson(res, 200, { + success: true, + ok: result.ok, + insertedCount: result.insertedCount, + matchedCount: result.matchedCount, + modifiedCount: result.modifiedCount, + deletedCount: result.deletedCount, + upsertedCount: result.upsertedCount, + }); + return; + } + + // --- createIndex / createIndexes --- + if (url === "/test/create-index" && method === "GET") { + const tempCol = db.collection("temp_create_index"); + await tempCol.deleteMany({}); + await tempCol.insertOne({ field1: "a", field2: "b" }); + const indexName = await tempCol.createIndex({ field1: 1 }); + const indexNames = await tempCol.createIndexes([{ key: { field2: 1 }, name: "field2_idx" }]); + await tempCol.drop(); + sendJson(res, 200, { success: true, indexName, indexNames }); + return; + } + + // --- dropIndex --- + if (url === "/test/drop-index" && method === "GET") { + const tempCol = db.collection("temp_drop_index"); + await tempCol.deleteMany({}); + await tempCol.insertOne({ field1: "a" }); + await tempCol.createIndex({ field1: 1 }, { name: "field1_drop_idx" }); + await tempCol.dropIndex("field1_drop_idx"); + sendJson(res, 200, { success: true }); + return; + } + + // --- listIndexes --- + if (url === "/test/list-indexes" && method === "GET") { + const indexes = await testUsers.listIndexes().toArray(); + sendJson(res, 200, { success: true, data: indexes }); + return; + } + + // --- db.command --- + if (url === "/test/db-command" && method === "GET") { + const result = await db.command({ ping: 1 }); + sendJson(res, 200, { success: true, data: result }); + return; + } + + // --- listCollections --- + if (url === "/test/list-collections" && method === "GET") { + const collections = await db.listCollections().toArray(); + sendJson(res, 200, { success: true, data: collections }); + return; + } + + // --- transaction (session) --- + if (url === "/test/transaction" && method === "GET") { + const session = client.startSession(); + try { + session.startTransaction(); + const txnCol = db.collection("temp_transaction"); + await txnCol.deleteMany({}, { session }); + await txnCol.insertOne({ name: "TxnUser", value: 1 }, { session }); + const found = await txnCol.findOne({ name: "TxnUser" }, { session }); + await txnCol.updateOne({ name: "TxnUser" }, { $set: { value: 2 } }, { session }); + await session.commitTransaction(); + await txnCol.drop(); + sendJson(res, 200, { success: true, data: found }); + } catch (error) { + await session.abortTransaction(); + throw error; + } finally { + await session.endSession(); + } + return; + } + + // --- ordered bulk operation --- + if (url === "/test/ordered-bulk" && method === "GET") { + const tempCol = db.collection("temp_ordered_bulk"); + await tempCol.deleteMany({}); + const bulk = tempCol.initializeOrderedBulkOp(); + bulk.insert({ name: "OrderedBulk1", value: 1 }); + bulk.insert({ name: "OrderedBulk2", value: 2 }); + bulk.insert({ name: "OrderedBulk3", value: 3 }); + bulk.find({ name: "OrderedBulk1" }).updateOne({ $set: { value: 10 } }); + bulk.find({ name: "OrderedBulk3" }).deleteOne(); + const result = await bulk.execute(); + await tempCol.drop(); + sendJson(res, 200, { + success: true, + ok: result.ok, + insertedCount: result.insertedCount, + modifiedCount: result.modifiedCount, + deletedCount: result.deletedCount, + }); + return; + } + + // --- unordered bulk operation --- + if (url === "/test/unordered-bulk" && method === "GET") { + const tempCol = db.collection("temp_unordered_bulk"); + await tempCol.deleteMany({}); + const bulk = tempCol.initializeUnorderedBulkOp(); + bulk.insert({ name: "UnorderedBulk1", value: 1 }); + bulk.insert({ name: "UnorderedBulk2", value: 2 }); + bulk.insert({ name: "UnorderedBulk3", value: 3 }); + const result = await bulk.execute(); + await tempCol.drop(); + sendJson(res, 200, { + success: true, + ok: result.ok, + insertedCount: result.insertedCount, + }); + return; + } + + // --- Mongoose: create --- + if (url === "/test/mongoose-create" && method === "GET") { + const tmpAuthor = await Author.create({ + name: "Temp Author", + email: "temp@example.com", + bio: "Temporary", + }); + const result = { id: tmpAuthor._id, name: tmpAuthor.name, email: tmpAuthor.email }; + await Author.deleteOne({ _id: tmpAuthor._id }); + sendJson(res, 200, { success: true, data: result }); + return; + } + + // --- Mongoose: create multiple docs (insertMany path) --- + if (url === "/test/mongoose-create-many" && method === "GET") { + const created = await Post.create([ + { + title: "TempMulti1", + body: "Multi body 1", + tags: ["temp-multi"], + views: 0, + published: false, + }, + { + title: "TempMulti2", + body: "Multi body 2", + tags: ["temp-multi"], + views: 0, + published: false, + }, + ]); + const result = created.map((d) => ({ title: d.title, id: d._id })); + await Post.deleteMany({ tags: "temp-multi" }); + sendJson(res, 200, { success: true, data: result, count: result.length }); + return; + } + + // --- Raw MongoDB: cursor map transform --- + if (url === "/test/cursor-map" && method === "GET") { + const mapped = await largeData + .find({ category: "even" }) + .map((doc) => ({ value: doc.value, doubled: doc.index * 2 })) + .toArray(); + sendJson(res, 200, { success: true, data: mapped }); + return; + } + + // --- Raw MongoDB: cursor async iterator (for await...of) --- + if (url === "/test/cursor-async-iterator" && method === "GET") { + const items: any[] = []; + const cursor = largeData.find({ category: "even" }).sort({ index: 1 }); + for await (const doc of cursor) { + items.push({ value: doc.value, index: doc.index }); + } + sendJson(res, 200, { success: true, data: items }); + return; + } + + // 404 for unknown routes + sendJson(res, 404, { error: "Not found" }); + } catch (error) { + console.error("Error handling request:", error); + sendJson(res, 500, { + success: false, + error: error instanceof Error ? error.message : String(error), + }); + } +}); + +// Initialize database first, then start server +async function main() { + await initializeDatabase(); + await initializeMongoose(); + server.listen(PORT, () => { + TuskDrift.markAppAsReady(); + console.log(`MongoDB integration test server running on port ${PORT}`); + console.log(`Test mode: ${process.env.TUSK_DRIFT_MODE}`); + }); +} + +main().catch((error) => { + console.error("Failed to start server:", error); + process.exit(1); +}); + +// Graceful shutdown +async function shutdown() { + console.log("Shutting down gracefully..."); + try { + await mongoose.disconnect(); + await client.close(); + } catch (error) { + console.error("Error during shutdown:", error); + } + process.exit(0); +} + +process.on("SIGTERM", shutdown); +process.on("SIGINT", shutdown); + +// Handle uncaught exceptions +process.on("uncaughtException", async (error) => { + console.error("Uncaught exception:", error); + await shutdown(); +}); + +process.on("unhandledRejection", async (reason, promise) => { + console.error("Unhandled rejection at:", promise, "reason:", reason); + await shutdown(); +}); diff --git a/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/src/tdInit.ts b/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/src/tdInit.ts new file mode 100644 index 00000000..d01058a3 --- /dev/null +++ b/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/src/tdInit.ts @@ -0,0 +1,15 @@ +import { register } from "node:module"; +import { pathToFileURL } from "node:url"; + +// Register the ESM loader +// This enables interception of ESM module imports +register("@use-tusk/drift-node-sdk/hook.mjs", pathToFileURL("./")); + +import { TuskDrift } from "@use-tusk/drift-node-sdk"; + +TuskDrift.initialize({ + apiKey: "api-key", + env: "dev", +}); + +export { TuskDrift }; diff --git a/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/src/test_requests.mjs b/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/src/test_requests.mjs new file mode 100644 index 00000000..b4dac25c --- /dev/null +++ b/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/src/test_requests.mjs @@ -0,0 +1,36 @@ +import { makeRequest, printRequestSummary } from "/app/test-utils.mjs"; + +await makeRequest("GET", "/health"); +await makeRequest("GET", "/test/insert-one"); +await makeRequest("GET", "/test/insert-many"); +await makeRequest("GET", "/test/find-one"); +await makeRequest("GET", "/test/find"); +await makeRequest("GET", "/test/find-with-options"); +await makeRequest("GET", "/test/find-cursor-next"); +await makeRequest("GET", "/test/update-one"); +await makeRequest("GET", "/test/update-many"); +await makeRequest("GET", "/test/delete-one"); +await makeRequest("GET", "/test/delete-many"); +await makeRequest("GET", "/test/replace-one"); +await makeRequest("GET", "/test/find-one-and-update"); +await makeRequest("GET", "/test/find-one-and-delete"); +await makeRequest("GET", "/test/find-one-and-replace"); +await makeRequest("GET", "/test/count-documents"); +await makeRequest("GET", "/test/estimated-count"); +await makeRequest("GET", "/test/distinct"); +await makeRequest("GET", "/test/aggregate"); +await makeRequest("GET", "/test/bulk-write"); +await makeRequest("GET", "/test/create-index"); +await makeRequest("GET", "/test/drop-index"); +await makeRequest("GET", "/test/list-indexes"); +await makeRequest("GET", "/test/db-command"); +await makeRequest("GET", "/test/list-collections"); +await makeRequest("GET", "/test/transaction"); +await makeRequest("GET", "/test/ordered-bulk"); +await makeRequest("GET", "/test/unordered-bulk"); +await makeRequest("GET", "/test/mongoose-create"); +await makeRequest("GET", "/test/mongoose-create-many"); +await makeRequest("GET", "/test/cursor-map"); +await makeRequest("GET", "/test/cursor-async-iterator"); + +printRequestSummary(); diff --git a/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/tsconfig.json b/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/tsconfig.json new file mode 100644 index 00000000..228eb708 --- /dev/null +++ b/src/instrumentation/libraries/mongodb/e2e-tests/esm-mongodb/tsconfig.json @@ -0,0 +1,17 @@ +{ + "compilerOptions": { + "target": "ES2020", + "module": "ES2020", + "lib": ["ES2020"], + "outDir": "./dist", + "rootDir": "./src", + "strict": true, + "esModuleInterop": true, + "skipLibCheck": true, + "forceConsistentCasingInFileNames": true, + "resolveJsonModule": true, + "moduleResolution": "node" + }, + "include": ["src/**/*"], + "exclude": ["node_modules", "dist"] +} diff --git a/src/instrumentation/libraries/mongodb/e2e-tests/run-all.sh b/src/instrumentation/libraries/mongodb/e2e-tests/run-all.sh new file mode 100755 index 00000000..4a7b2cdc --- /dev/null +++ b/src/instrumentation/libraries/mongodb/e2e-tests/run-all.sh @@ -0,0 +1,14 @@ +#!/bin/bash + +# Exit on error +set -e + +# Get the directory where this script is located +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +# Source common E2E helpers +source "$SCRIPT_DIR/../../e2e-common/e2e-helpers.sh" + +# Run all E2E tests for mongodb +# Accepts optional base port parameter (default: 3000) +run_all_e2e_tests "$SCRIPT_DIR" "mongodb" "${1:-3000}" diff --git a/src/instrumentation/libraries/mongodb/handlers/ConnectionHandler.ts b/src/instrumentation/libraries/mongodb/handlers/ConnectionHandler.ts new file mode 100644 index 00000000..94c8caee --- /dev/null +++ b/src/instrumentation/libraries/mongodb/handlers/ConnectionHandler.ts @@ -0,0 +1,277 @@ +import { SpanKind, SpanStatusCode } from "@opentelemetry/api"; +import { SpanUtils, SpanInfo } from "../../../../core/tracing/SpanUtils"; +import { TuskDriftMode } from "../../../../core/TuskDrift"; +import { handleRecordMode, handleReplayMode } from "../../../core/utils/modeUtils"; +import { PackageType } from "@use-tusk/drift-schemas/core/span"; +import { logger } from "../../../../core/utils"; + +export class ConnectionHandler { + constructor( + private mode: TuskDriftMode, + private instrumentationName: string, + private isAppReady: () => boolean, + ) {} + + /** + * Handle MongoClient.connect() calls. + * - RECORD: Call original connect, create span recording connection. + * - REPLAY: Skip TCP connection, resolve immediately with the client. + * - DISABLED: Passthrough. + */ + handleConnect(original: Function, thisArg: any, args: any[]): any { + const inputValue = this.extractConnectInputValue(thisArg); + + if (this.mode === TuskDriftMode.REPLAY) { + return handleReplayMode({ + noOpRequestHandler: () => { + return Promise.resolve(thisArg); + }, + isServerRequest: false, + replayModeHandler: () => { + return SpanUtils.createAndExecuteSpan( + this.mode, + () => Promise.resolve(thisArg), + { + name: "mongodb.connect", + kind: SpanKind.CLIENT, + submodule: "connect", + packageType: PackageType.MONGODB, + packageName: "mongodb", + instrumentationName: this.instrumentationName, + inputValue: inputValue, + isPreAppStart: !this.isAppReady(), + }, + (spanInfo: SpanInfo) => { + return this.handleReplayConnect(spanInfo, thisArg); + }, + ); + }, + }); + } else if (this.mode === TuskDriftMode.RECORD) { + return handleRecordMode({ + originalFunctionCall: () => original.apply(thisArg, args), + recordModeHandler: ({ isPreAppStart }) => { + return SpanUtils.createAndExecuteSpan( + this.mode, + () => original.apply(thisArg, args), + { + name: "mongodb.connect", + kind: SpanKind.CLIENT, + submodule: "connect", + packageType: PackageType.MONGODB, + packageName: "mongodb", + instrumentationName: this.instrumentationName, + inputValue: inputValue, + isPreAppStart, + }, + (spanInfo: SpanInfo) => { + return this.handleRecordConnect(spanInfo, original, thisArg, args); + }, + ); + }, + spanKind: SpanKind.CLIENT, + }); + } else { + return original.apply(thisArg, args); + } + } + + /** + * Handle MongoClient.close() calls. + * - RECORD: Call original close, create span. + * - REPLAY: No-op (no real connection to close). + * - DISABLED: Passthrough. + */ + handleClose(original: Function, thisArg: any, args: any[]): any { + if (this.mode === TuskDriftMode.REPLAY) { + return handleReplayMode({ + noOpRequestHandler: () => { + return Promise.resolve(); + }, + isServerRequest: false, + replayModeHandler: () => { + logger.debug(`[${this.instrumentationName}] Replaying MongoDB close (no-op)`); + return Promise.resolve(); + }, + }); + } else if (this.mode === TuskDriftMode.RECORD) { + return handleRecordMode({ + originalFunctionCall: () => original.apply(thisArg, args), + recordModeHandler: ({ isPreAppStart }) => { + return SpanUtils.createAndExecuteSpan( + this.mode, + () => original.apply(thisArg, args), + { + name: "mongodb.close", + kind: SpanKind.CLIENT, + submodule: "client-close", + packageType: PackageType.MONGODB, + packageName: "mongodb", + instrumentationName: this.instrumentationName, + inputValue: {}, + isPreAppStart, + stopRecordingChildSpans: true, + }, + (spanInfo: SpanInfo) => { + return this.handleRecordClose(spanInfo, original, thisArg, args); + }, + ); + }, + spanKind: SpanKind.CLIENT, + }); + } else { + return original.apply(thisArg, args); + } + } + + /** + * Handle MongoClient.db() calls. + * db() is synchronous and purely in-memory — creates a Db object referencing + * the client. No network I/O, no span needed. Passthrough in all modes. + */ + handleDb(original: Function, thisArg: any, args: any[]): any { + return original.apply(thisArg, args); + } + + // --------------------------------------------------------------------------- + // Private helpers + // --------------------------------------------------------------------------- + + private handleRecordConnect( + spanInfo: SpanInfo, + original: Function, + thisArg: any, + args: any[], + ): Promise { + const connectPromise = original.apply(thisArg, args) as Promise; + + return connectPromise + .then((result: any) => { + try { + logger.debug( + `[${this.instrumentationName}] MongoDB connection created successfully (${SpanUtils.getTraceInfo()})`, + ); + SpanUtils.addSpanAttributes(spanInfo.span, { + outputValue: { connected: true }, + }); + SpanUtils.endSpan(spanInfo.span, { code: SpanStatusCode.OK }); + } catch (error) { + logger.error( + `[${this.instrumentationName}] Error adding span attributes for connect:`, + error, + ); + } + return result; + }) + .catch((error: any) => { + try { + logger.error( + `[${this.instrumentationName}] MongoDB connection failed (${SpanUtils.getTraceInfo()}):`, + error, + ); + SpanUtils.addSpanAttributes(spanInfo.span, { + outputValue: { error: error?.message || "Unknown error" }, + }); + SpanUtils.endSpan(spanInfo.span, { + code: SpanStatusCode.ERROR, + message: error?.message || "Connection failed", + }); + } catch (spanError) { + logger.error( + `[${this.instrumentationName}] Error recording span for connect error:`, + spanError, + ); + } + throw error; + }); + } + + private handleReplayConnect(spanInfo: SpanInfo, thisArg: any): Promise { + logger.debug( + `[${this.instrumentationName}] Replaying MongoDB connection (skipping TCP connect)`, + ); + + try { + SpanUtils.addSpanAttributes(spanInfo.span, { + outputValue: { connected: true, replayed: true }, + }); + SpanUtils.endSpan(spanInfo.span, { code: SpanStatusCode.OK }); + } catch (error) { + logger.error(`[${this.instrumentationName}] Error ending replay connect span:`, error); + } + + return Promise.resolve(thisArg); + } + + private handleRecordClose( + spanInfo: SpanInfo, + original: Function, + thisArg: any, + args: any[], + ): Promise { + const closePromise = original.apply(thisArg, args) as Promise; + + return closePromise + .then(() => { + try { + logger.debug( + `[${this.instrumentationName}] MongoDB connection closed successfully (${SpanUtils.getTraceInfo()})`, + ); + SpanUtils.addSpanAttributes(spanInfo.span, { + outputValue: { closed: true }, + }); + SpanUtils.endSpan(spanInfo.span, { code: SpanStatusCode.OK }); + } catch (error) { + logger.error( + `[${this.instrumentationName}] Error adding span attributes for close:`, + error, + ); + } + }) + .catch((error: any) => { + try { + SpanUtils.addSpanAttributes(spanInfo.span, { + outputValue: { error: error?.message || "Unknown error" }, + }); + SpanUtils.endSpan(spanInfo.span, { + code: SpanStatusCode.ERROR, + message: error?.message || "Close failed", + }); + } catch (spanError) { + logger.error( + `[${this.instrumentationName}] Error recording span for close error:`, + spanError, + ); + } + throw error; + }); + } + + /** + * Sanitize a MongoDB connection string by removing password credentials. + */ + private sanitizeConnectionString(connectionString: string): string { + try { + const url = new URL(connectionString); + if (url.password) { + url.password = "***"; + } + return url.toString(); + } catch { + return connectionString.replace(/(:\/\/[^:]+):([^@]+)@/, "$1:***@"); + } + } + + /** + * Extract the connection input value from a MongoClient instance. + * The connection string is stored in this.s.url (set by the MongoClient constructor). + */ + private extractConnectInputValue(mongoClient: any): Record { + const rawUrl = mongoClient?.s?.url; + const sanitizedUrl = rawUrl ? this.sanitizeConnectionString(rawUrl) : undefined; + + return { + connectionString: sanitizedUrl, + }; + } +} diff --git a/src/instrumentation/libraries/mongodb/index.ts b/src/instrumentation/libraries/mongodb/index.ts new file mode 100644 index 00000000..18776dc6 --- /dev/null +++ b/src/instrumentation/libraries/mongodb/index.ts @@ -0,0 +1 @@ +export * from "./Instrumentation"; diff --git a/src/instrumentation/libraries/mongodb/mocks/FakeCursor.ts b/src/instrumentation/libraries/mongodb/mocks/FakeCursor.ts new file mode 100644 index 00000000..e61ffb4b --- /dev/null +++ b/src/instrumentation/libraries/mongodb/mocks/FakeCursor.ts @@ -0,0 +1,295 @@ +/** + * Fake MongoDB cursors for replay mode. + * + * These implement the key cursor interface methods so that application code + * interacting with a cursor works correctly during replay without hitting + * a real database. Supports builder-pattern chaining (sort, limit, skip, etc.) + * and all terminal iteration methods (toArray, next, forEach, async iterator). + * + * Lazy mock loading: The constructor accepts an optional `mockDataLoader` + * function. When provided, mock data is not loaded until the first terminal + * method is called. This preserves the synchronous return signature of + * find()/aggregate() while deferring the async mock lookup. + */ +export class TdFakeFindCursor { + protected documents: any[]; + private index: number = 0; + private _mockDataLoader: (() => Promise) | null; + private _mockLoadPromise: Promise | null = null; + private _mockLoaded: boolean; + + constructor(documents: any[] = [], mockDataLoader?: () => Promise) { + this.documents = documents; + this._mockDataLoader = mockDataLoader || null; + this._mockLoaded = !mockDataLoader; + } + + /** + * Lazily load mock data on first terminal method call. + * Subsequent calls return the same cached promise. + */ + private async _ensureMockLoaded(): Promise { + if (this._mockLoaded) return; + if (!this._mockLoadPromise && this._mockDataLoader) { + this._mockLoadPromise = this._mockDataLoader().then((docs) => { + this.documents = docs; + this._mockLoaded = true; + return docs; + }); + } + if (this._mockLoadPromise) { + await this._mockLoadPromise; + } + } + + // --- Terminal methods (must await mock loading) --- + + async toArray(): Promise { + await this._ensureMockLoaded(); + return [...this.documents]; + } + + async next(): Promise { + await this._ensureMockLoaded(); + return this.documents[this.index++] ?? null; + } + + async tryNext(): Promise { + await this._ensureMockLoaded(); + return this.documents[this.index++] ?? null; + } + + async hasNext(): Promise { + await this._ensureMockLoaded(); + return this.index < this.documents.length; + } + + async forEach(fn: (doc: any) => void): Promise { + await this._ensureMockLoaded(); + this.documents.forEach(fn); + } + + async *[Symbol.asyncIterator](): AsyncGenerator { + await this._ensureMockLoaded(); + for (const doc of this.documents) { + yield doc; + } + } + + // --- Builder methods (return this for chaining) --- + + filter(): this { + return this; + } + sort(): this { + return this; + } + limit(): this { + return this; + } + skip(): this { + return this; + } + project(): this { + return this; + } + hint(): this { + return this; + } + batchSize(): this { + return this; + } + maxTimeMS(): this { + return this; + } + collation(): this { + return this; + } + comment(): this { + return this; + } + min(): this { + return this; + } + max(): this { + return this; + } + returnKey(): this { + return this; + } + showRecordId(): this { + return this; + } + addQueryModifier(): this { + return this; + } + maxAwaitTimeMS(): this { + return this; + } + allowDiskUse(): this { + return this; + } + addCursorFlag(): this { + return this; + } + withReadPreference(): this { + return this; + } + withReadConcern(): this { + return this; + } + + map(): this { + // No-op: recorded mock data already contains post-transform results. + // The real cursor's terminal methods (toArray, next, etc.) apply the + // map transform before the instrumentation captures the output, so + // re-applying the transform here would corrupt the data. + return this; + } + + // --- Lifecycle methods --- + + async close(): Promise {} + + rewind(): void { + this.index = 0; + } + + clone(): TdFakeFindCursor { + if (this._mockLoaded) { + return new TdFakeFindCursor([...this.documents]); + } + return new TdFakeFindCursor([], this._mockDataLoader || undefined); + } + + // Stream support — returns a minimal async iterable + stream(): AsyncIterable { + return this[Symbol.asyncIterator]() as any; + } +} + +/** + * Fake MongoDB AggregationCursor for replay mode. + * + * Extends TdFakeFindCursor with pipeline builder methods that are no-ops + * during replay (the recorded result set is already computed). + */ +export class TdFakeAggregationCursor extends TdFakeFindCursor { + constructor(documents: any[] = [], mockDataLoader?: () => Promise) { + super(documents, mockDataLoader); + } + + // Pipeline builder methods (return this for chaining) + addStage(): this { + return this; + } + match(): this { + return this; + } + group(): this { + return this; + } + lookup(): this { + return this; + } + unwind(): this { + return this; + } + addFields(): this { + return this; + } + out(): this { + return this; + } + merge(): this { + return this; + } + redact(): this { + return this; + } + geoNear(): this { + return this; + } +} + +/** + * Fake MongoDB ChangeStream for replay mode. + * + * ChangeStreams are long-lived event-based streams. In replay mode, no real + * server connection exists so we return this minimal stub that: + * - Is EventEmitter-compatible (on/once/off/removeListener) + * - close() is a no-op + * - The async iterator yields nothing + * - hasNext() returns false, next() returns null + * + * This prevents crashes when application code creates a ChangeStream + * during replay, without attempting to replay individual change events. + */ +export class TdFakeChangeStream { + private _closed: boolean = false; + private _listeners: Map = new Map(); + + on(event: string, listener: Function): this { + if (!this._listeners.has(event)) { + this._listeners.set(event, []); + } + this._listeners.get(event)!.push(listener); + return this; + } + + once(event: string, listener: Function): this { + const wrappedListener = (...args: any[]) => { + this.off(event, wrappedListener); + listener(...args); + }; + return this.on(event, wrappedListener); + } + + off(event: string, listener: Function): this { + const eventListeners = this._listeners.get(event); + if (eventListeners) { + const index = eventListeners.indexOf(listener); + if (index !== -1) { + eventListeners.splice(index, 1); + } + } + return this; + } + + removeListener(event: string, listener: Function): this { + return this.off(event, listener); + } + + removeAllListeners(event?: string): this { + if (event) { + this._listeners.delete(event); + } else { + this._listeners.clear(); + } + return this; + } + + async close(): Promise { + this._closed = true; + } + + get closed(): boolean { + return this._closed; + } + + async hasNext(): Promise { + return false; + } + + async next(): Promise { + return null; + } + + async *[Symbol.asyncIterator](): AsyncGenerator { + // No events to yield in replay mode + } + + stream(): AsyncIterable { + return this[Symbol.asyncIterator]() as any; + } +} diff --git a/src/instrumentation/libraries/mongodb/mocks/FakeTopology.ts b/src/instrumentation/libraries/mongodb/mocks/FakeTopology.ts new file mode 100644 index 00000000..9eec61af --- /dev/null +++ b/src/instrumentation/libraries/mongodb/mocks/FakeTopology.ts @@ -0,0 +1,31 @@ +import { EventEmitter } from "events"; + +/** + * Fake MongoDB Topology for replay mode. + * + * When BulkOperationBase is constructed, it calls getTopology(collection) + * which requires a connected topology. In replay mode, no real connection + * exists, so we inject this fake topology to satisfy the constructor's + * requirements without hitting a real server. + * + * The constructor reads: + * - topology.lastHello() -> returns {} so all size limits use defaults + * - topology.lastIsMaster() -> returns {} (legacy compat) + * - topology.s.options -> returns {} so autoEncryption check is false + */ +export class TdFakeTopology extends EventEmitter { + s: { options: Record }; + + constructor() { + super(); + this.s = { options: {} }; + } + + lastHello(): Record { + return {}; + } + + lastIsMaster(): Record { + return {}; + } +} diff --git a/src/instrumentation/libraries/mongodb/types.ts b/src/instrumentation/libraries/mongodb/types.ts new file mode 100644 index 00000000..304ac972 --- /dev/null +++ b/src/instrumentation/libraries/mongodb/types.ts @@ -0,0 +1,76 @@ +import { TdInstrumentationConfig } from "../../core/baseClasses/TdInstrumentationAbstract"; +import { TuskDriftMode } from "../../../core/TuskDrift"; + +/** + * Input value for MongoDB collection/db operations (findOne, insertOne, aggregate, etc.) + */ +export interface MongodbCommandInputValue { + /** The MongoDB operation name (e.g., "findOne", "insertOne", "aggregate") */ + command: string; + /** The collection name */ + collection?: string; + /** The database name */ + database?: string; + /** Command-specific arguments (filter, document, pipeline, options, etc.) */ + commandArgs?: Record; + [key: string]: unknown; +} + +/** + * Input value for MongoClient.connect operations + */ +export interface MongodbConnectInputValue { + connectionString?: string; + options?: Record; + [key: string]: unknown; +} + +/** + * Module exports shape for the mongodb package + */ +export interface MongodbModuleExports { + MongoClient?: any; + Collection?: any; + Db?: any; + default?: any; + [key: string]: any; +} + +/** + * Configuration for MongoDB instrumentation + */ +export interface MongodbInstrumentationConfig extends TdInstrumentationConfig { + mode?: TuskDriftMode; +} + +/** + * MongoDB document type (generic object) + */ +export type MongodbDocument = Record; + +/** + * MongoDB operation result — varies by operation type + */ +export interface MongodbOutputValue { + /** For find/aggregate: array of returned documents */ + documents?: MongodbDocument[]; + /** For insertOne: the inserted ID */ + insertedId?: any; + /** For insertMany: map of index to inserted ID */ + insertedIds?: Record; + /** For insertMany: count of inserted documents */ + insertedCount?: number; + /** For update: count of matched documents */ + matchedCount?: number; + /** For update: count of modified documents */ + modifiedCount?: number; + /** For delete: count of deleted documents */ + deletedCount?: number; + /** For update with upsert: count of upserted documents */ + upsertedCount?: number; + /** For update with upsert: the upserted ID */ + upsertedId?: any; + /** Whether the operation was acknowledged by the server */ + acknowledged?: boolean; + [key: string]: unknown; +} diff --git a/src/instrumentation/libraries/mongodb/utils/bsonConversion.ts b/src/instrumentation/libraries/mongodb/utils/bsonConversion.ts new file mode 100644 index 00000000..62785616 --- /dev/null +++ b/src/instrumentation/libraries/mongodb/utils/bsonConversion.ts @@ -0,0 +1,528 @@ +import { SpanUtils, SpanInfo } from "../../../../core/tracing/SpanUtils"; +import { logger } from "../../../../core/utils"; + +// --------------------------------------------------------------------------- +// BSON type detection helpers +// --------------------------------------------------------------------------- + +function isObjectId(value: any): boolean { + return ( + typeof value === "object" && + value !== null && + (value._bsontype === "ObjectId" || value._bsontype === "ObjectID") && + typeof value.toHexString === "function" + ); +} + +function isUUID(value: any): boolean { + // UUID extends Binary so _bsontype is "Binary"; disambiguate via constructor name + return ( + typeof value === "object" && + value !== null && + value.constructor?.name === "UUID" && + value._bsontype === "Binary" && + typeof value.toString === "function" + ); +} + +function isBinary(value: any): boolean { + return ( + typeof value === "object" && + value !== null && + value._bsontype === "Binary" && + value.constructor?.name !== "UUID" + ); +} + +function isTimestamp(value: any): boolean { + return ( + typeof value === "object" && + value !== null && + value._bsontype === "Timestamp" + ); +} + +function isDecimal128(value: any): boolean { + return ( + typeof value === "object" && + value !== null && + value._bsontype === "Decimal128" && + typeof value.toString === "function" + ); +} + +function isLong(value: any): boolean { + return ( + typeof value === "object" && + value !== null && + value._bsontype === "Long" + ); +} + +function isDouble(value: any): boolean { + return ( + typeof value === "object" && + value !== null && + value._bsontype === "Double" + ); +} + +function isInt32(value: any): boolean { + return ( + typeof value === "object" && + value !== null && + value._bsontype === "Int32" + ); +} + +function isCode(value: any): boolean { + return ( + typeof value === "object" && + value !== null && + value._bsontype === "Code" + ); +} + +function isDBRef(value: any): boolean { + return ( + typeof value === "object" && + value !== null && + value._bsontype === "DBRef" + ); +} + +function isMaxKey(value: any): boolean { + return ( + typeof value === "object" && + value !== null && + value._bsontype === "MaxKey" + ); +} + +function isMinKey(value: any): boolean { + return ( + typeof value === "object" && + value !== null && + value._bsontype === "MinKey" + ); +} + +function isBSONRegExp(value: any): boolean { + return ( + typeof value === "object" && + value !== null && + value._bsontype === "BSONRegExp" + ); +} + +function isBSONSymbol(value: any): boolean { + return ( + typeof value === "object" && + value !== null && + value._bsontype === "BSONSymbol" + ); +} + +// --------------------------------------------------------------------------- +// BSON constructor resolution for reconstruction +// --------------------------------------------------------------------------- + +let cachedBsonModule: any = null; + +function getBsonConstructors(moduleExports?: any): any { + // Prefer module exports (mongodb re-exports all BSON types) + if (moduleExports?.ObjectId) { + return moduleExports; + } + + // Use cached bson module if available + if (cachedBsonModule) { + return cachedBsonModule; + } + + // Fallback: try to require bson (always available as mongodb dependency) + try { + cachedBsonModule = require("bson"); + return cachedBsonModule; + } catch { + logger.warn( + `[MongodbInstrumentation] Could not load BSON constructors for reconstruction`, + ); + return null; + } +} + +// --------------------------------------------------------------------------- +// Sanitization (BSON instances -> JSON-safe marker objects) +// --------------------------------------------------------------------------- + +function _sanitize(value: any, seen: WeakSet): any { + // Primitives and null/undefined pass through + if (value === null || value === undefined) { + return value; + } + if (typeof value !== "object" && typeof value !== "function") { + return value; + } + // Date is handled natively by JSON.stringify + if (value instanceof Date) { + return value; + } + + // --- BSON type checks (UUID before Binary is critical) --- + try { + if (isObjectId(value)) { + return { __bsonType: "ObjectId", value: value.toHexString() }; + } + if (isUUID(value)) { + return { __bsonType: "UUID", value: value.toString() }; + } + if (isBinary(value)) { + return { + __bsonType: "Binary", + base64: value.toString("base64"), + subType: value.sub_type, + }; + } + if (isTimestamp(value)) { + return { + __bsonType: "Timestamp", + t: typeof value.t === "number" ? value.t : Number(value.getHighBits?.() ?? 0), + i: typeof value.i === "number" ? value.i : Number(value.getLowBits?.() ?? 0), + }; + } + if (isDecimal128(value)) { + return { __bsonType: "Decimal128", value: value.toString() }; + } + if (isLong(value)) { + return { + __bsonType: "Long", + low: value.low, + high: value.high, + unsigned: !!value.unsigned, + }; + } + if (isDouble(value)) { + return typeof value.valueOf === "function" ? value.valueOf() : value.value; + } + if (isInt32(value)) { + return typeof value.valueOf === "function" ? value.valueOf() : value.value; + } + if (isCode(value)) { + return { + __bsonType: "Code", + code: value.code, + scope: value.scope ? _sanitize(value.scope, seen) : null, + }; + } + if (isDBRef(value)) { + return { + __bsonType: "DBRef", + collection: value.collection, + oid: _sanitize(value.oid, seen), + db: value.db || undefined, + }; + } + if (isMaxKey(value)) { + return { __bsonType: "MaxKey" }; + } + if (isMinKey(value)) { + return { __bsonType: "MinKey" }; + } + if (isBSONRegExp(value)) { + return { + __bsonType: "BSONRegExp", + pattern: value.pattern, + options: value.options, + }; + } + if (isBSONSymbol(value)) { + return { + __bsonType: "BSONSymbol", + value: typeof value.valueOf === "function" ? value.valueOf() : String(value), + }; + } + } catch (error) { + logger.warn( + `[MongodbInstrumentation] Error sanitizing BSON value, falling back to String():`, + error, + ); + return String(value); + } + + // --- Native RegExp — serializes to {} via JSON, so handle explicitly --- + if (value instanceof RegExp) { + return { __bsonType: "NativeRegExp", pattern: value.source, flags: value.flags }; + } + + // --- Node.js Buffer — serialize as compact base64 --- + if (Buffer.isBuffer(value)) { + return { __bsonType: "Buffer", base64: value.toString("base64") }; + } + + // --- Circular reference detection --- + if (seen.has(value)) { + return "[Circular]"; + } + seen.add(value); + + // --- Array handling --- + if (Array.isArray(value)) { + return value.map((item) => _sanitize(item, seen)); + } + + // --- Plain object handling --- + if (typeof value === "object") { + const result: any = {}; + for (const key of Object.keys(value)) { + result[key] = _sanitize(value[key], seen); + } + // Capture getter properties from the prototype (e.g. BulkWriteResult.ok) + const proto = Object.getPrototypeOf(value); + if (proto && proto !== Object.prototype) { + const descriptors = Object.getOwnPropertyDescriptors(proto); + for (const [key, desc] of Object.entries(descriptors)) { + if (desc.get && key !== "constructor" && !(key in result)) { + try { + const val = value[key]; + if (val !== undefined && typeof val !== "function") { + result[key] = _sanitize(val, seen); + } + } catch { + // Getter may throw; skip it + } + } + } + } + return result; + } + + return value; +} + +/** + * Sanitize BSON types to JSON-serializable marker representations. + * + * Recursively walks the input value and converts BSON type instances + * (ObjectId, Binary, UUID, Timestamp, etc.) to plain objects with a + * `__bsonType` discriminator that can survive JSON round-tripping. + */ +export function sanitizeBsonValue(value: any): any { + return _sanitize(value, new WeakSet()); +} + +// --------------------------------------------------------------------------- +// Reconstruction (JSON-safe markers -> BSON instances) +// --------------------------------------------------------------------------- + +/** + * Reconstruct BSON types from their JSON-safe marker representations. + * Used during replay mode to restore the correct BSON type instances + * so that application code calling .toString(), .equals(), or instanceof works correctly. + */ +export function reconstructBsonValue(value: any, moduleExports?: any): any { + if (value === null || value === undefined) { + return value; + } + if (typeof value !== "object") { + return value; + } + + // Array: recurse into each element + if (Array.isArray(value)) { + return value.map((item) => reconstructBsonValue(item, moduleExports)); + } + + // Detect __bsonType marker and reconstruct + if (value.__bsonType) { + const bson = getBsonConstructors(moduleExports); + if (!bson) { + logger.warn( + `[MongodbInstrumentation] Cannot reconstruct BSON type "${value.__bsonType}" — no constructors available`, + ); + return value; + } + + try { + switch (value.__bsonType) { + case "ObjectId": + if (bson.ObjectId) return new bson.ObjectId(value.value); + break; + case "UUID": + if (bson.UUID) return new bson.UUID(value.value); + break; + case "Binary": + if (bson.Binary) { + return new bson.Binary( + Buffer.from(value.base64, "base64"), + value.subType, + ); + } + break; + case "Timestamp": + if (bson.Timestamp) { + return new bson.Timestamp({ t: value.t, i: value.i }); + } + break; + case "Decimal128": + if (bson.Decimal128) { + return bson.Decimal128.fromString + ? bson.Decimal128.fromString(value.value) + : new bson.Decimal128(value.value); + } + break; + case "Long": + if (bson.Long) { + return new bson.Long(value.low, value.high, value.unsigned); + } + break; + case "Code": + if (bson.Code) { + const scope = value.scope + ? reconstructBsonValue(value.scope, moduleExports) + : undefined; + return new bson.Code(value.code, scope); + } + break; + case "DBRef": + if (bson.DBRef) { + const oid = reconstructBsonValue(value.oid, moduleExports); + return new bson.DBRef(value.collection, oid, value.db); + } + break; + case "MaxKey": + if (bson.MaxKey) return new bson.MaxKey(); + break; + case "MinKey": + if (bson.MinKey) return new bson.MinKey(); + break; + case "BSONRegExp": + if (bson.BSONRegExp) { + return new bson.BSONRegExp(value.pattern, value.options); + } + break; + case "BSONSymbol": + if (bson.BSONSymbol) return new bson.BSONSymbol(value.value); + break; + case "NativeRegExp": + return new RegExp(value.pattern, value.flags || ""); + case "Buffer": + return Buffer.from(value.base64, "base64"); + default: + logger.warn( + `[MongodbInstrumentation] Unknown BSON marker type: "${value.__bsonType}"`, + ); + return value; + } + } catch (error) { + logger.warn( + `[MongodbInstrumentation] Error reconstructing BSON type "${value.__bsonType}":`, + error, + ); + return value; + } + + // Constructor not found for this type + logger.warn( + `[MongodbInstrumentation] BSON constructor not available for type "${value.__bsonType}"`, + ); + return value; + } + + // Plain object: recurse into each property + const result: any = {}; + for (const key of Object.keys(value)) { + result[key] = reconstructBsonValue(value[key], moduleExports); + } + return result; +} + +// --------------------------------------------------------------------------- +// Span attribute helpers +// --------------------------------------------------------------------------- + +/** + * Add output attributes to a span for MongoDB operation results. + * Sanitizes BSON types before adding to ensure JSON-serializability. + */ +export function addOutputAttributesToSpan( + spanInfo: SpanInfo, + result: any, +): void { + try { + const sanitized = sanitizeBsonValue(result); + SpanUtils.addSpanAttributes(spanInfo.span, { outputValue: sanitized }); + } catch (error) { + logger.error( + `[MongodbInstrumentation] Error adding output attributes to span:`, + error, + ); + } +} + +/** + * Wrap cursor result documents in an object for recording. + * + * The Tusk CLI serializes span outputValue into a protobuf Struct for mock + * responses. Top-level arrays cannot be represented as a Struct field "body" + * (the CLI omits them), so we wrap the documents array inside a plain object. + * Use `unwrapCursorOutput` during replay to extract the documents. + */ +export function wrapCursorOutput(documents: any[]): any { + return { __cursorDocuments: documents }; +} + +/** + * Unwrap cursor result documents from the recording wrapper. + * Handles both the wrapped format (`{ __cursorDocuments: [...] }`) and + * a direct array fallback for forward compatibility. + */ +export function unwrapCursorOutput(result: any): any[] { + if (result && typeof result === "object" && Array.isArray(result.__cursorDocuments)) { + return result.__cursorDocuments; + } + if (Array.isArray(result)) { + return result; + } + return []; +} + +/** + * Wrap a direct method result so it survives protobuf Struct serialization. + * + * The Tusk CLI can only store objects in response.body (protobuf Struct). + * Non-object values (numbers, strings, booleans, arrays) are lost. + * This wraps them in a plain object; use `unwrapDirectOutput` during replay. + * Object results are left as-is since they already work with Struct. + */ +export function wrapDirectOutput(value: any): any { + if (value !== null && value !== undefined && typeof value === "object" && !Array.isArray(value)) { + return value; + } + return { __directResult: value }; +} + +/** + * Unwrap a direct method result from the recording wrapper. + * Handles both the wrapped format and passthrough for plain objects. + */ +export function unwrapDirectOutput(value: any): any { + if (value && typeof value === "object" && "__directResult" in value) { + return value.__directResult; + } + return value; +} + +/** + * Strips the `session` property from MongoDB options objects. + * Sessions are transient runtime objects that cannot be serialized + * and are not meaningful for mock matching. + */ +export function sanitizeOptions(options: any): any { + if (!options || typeof options !== "object") return options; + if (options.session) { + const { session, ...rest } = options; + return rest; + } + return options; +}