feat: add mysql2 responsehook (#915)

* supprt responsehook config * add specs * response hook receives an info object
open-telemetry · Feb 27, 2022 · f436601 · f436601
1 parent 2229cbe
commit f436601
Show file tree

Hide file tree

Showing 4 changed files with 168 additions and 6 deletions.
diff --git a/plugins/node/opentelemetry-instrumentation-mysql2/README.md b/plugins/node/opentelemetry-instrumentation-mysql2/README.md
@@ -43,6 +43,14 @@ registerInstrumentations({
 })
 ```
 
+### MySQL2 Instrumentation Options
+
+You can set the following instrumentation options:
+
+| Options | Type | Description |
+| ------- | ---- | ----------- |
+| `responseHook` | `MySQL2InstrumentationExecutionResponseHook` (function) | Function for adding custom attributes from db response |
+
 ## Useful links
 
 - For more information on OpenTelemetry, visit: <https://opentelemetry.io/>

diff --git a/plugins/node/opentelemetry-instrumentation-mysql2/src/instrumentation.ts b/plugins/node/opentelemetry-instrumentation-mysql2/src/instrumentation.ts
@@ -19,6 +19,7 @@ import {
   InstrumentationBase,
   InstrumentationNodeModuleDefinition,
   isWrapped,
+  safeExecuteInTheMiddle,
 } from '@opentelemetry/instrumentation';
 import {
   DbSystemValues,
@@ -109,13 +110,29 @@ export class MySQL2Instrumentation extends InstrumentationBase<
             ),
           },
         });
-        const endSpan = once((err?: any) => {
+        const endSpan = once((err?: any, results?: any) => {
           if (err) {
             span.setStatus({
               code: api.SpanStatusCode.ERROR,
               message: err.message,
             });
+          } else {
+            const config: MySQL2InstrumentationConfig = thisPlugin._config;
+            if (typeof config.responseHook === 'function') {
+              safeExecuteInTheMiddle(
+                () => {
+                  config.responseHook!(span, { queryResults: results });
+                },
+                err => {
+                  if (err) {
+                    thisPlugin._diag.warn('Failed executing responseHook', err);
+                  }
+                },
+                true
+              );
+            }
           }
+
           span.end();
         });
 
@@ -138,8 +155,8 @@ export class MySQL2Instrumentation extends InstrumentationBase<
             .once('error', err => {
               endSpan(err);
             })
-            .once('result', () => {
-              endSpan();
+            .once('result', results => {
+              endSpan(undefined, results);
             });
 
           return streamableQuery;
@@ -171,7 +188,7 @@ export class MySQL2Instrumentation extends InstrumentationBase<
         results?: any,
         fields?: mysqlTypes.FieldPacket[]
       ) {
-        endSpan(err);
+        endSpan(err, results);
         return originalCallback(...arguments);
       };
     };

diff --git a/plugins/node/opentelemetry-instrumentation-mysql2/src/types.ts b/plugins/node/opentelemetry-instrumentation-mysql2/src/types.ts
@@ -15,5 +15,22 @@
  */
 
 import { InstrumentationConfig } from '@opentelemetry/instrumentation';
+import type { Span } from '@opentelemetry/api';
 
-export type MySQL2InstrumentationConfig = InstrumentationConfig;
+export interface MySQL2ResponseHookInformation {
+  queryResults: any;
+}
+
+export interface MySQL2InstrumentationExecutionResponseHook {
+  (span: Span, responseHookInfo: MySQL2ResponseHookInformation): void;
+}
+
+export interface MySQL2InstrumentationConfig extends InstrumentationConfig {
+  /**
+   * Hook that allows adding custom span attributes based on the data
+   * returned MySQL2 queries.
+   *
+   * @default undefined
+   */
+  responseHook?: MySQL2InstrumentationExecutionResponseHook;
+}
diff --git a/plugins/node/opentelemetry-instrumentation-mysql2/test/mysql.test.ts b/plugins/node/opentelemetry-instrumentation-mysql2/test/mysql.test.ts
@@ -29,7 +29,7 @@ import {
   SimpleSpanProcessor,
 } from '@opentelemetry/sdk-trace-base';
 import * as assert from 'assert';
-import { MySQL2Instrumentation } from '../src';
+import { MySQL2Instrumentation, MySQL2InstrumentationConfig } from '../src';
 
 const LIB_VERSION = testUtils.getPackageVersion('mysql2');
 const port = Number(process.env.MYSQL_PORT) || 33306;
@@ -647,6 +647,126 @@ describe('mysql@2.x', () => {
       );
     });
   });
+
+  describe('#responseHook', () => {
+    const queryResultAttribute = 'query_result';
+
+    after(() => {
+      instrumentation.setConfig({});
+    });
+
+    describe('invalid repsonse hook', () => {
+      before(() => {
+        instrumentation.disable();
+        instrumentation.setTracerProvider(provider);
+        const config: MySQL2InstrumentationConfig = {
+          responseHook: (span, responseHookInfo) => {
+            throw new Error('random failure!');
+          },
+        };
+        instrumentation.setConfig(config);
+        instrumentation.enable();
+      });
+
+      it('should not affect the behavior of the query', done => {
+        const span = provider.getTracer('default').startSpan('test span');
+        context.with(trace.setSpan(context.active(), span), () => {
+          const sql = 'SELECT 1+1 as solution';
+          connection.query(sql, (err, res: mysqlTypes.RowDataPacket[]) => {
+            assert.ifError(err);
+            assert.ok(res);
+            assert.strictEqual(res[0].solution, 2);
+            done();
+          });
+        });
+      });
+    });
+
+    describe('valid response hook', () => {
+      before(() => {
+        instrumentation.disable();
+        instrumentation.setTracerProvider(provider);
+        const config: MySQL2InstrumentationConfig = {
+          responseHook: (span, responseHookInfo) => {
+            span.setAttribute(
+              queryResultAttribute,
+              JSON.stringify(responseHookInfo.queryResults)
+            );
+          },
+        };
+        instrumentation.setConfig(config);
+        instrumentation.enable();
+      });
+
+      it('should extract data from responseHook - connection', done => {
+        const span = provider.getTracer('default').startSpan('test span');
+        context.with(trace.setSpan(context.active(), span), () => {
+          const sql = 'SELECT 1+1 as solution';
+          connection.query(sql, (err, res: mysqlTypes.RowDataPacket[]) => {
+            assert.ifError(err);
+            assert.ok(res);
+            assert.strictEqual(res[0].solution, 2);
+            const spans = memoryExporter.getFinishedSpans();
+            assert.strictEqual(spans.length, 1);
+            assertSpan(spans[0], sql);
+            assert.strictEqual(
+              spans[0].attributes[queryResultAttribute],
+              JSON.stringify(res)
+            );
+            done();
+          });
+        });
+      });
+
+      it('should extract data from responseHook - pool', done => {
+        const span = provider.getTracer('default').startSpan('test span');
+        context.with(trace.setSpan(context.active(), span), () => {
+          const sql = 'SELECT 1+1 as solution';
+          pool.getConnection((err, conn) => {
+            conn.query(sql, (err, res: mysqlTypes.RowDataPacket[]) => {
+              assert.ifError(err);
+              assert.ok(res);
+              assert.strictEqual(res[0].solution, 2);
+              const spans = memoryExporter.getFinishedSpans();
+              assert.strictEqual(spans.length, 1);
+              assertSpan(spans[0], sql);
+              assert.strictEqual(
+                spans[0].attributes[queryResultAttribute],
+                JSON.stringify(res)
+              );
+              done();
+            });
+          });
+        });
+      });
+
+      it('should extract data from responseHook - poolCluster', done => {
+        poolCluster.getConnection((err, poolClusterConnection) => {
+          assert.ifError(err);
+          const span = provider.getTracer('default').startSpan('test span');
+          context.with(trace.setSpan(context.active(), span), () => {
+            const sql = 'SELECT 1+1 as solution';
+            poolClusterConnection.query(
+              sql,
+              (err, res: mysqlTypes.RowDataPacket[]) => {
+                assert.ifError(err);
+                assert.ok(res);
+                assert.strictEqual(res[0].solution, 2);
+                const spans = memoryExporter.getFinishedSpans();
+                assert.strictEqual(spans.length, 1);
+                assertSpan(spans[0], sql);
+                assert.strictEqual(
+                  spans[0].attributes[queryResultAttribute],
+                  JSON.stringify(res)
+                );
+                done();
+              }
+            );
+          });
+        });
+      });
+    });
+  });
 });
 
 function assertSpan(