feat(mongodb instrumentation): added response hook option

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
 

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,29 @@ 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,13 +451,17 @@ 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) {
         span.setStatus({
           code: SpanStatusCode.ERROR,
           message: error.message,
         });
+      } else {
+        const result = args[1] as CommandResult;
+        instrumentation._handleExecutionResult(span, result);
       }
       span.end();
 

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<T> = (...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?: unknown;
+  connection?: unknown;
+  message?: unknown;
+};
+
 // https://github.com/mongodb/node-mongodb-native/blob/3.6/lib/core/wireprotocol/index.js
 export type WireProtocolInternal = {
   insert: (

plugins/node/opentelemetry-instrumentation-mongodb/test/mongodb.test.ts

@@ -16,15 +16,16 @@
 
 // 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 {
   InMemorySpanExporter,
   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,93 @@ 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();