feat(tracing): Add transaction.setContext method (#6154)

Author: AbhiPrasad

packages/opentelemetry-node/test/spanprocessor.test.ts

@@ -7,7 +7,7 @@ import { SemanticAttributes, SemanticResourceAttributes } from '@opentelemetry/s
 import { createTransport, Hub, makeMain } from '@sentry/core';
 import { NodeClient } from '@sentry/node';
 import { addExtensionMethods, Span as SentrySpan, SpanStatusType, Transaction } from '@sentry/tracing';
-import { Contexts, Scope } from '@sentry/types';
+import { Scope } from '@sentry/types';
 import { resolvedSyncPromise } from '@sentry/utils';
 
 import { SENTRY_SPAN_PROCESSOR_MAP, SentrySpanProcessor } from '../src/spanprocessor';
@@ -55,21 +55,22 @@ describe('SentrySpanProcessor', () => {
   }
 
   function getContext(transaction: Transaction) {
-    const transactionWithContext = transaction as unknown as Transaction & { _contexts: Contexts };
+    const transactionWithContext = transaction as unknown as Transaction;
+    // @ts-ignore accessing private property
     return transactionWithContext._contexts;
   }
 
   // monkey-patch finish to store the context at finish time
   function monkeyPatchTransactionFinish(transaction: Transaction) {
-    const monkeyPatchedTransaction = transaction as Transaction & { _contexts: Contexts };
+    const monkeyPatchedTransaction = transaction as Transaction;
 
     // eslint-disable-next-line @typescript-eslint/unbound-method
     const originalFinish = monkeyPatchedTransaction.finish;
+    // @ts-ignore accessing private property
     monkeyPatchedTransaction._contexts = {};
     monkeyPatchedTransaction.finish = function (endTimestamp?: number | undefined) {
-      monkeyPatchedTransaction._contexts = (
-        transaction._hub.getScope() as unknown as Scope & { _contexts: Contexts }
-      )._contexts;
+      // @ts-ignore accessing private property
+      monkeyPatchedTransaction._contexts = (transaction._hub.getScope() as unknown as Scope)._contexts;
 
       return originalFinish.apply(monkeyPatchedTransaction, [endTimestamp]);
     };

packages/tracing/src/transaction.ts

@@ -1,5 +1,7 @@
 import { getCurrentHub, Hub } from '@sentry/core';
 import {
+  Context,
+  Contexts,
   DynamicSamplingContext,
   Event,
   Measurements,
@@ -25,6 +27,8 @@ export class Transaction extends SpanClass implements TransactionInterface {
 
   private _measurements: Measurements = {};
 
+  private _contexts: Contexts = {};
+
   private _trimEnd?: boolean;
 
   private _frozenDynamicSamplingContext: Readonly<Partial<DynamicSamplingContext>> | undefined = undefined;
@@ -105,6 +109,18 @@ export class Transaction extends SpanClass implements TransactionInterface {
     this.spanRecorder.add(this);
   }
 
+  /**
+   * @inheritDoc
+   */
+  public setContext(key: string, context: Context | null): void {
+    if (context === null) {
+      // eslint-disable-next-line @typescript-eslint/no-dynamic-delete
+      delete this._contexts[key];
+    } else {
+      this._contexts[key] = context;
+    }
+  }
+
   /**
    * @inheritDoc
    */
@@ -163,6 +179,8 @@ export class Transaction extends SpanClass implements TransactionInterface {
 
     const transaction: Event = {
       contexts: {
+        ...this._contexts,
+        // We don't want to override trace context
         trace: this.getTraceContext(),
       },
       spans: finishedSpans,

packages/tracing/test/transaction.test.ts

@@ -1,6 +1,14 @@
+import { BrowserClient, Hub } from '@sentry/browser';
+
+import { addExtensionMethods } from '../src';
 import { Transaction } from '../src/transaction';
+import { getDefaultBrowserClientOptions } from './testutils';
 
 describe('`Transaction` class', () => {
+  beforeAll(() => {
+    addExtensionMethods();
+  });
+
   describe('transaction name source', () => {
     it('sets source in constructor if provided', () => {
       const transaction = new Transaction({ name: 'dogpark', metadata: { source: 'route' } });
@@ -192,4 +200,127 @@ describe('`Transaction` class', () => {
       });
     });
   });
+
+  describe('setContext', () => {
+    it('sets context', () => {
+      const transaction = new Transaction({ name: 'dogpark' });
+      transaction.setContext('foo', {
+        key: 'val',
+        key2: 'val2',
+      });
+
+      // @ts-ignore accessing private property
+      expect(transaction._contexts).toEqual({
+        foo: {
+          key: 'val',
+          key2: 'val2',
+        },
+      });
+    });
+
+    it('overwrites context', () => {
+      const transaction = new Transaction({ name: 'dogpark' });
+      transaction.setContext('foo', {
+        key: 'val',
+        key2: 'val2',
+      });
+      transaction.setContext('foo', {
+        key3: 'val3',
+      });
+
+      // @ts-ignore accessing private property
+      expect(transaction._contexts).toEqual({
+        foo: {
+          key3: 'val3',
+        },
+      });
+    });
+
+    it('merges context', () => {
+      const transaction = new Transaction({ name: 'dogpark' });
+      transaction.setContext('foo', {
+        key: 'val',
+        key2: 'val2',
+      });
+      transaction.setContext('bar', {
+        anotherKey: 'anotherVal',
+      });
+
+      // @ts-ignore accessing private property
+      expect(transaction._contexts).toEqual({
+        foo: {
+          key: 'val',
+          key2: 'val2',
+        },
+        bar: {
+          anotherKey: 'anotherVal',
+        },
+      });
+    });
+
+    it('deletes context', () => {
+      const transaction = new Transaction({ name: 'dogpark' });
+      transaction.setContext('foo', {
+        key: 'val',
+        key2: 'val2',
+      });
+      transaction.setContext('foo', null);
+
+      // @ts-ignore accessing private property
+      expect(transaction._contexts).toEqual({});
+    });
+
+    it('sets contexts on the event', () => {
+      const options = getDefaultBrowserClientOptions({ tracesSampleRate: 1 });
+      const client = new BrowserClient(options);
+      const hub = new Hub(client);
+
+      jest.spyOn(hub, 'captureEvent');
+
+      const transaction = hub.startTransaction({ name: 'dogpark' });
+      transaction.setContext('foo', { key: 'val' });
+      transaction.finish();
+
+      // eslint-disable-next-line @typescript-eslint/unbound-method
+      expect(hub.captureEvent).toHaveBeenCalledTimes(1);
+      // eslint-disable-next-line @typescript-eslint/unbound-method
+      expect(hub.captureEvent).toHaveBeenLastCalledWith(
+        expect.objectContaining({
+          contexts: {
+            foo: { key: 'val' },
+            trace: {
+              span_id: transaction.spanId,
+              trace_id: transaction.traceId,
+            },
+          },
+        }),
+      );
+    });
+
+    it('does not override trace context', () => {
+      const options = getDefaultBrowserClientOptions({ tracesSampleRate: 1 });
+      const client = new BrowserClient(options);
+      const hub = new Hub(client);
+
+      jest.spyOn(hub, 'captureEvent');
+
+      const transaction = hub.startTransaction({ name: 'dogpark' });
+      transaction.setContext('trace', { key: 'val' });
+      transaction.finish();
+
+      // eslint-disable-next-line @typescript-eslint/unbound-method
+      expect(hub.captureEvent).toHaveBeenCalledTimes(1);
+      // eslint-disable-next-line @typescript-eslint/unbound-method
+      expect(hub.captureEvent).toHaveBeenLastCalledWith(
+        expect.objectContaining({
+          contexts: {
+            trace: {
+              span_id: transaction.spanId,
+              trace_id: transaction.traceId,
+            },
+          },
+        }),
+      );
+    });
+  });
 });

packages/types/src/transaction.ts

@@ -1,3 +1,4 @@
+import { Context } from './context';
 import { DynamicSamplingContext } from './envelope';
 import { Instrumenter } from './instrumenter';
 import { MeasurementUnit } from './measurement';
@@ -82,6 +83,11 @@ export interface Transaction extends TransactionContext, Span {
    */
   setName(name: string, source?: TransactionMetadata['source']): void;
 
+  /**
+   * Set the context of a transaction event
+   */
+  setContext(key: string, context: Context): void;
+
   /**
    * Set observed measurement for this transaction.
    *