feat(instrumentation-mysql2): Add hook for setting span name

open-telemetry · Nov 26, 2024 · ad48e45 · ad48e45
1 parent d5215f3
commit ad48e45
Show file tree

Hide file tree

Showing 3 changed files with 151 additions and 2 deletions.
diff --git a/plugins/node/opentelemetry-instrumentation-mysql2/src/instrumentation.ts b/plugins/node/opentelemetry-instrumentation-mysql2/src/instrumentation.ts
@@ -28,7 +28,7 @@ import {
 } from '@opentelemetry/semantic-conventions';
 import { addSqlCommenterComment } from '@opentelemetry/sql-common';
 import type * as mysqlTypes from 'mysql2';
-import { MySQL2InstrumentationConfig } from './types';
+import { MySQL2InstrumentationConfig, MySQL2RequestInfo } from './types';
 import {
   getConnectionAttributes,
   getDbStatement,
@@ -104,7 +104,30 @@ export class MySQL2Instrumentation extends InstrumentationBase<MySQL2Instrumenta
           values = [_valuesOrCallback];
         }
 
-        const span = thisPlugin.tracer.startSpan(getSpanName(query), {
+        const defaultSpanName = getSpanName(query);
+        const mysql2RequestInfo: MySQL2RequestInfo = {
+          query,
+          values,
+          database: this.config.database,
+          host: this.config.host,
+        };
+        const spanNameHook = thisPlugin.getConfig().spanNameHook;
+        let spanName = defaultSpanName;
+        if (spanNameHook) {
+          spanName =
+            safeExecuteInTheMiddle(
+              () => spanNameHook(mysql2RequestInfo, defaultSpanName),
+              (err, result) => {
+                if (err) {
+                  thisPlugin._diag.warn('Failed executing spanNameHook', err);
+                }
+                return result;
+              },
+              true
+            ) ?? defaultSpanName;
+        }
+
+        const span = thisPlugin.tracer.startSpan(spanName, {
           kind: api.SpanKind.CLIENT,
           attributes: {
             ...MySQL2Instrumentation.COMMON_ATTRIBUTES,

diff --git a/plugins/node/opentelemetry-instrumentation-mysql2/src/types.ts b/plugins/node/opentelemetry-instrumentation-mysql2/src/types.ts
@@ -16,6 +16,7 @@
 
 import { InstrumentationConfig } from '@opentelemetry/instrumentation';
 import type { Span } from '@opentelemetry/api';
+import type { Query, QueryOptions } from 'mysql2';
 
 export interface MySQL2ResponseHookInformation {
   queryResults: any;
@@ -25,6 +26,22 @@ export interface MySQL2InstrumentationExecutionResponseHook {
   (span: Span, responseHookInfo: MySQL2ResponseHookInformation): void;
 }
 
+export interface MySQL2RequestInfo {
+  host?: string;
+  database?: string;
+  query: string | Query | QueryOptions;
+  values?: unknown[];
+}
+
+export type SpanNameHook = (
+  info: MySQL2RequestInfo,
+  /**
+   * If no decision is taken based on RequestInfo, the default name
+   * supplied by the instrumentation can be used instead.
+   */
+  defaultName: string
+) => string;
+
 export interface MySQL2InstrumentationConfig extends InstrumentationConfig {
   /**
    * Hook that allows adding custom span attributes based on the data
@@ -34,6 +51,11 @@ export interface MySQL2InstrumentationConfig extends InstrumentationConfig {
    */
   responseHook?: MySQL2InstrumentationExecutionResponseHook;
 
+  /**
+   * Hook to override the name for an SQL span
+   */
+  spanNameHook?: SpanNameHook;
+
   /**
    * If true, queries are modified to also include a comment with
    * the tracing context, following the {@link https://github.com/open-telemetry/opentelemetry-sqlcommenter sqlcommenter} format

diff --git a/plugins/node/opentelemetry-instrumentation-mysql2/test/mysql.test.ts b/plugins/node/opentelemetry-instrumentation-mysql2/test/mysql.test.ts
@@ -1100,6 +1100,110 @@ describe('mysql2', () => {
     });
   });
 
+  describe('#spanNameHook', () => {
+    it('span hook gets request info as parameters', done => {
+      const config: MySQL2InstrumentationConfig = {
+        spanNameHook: (info, defaultName) => {
+          assert.strictEqual(defaultName, 'SELECT');
+          assert.deepStrictEqual(info, {
+            database: 'test_db',
+            host: '127.0.0.1',
+            query: 'SELECT ? as solution',
+            values: ['otel-user'],
+          });
+          return defaultName;
+        },
+      };
+      instrumentation.setConfig(config);
+
+      const span = provider.getTracer('default').startSpan('test span');
+      context.with(trace.setSpan(context.active(), span), () => {
+        const sql = 'SELECT ? as solution';
+        const query = connection.query(sql, ['otel-user']);
+
+        query.on('end', () => {
+          done();
+        });
+      });
+    });
+
+    describe('valid span name hook', () => {
+      beforeEach(() => {
+        const config: MySQL2InstrumentationConfig = {
+          spanNameHook: (info, defaultName) => {
+            const query =
+              typeof info.query === 'string' ? info.query : info.query.sql;
+            const prioritizedSqlVerbs = [
+              'DROP',
+              'DELETE',
+              'INSERT',
+              'UPDATE',
+              'SELECT',
+            ];
+            for (const verb of prioritizedSqlVerbs) {
+              if (query.includes(verb)) {
+                return verb;
+              }
+            }
+            return 'UNKNOWN';
+          },
+        };
+        instrumentation.setConfig(config);
+      });
+
+      it('should set span name using spanNameHook', done => {
+        const span = provider.getTracer('default').startSpan('test span');
+        context.with(trace.setSpan(context.active(), span), () => {
+          const sql =
+            'WITH number AS (SELECT 1+1 as solution) SELECT solution FROM number';
+          connection.query(
+            sql,
+            ['otel-user'],
+            (err, res: mysqlTypes.RowDataPacket[]) => {
+              assert.ifError(err);
+              assert.ok(res);
+              console.log(res[0]);
+              assert.strictEqual(res[0].solution, 2);
+              const spans = memoryExporter.getFinishedSpans();
+              assert.strictEqual(spans.length, 1);
+              assertSpan(spans[0], sql);
+              assert.strictEqual(spans[0].name, 'SELECT');
+              done();
+            }
+          );
+        });
+      });
+    });
+
+    describe('invalid span name hook', () => {
+      beforeEach(() => {
+        const config: MySQL2InstrumentationConfig = {
+          spanNameHook: () => {
+            throw new Error('could not decide on a name');
+          },
+        };
+        instrumentation.setConfig(config);
+      });
+
+      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);
+            const spans = memoryExporter.getFinishedSpans();
+            assert.strictEqual(spans.length, 1);
+            assertSpan(spans[0], sql);
+            assert.strictEqual(spans[0].name, 'SELECT');
+            done();
+          });
+        });
+      });
+    });
+  });
+
   describe('#responseHook', () => {
     const queryResultAttribute = 'query_result';