feat(tedious): optional context propagation to SQL Server using SET CONTEXT_INFO (#3141)

Author: mhennoch

packages/instrumentation-tedious/README.md

@@ -56,6 +56,13 @@ Attributes collected:
 | `net.peer.name`         | Remote hostname or similar.                                                    |
 | `net.peer.port`         | Remote port number.                                                            |
 
+### Trace Context Propagation
+
+Database trace context propagation can be enabled by setting `enableTraceContextPropagation`to `true`.
+This uses the [SET CONTEXT_INFO](https://learn.microsoft.com/en-us/sql/t-sql/statements/set-context-info-transact-sql?view=sql-server-ver16)
+command to set [traceparent](https://www.w3.org/TR/trace-context/#traceparent-header)information
+for the current connection, which results in **an additional round-trip to the database**.
+
 ## Useful links
 
 - For more information on OpenTelemetry, visit: <https://opentelemetry.io/>

packages/instrumentation-tedious/src/instrumentation.ts

@@ -40,6 +40,11 @@ import { PACKAGE_NAME, PACKAGE_VERSION } from './version';
 const CURRENT_DATABASE = Symbol(
   'opentelemetry.instrumentation-tedious.current-database'
 );
+
+export const INJECTED_CTX = Symbol(
+  'opentelemetry.instrumentation-tedious.context-info-injected'
+);
+
 const PATCHED_METHODS = [
   'callProcedure',
   'execSql',
@@ -89,7 +94,7 @@ export class TediousInstrumentation extends InstrumentationBase<TediousInstrumen
             this._wrap(
               ConnectionPrototype,
               method,
-              this._patchQuery(method) as any
+              this._patchQuery(method, moduleExports) as any
             );
           }
 
@@ -127,11 +132,61 @@ export class TediousInstrumentation extends InstrumentationBase<TediousInstrumen
     };
   }
 
-  private _patchQuery(operation: string) {
+  private _buildTraceparent(span: api.Span): string {
+    const sc = span.spanContext();
+    return `00-${sc.traceId}-${sc.spanId}-0${Number(sc.traceFlags || api.TraceFlags.NONE).toString(16)}`;
+  }
+
+  /**
+   * Fire a one-off `SET CONTEXT_INFO @opentelemetry_traceparent` on the same
+   * connection. Marks the request with INJECTED_CTX so our patch skips it.
+   */
+  private _injectContextInfo(
+    connection: any,
+    tediousModule: typeof tedious,
+    traceparent: string
+  ): Promise<void> {
+    return new Promise(resolve => {
+      try {
+        const sql = 'set context_info @opentelemetry_traceparent';
+        const req = new tediousModule.Request(sql, (_err: any) => {
+          resolve();
+        });
+        Object.defineProperty(req, INJECTED_CTX, { value: true });
+        const buf = Buffer.from(traceparent, 'utf8');
+        req.addParameter(
+          'opentelemetry_traceparent',
+          (tediousModule as any).TYPES.VarBinary,
+          buf,
+          { length: buf.length }
+        );
+
+        connection.execSql(req);
+      } catch {
+        resolve();
+      }
+    });
+  }
+
+  private _shouldInjectFor(operation: string): boolean {
+    return (
+      operation === 'execSql' ||
+      operation === 'execSqlBatch' ||
+      operation === 'callProcedure' ||
+      operation === 'execute'
+    );
+  }
+
+  private _patchQuery(operation: string, tediousModule: typeof tedious) {
     return (originalMethod: UnknownFunction): UnknownFunction => {
       const thisPlugin = this;
 
       function patchedMethod(this: ApproxConnection, request: ApproxRequest) {
+        // Skip our own injected request
+        if ((request as any)?.[INJECTED_CTX]) {
+          return originalMethod.apply(this, arguments as unknown as any[]);
+        }
+
         if (!(request instanceof EventEmitter)) {
           thisPlugin._diag.warn(
             `Unexpected invocation of patched ${operation} method. Span not recorded`
@@ -207,12 +262,27 @@ export class TediousInstrumentation extends InstrumentationBase<TediousInstrumen
           thisPlugin._diag.error('Expected request.callback to be a function');
         }
 
-        return api.context.with(
-          api.trace.setSpan(api.context.active(), span),
-          originalMethod,
-          this,
-          ...arguments
-        );
+        const runUserRequest = () => {
+          return api.context.with(
+            api.trace.setSpan(api.context.active(), span),
+            originalMethod,
+            this,
+            ...arguments
+          );
+        };
+
+        const cfg = thisPlugin.getConfig();
+        const shouldInject =
+          cfg.enableTraceContextPropagation &&
+          thisPlugin._shouldInjectFor(operation);
+
+        if (!shouldInject) return runUserRequest();
+
+        const traceparent = thisPlugin._buildTraceparent(span);
+        
+        void thisPlugin
+          ._injectContextInfo(this, tediousModule, traceparent)
+          .finally(runUserRequest);
       }
 
       Object.defineProperty(patchedMethod, 'length', {

packages/instrumentation-tedious/src/types.ts

@@ -15,5 +15,11 @@
  */
 
 import { InstrumentationConfig } from '@opentelemetry/instrumentation';
-
-export type TediousInstrumentationConfig = InstrumentationConfig;
+export interface TediousInstrumentationConfig extends InstrumentationConfig {
+  /**
+   * If true, injects the current DB span's W3C traceparent into SQL Server
+   * session state via `SET CONTEXT_INFO @opentelemetry_traceparent` (varbinary).
+   * Off by default to avoid the extra round-trip per request.
+   */
+  enableTraceContextPropagation?: boolean;
+}

packages/instrumentation-tedious/test/api.ts

@@ -17,7 +17,7 @@
 import * as assert from 'assert';
 import { promisify } from 'util';
 import type { Connection, Request, TYPES } from 'tedious';
-
+import { INJECTED_CTX } from '../src/instrumentation';
 type Method = keyof Connection & ('execSql' | 'execSqlBatch' | 'prepare');
 export type tedious = {
   Connection: typeof Connection;
@@ -67,7 +67,8 @@ export const makeApi = (tedious: tedious) => {
   const query = (
     connection: Connection,
     params: string,
-    method: Method = 'execSql'
+    method: Method = 'execSql',
+    noTracking?: boolean
   ): Promise<any[]> => {
     return new Promise((resolve, reject) => {
       const result: any[] = [];
@@ -78,7 +79,9 @@ export const makeApi = (tedious: tedious) => {
           resolve(result);
         }
       });
-
+      if (noTracking) {
+        Object.defineProperty(request, INJECTED_CTX, { value: true });
+      }
       // request.on('returnValue', console.log.bind(console, /*request.sqlTextOrProcedure,*/ 'returnValue:'));
       // request.on('error', console.log.bind(console, /*request.sqlTextOrProcedure,*/ 'error:'));
       // request.on('row', console.log.bind(console, /*request.sqlTextOrProcedure,*/ 'row:'));
@@ -314,7 +317,9 @@ export const makeApi = (tedious: tedious) => {
       if exists(SELECT * FROM sysobjects WHERE name='test_transact' AND xtype='U') DROP TABLE ${transaction.tableName};
       if exists(SELECT * FROM sysobjects WHERE name='test_proced' AND xtype='U') DROP PROCEDURE ${storedProcedure.procedureName};
       if exists(SELECT * FROM sys.databases WHERE name = 'temp_otel_db') DROP DATABASE temp_otel_db;
-    `.trim()
+    `.trim(),
+      'execSql',
+      true
     );
   };
 

packages/instrumentation-tedious/test/instrumentation.test.ts

@@ -47,8 +47,6 @@ const user = process.env.MSSQL_USER || 'sa';
 const password = process.env.MSSQL_PASSWORD || 'mssql_passw0rd';
 
 const instrumentation = new TediousInstrumentation();
-instrumentation.enable();
-instrumentation.disable();
 
 const config: any = {
   userName: user,
@@ -308,6 +306,68 @@ describe('tedious', () => {
       table: 'test_bulk',
     });
   });
+
+  describe('trace context propagation via CONTEXT_INFO', () => {
+    function traceparentFromSpan(span: ReadableSpan) {
+      const sc = span.spanContext();
+      const flags = sc.traceFlags & 0x01 ? '01' : '00';
+      return `00-${sc.traceId}-${sc.spanId}-${flags}`;
+    }
+
+    beforeEach(() => {
+      instrumentation.setConfig({
+        enableTraceContextPropagation: true,
+      });
+    });
+
+    after(() => {
+      instrumentation.setConfig({ enableTraceContextPropagation: false });
+    });
+
+    it('injects DB-span traceparent for execSql', async function () {
+      const sql =
+        "SELECT REPLACE(CONVERT(varchar(128), CONTEXT_INFO()), CHAR(0), '') AS traceparent";
+      const result = await tedious.query(connection, sql);
+
+      const spans = memoryExporter.getFinishedSpans();
+      assert.strictEqual(spans.length, 1);
+      const expectedTp = traceparentFromSpan(spans[0]);
+      assert.strictEqual(
+        result[0],
+        expectedTp,
+        'CONTEXT_INFO traceparent should match DB span'
+      );
+    });
+
+    it('injects for execSqlBatch', async function () {
+      const batch = `
+          SELECT REPLACE(CONVERT(varchar(128), CONTEXT_INFO()), CHAR(0), '') AS tp;
+          SELECT 42;
+        `;
+      const result = await tedious.query(connection, batch, 'execSqlBatch');
+
+      assert.deepStrictEqual(result, [result[0], 42]);
+
+      const spans = memoryExporter.getFinishedSpans();
+      assert.strictEqual(spans.length, 1);
+      const expectedTp = traceparentFromSpan(spans[0]);
+      assert.strictEqual(result[0], expectedTp);
+    });
+
+    it('when disabled, CONTEXT_INFO stays empty', async function () {
+      instrumentation.setConfig({
+        enableTraceContextPropagation: false,
+      });
+
+      const [val] = await tedious.query(
+        connection,
+        "SELECT REPLACE(CONVERT(varchar(128), CONTEXT_INFO()), CHAR(0), '')"
+      );
+      assert.strictEqual(val, null);
+      const spans = memoryExporter.getFinishedSpans();
+      assert.strictEqual(spans.length, 1);
+    });
+  });
 });
 
 const assertMatch = (actual: string | undefined, expected: RegExp) => {