feat(tracing): Add transaction.setContext method

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 = { ...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,13 @@
 import { Transaction } from '../src/transaction';
+import { addExtensionMethods } from '../src';
+import { getDefaultBrowserClientOptions } from './testutils';
+import { BrowserClient, Hub } from '@sentry/browser';
 
 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' } });
@@ -180,4 +187,123 @@ 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();
+
+      expect(hub.captureEvent).toHaveBeenCalledTimes(1);
+      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();
+
+      expect(hub.captureEvent).toHaveBeenCalledTimes(1);
+      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 { MeasurementUnit } from './measurement';
 import { ExtractedNodeRequestData, Primitive, WorkerLocation } from './misc';
@@ -76,6 +77,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.
    *