diff --git a/plugins/node/opentelemetry-instrumentation-mongodb/README.md b/plugins/node/opentelemetry-instrumentation-mongodb/README.md index 9fe1d7d291..f31c88feca 100644 --- a/plugins/node/opentelemetry-instrumentation-mongodb/README.md +++ b/plugins/node/opentelemetry-instrumentation-mongodb/README.md @@ -50,6 +50,7 @@ Mongodb instrumentation has few options available to choose from. You can set th | Options | Type | Description | | ------- | ---- | ----------- | | [`enhancedDatabaseReporting`](https://github.com/open-telemetry/opentelemetry-js/blob/main/packages/opentelemetry-api/src/trace/instrumentation/instrumentation.ts#L91) | `string` | If true, additional information about query parameters and results will be attached (as `attributes`) to spans representing database operations | +| `responseHook` | `MongoDBInstrumentationExecutionResponseHook` (function) | Function for adding custom attributes from db response | ## Useful links diff --git a/plugins/node/opentelemetry-instrumentation-mongodb/src/instrumentation.ts b/plugins/node/opentelemetry-instrumentation-mongodb/src/instrumentation.ts index 5fa9f9c3c4..a72909ca9b 100644 --- a/plugins/node/opentelemetry-instrumentation-mongodb/src/instrumentation.ts +++ b/plugins/node/opentelemetry-instrumentation-mongodb/src/instrumentation.ts @@ -27,6 +27,7 @@ import { InstrumentationNodeModuleDefinition, InstrumentationNodeModuleFile, isWrapped, + safeExecuteInTheMiddle, } from '@opentelemetry/instrumentation'; import { SemanticAttributes } from '@opentelemetry/semantic-conventions'; import type * as mongodb from 'mongodb'; @@ -37,6 +38,7 @@ import { MongoInternalCommand, MongoInternalTopology, WireProtocolInternal, + CommandResult, } from './types'; import { VERSION } from './version'; @@ -417,6 +419,32 @@ export class MongoDBInstrumentation extends InstrumentationBase< span.setAttribute(SemanticAttributes.DB_STATEMENT, JSON.stringify(query)); } + /** + * Triggers the response hook in case it is defined. + * @param span The span to add the results to. + * @param config The MongoDB instrumentation config object + * @param result The command result + */ + private _handleExecutionResult( + span: Span, + result: CommandResult + ) { + const config: MongoDBInstrumentationConfig = this.getConfig(); + if (typeof config.responseHook === 'function') { + safeExecuteInTheMiddle( + () => { + config.responseHook!(span, { data: result }); + }, + err => { + if (err) { + this._diag.error('Error running response hook', err); + } + }, + true + ); + } + } + /** * Ends a created span. * @param span The created span to end. @@ -426,6 +454,7 @@ export class MongoDBInstrumentation extends InstrumentationBase< // mongodb is using "tick" when calling a callback, this way the context // in final callback (resultHandler) is lost const activeContext = context.active(); + const instrumentation = this; return function patchedEnd(this: {}, ...args: unknown[]) { const error = args[0]; if (error instanceof Error) { @@ -433,6 +462,9 @@ export class MongoDBInstrumentation extends InstrumentationBase< code: SpanStatusCode.ERROR, message: error.message, }); + } else { + const result = args[1] as CommandResult; + instrumentation._handleExecutionResult(span, result); } span.end(); diff --git a/plugins/node/opentelemetry-instrumentation-mongodb/src/types.ts b/plugins/node/opentelemetry-instrumentation-mongodb/src/types.ts index 6948928d9a..e86a9e0a9d 100644 --- a/plugins/node/opentelemetry-instrumentation-mongodb/src/types.ts +++ b/plugins/node/opentelemetry-instrumentation-mongodb/src/types.ts @@ -15,6 +15,11 @@ */ import { InstrumentationConfig } from '@opentelemetry/instrumentation'; +import { Span } from '@opentelemetry/api'; + +export interface MongoDBInstrumentationExecutionResponseHook { + (span: Span, responseInfo: MongoResponseHookInformation): void; +} export interface MongoDBInstrumentationConfig extends InstrumentationConfig { /** @@ -23,6 +28,14 @@ export interface MongoDBInstrumentationConfig extends InstrumentationConfig { * database operations. */ enhancedDatabaseReporting?: boolean; + + /** + * Hook that allows adding custom span attributes based on the data + * returned from MongoDB actions. + * + * @default undefined + */ + responseHook?: MongoDBInstrumentationExecutionResponseHook; } export type Func = (...args: unknown[]) => T; @@ -43,6 +56,17 @@ export type CursorState = { cmd: MongoInternalCommand } & Record< unknown >; +export interface MongoResponseHookInformation { + data: CommandResult; +} + +// https://github.com/mongodb/node-mongodb-native/blob/3.6/lib/core/connection/command_result.js +export type CommandResult = { + result?: object; + connection?: object; + message?: object; +}; + // https://github.com/mongodb/node-mongodb-native/blob/3.6/lib/core/wireprotocol/index.js export type WireProtocolInternal = { insert: ( diff --git a/plugins/node/opentelemetry-instrumentation-mongodb/test/mongodb.test.ts b/plugins/node/opentelemetry-instrumentation-mongodb/test/mongodb.test.ts index 8a2e73f3eb..cac5a334a5 100644 --- a/plugins/node/opentelemetry-instrumentation-mongodb/test/mongodb.test.ts +++ b/plugins/node/opentelemetry-instrumentation-mongodb/test/mongodb.test.ts @@ -16,7 +16,7 @@ // for testing locally "npm run docker:start" -import { context, trace, SpanKind } from '@opentelemetry/api'; +import { context, trace, SpanKind, Span } from '@opentelemetry/api'; import { BasicTracerProvider } from '@opentelemetry/tracing'; import { AsyncHooksContextManager } from '@opentelemetry/context-async-hooks'; import { @@ -24,7 +24,8 @@ import { SimpleSpanProcessor, } from '@opentelemetry/tracing'; import * as assert from 'assert'; -import { MongoDBInstrumentation } from '../src'; +import { MongoDBInstrumentation, MongoDBInstrumentationConfig } from '../src'; +import { MongoResponseHookInformation } from '../src/types'; const instrumentation = new MongoDBInstrumentation(); instrumentation.enable(); @@ -34,6 +35,10 @@ import * as mongodb from 'mongodb'; import { assertSpans, accessCollection } from './utils'; describe('MongoDBInstrumentation', () => { + function create(config: MongoDBInstrumentationConfig = {}) { + instrumentation.setConfig(config); + instrumentation.enable(); + } // For these tests, mongo must be running. Add RUN_MONGODB_TESTS to run // these tests. const RUN_MONGODB_TESTS = process.env.RUN_MONGODB_TESTS as string; @@ -244,6 +249,88 @@ describe('MongoDBInstrumentation', () => { }); }); + describe('when specifying a responseHook configuration', () => { + const dataAttributeName = 'mongodb_data'; + beforeEach(() => { + memoryExporter.reset(); + }); + + describe('with a valid function', () => { + beforeEach(() => { + create({ + responseHook: (span: Span, result: MongoResponseHookInformation) => { + span.setAttribute(dataAttributeName, JSON.stringify(result.data.result)); + }, + }); + }); + + it('should attach response hook data to the resulting span for insert function', done => { + const insertData = [{ a: 1 }, { a: 2 }, { a: 3 }]; + const span = provider.getTracer('default').startSpan('insertRootSpan'); + context.with(trace.setSpan(context.active(), span), () => { + collection.insertMany(insertData, (err, result) => { + span.end(); + assert.ifError(err); + const spans = memoryExporter.getFinishedSpans(); + const insertSpan = spans[0]; + + assert.deepStrictEqual( + JSON.parse(insertSpan.attributes[dataAttributeName] as string), + result.result + ); + + done(); + }); + }); + }); + + it('should attach response hook data to the resulting span for find function', done => { + const span = provider.getTracer('default').startSpan('findRootSpan'); + context.with(trace.setSpan(context.active(), span), () => { + collection.find({ a: 1 }).toArray((err, results) => { + span.end(); + assert.ifError(err); + const spans = memoryExporter.getFinishedSpans(); + const findSpan = spans[0]; + const hookAttributeValue = JSON.parse(findSpan.attributes[dataAttributeName] as string); + + assert.strictEqual( + hookAttributeValue?.cursor?.firstBatch[0]._id, + results[0]._id.toString() + ); + + done(); + }); + }); + }); + }); + + describe('with an invalid function', () => { + beforeEach(() => { + create({ + responseHook: (span: Span, result: MongoResponseHookInformation) => { + throw 'some error'; + }, + }); + }); + + it('should not do any harm when throwing an exception', done => { + const span = provider.getTracer('default').startSpan('findRootSpan'); + context.with(trace.setSpan(context.active(), span), () => { + collection.find({ a: 1 }).toArray((err, results) => { + span.end(); + const spans = memoryExporter.getFinishedSpans(); + + assert.ifError(err); + assertSpans(spans, 'mongodb.find', SpanKind.CLIENT); + + done(); + }); + }); + }); + }); + }); + describe('Mixed operations with callback', () => { beforeEach(() => { memoryExporter.reset();