feat(tracing): Add context update methods to Span and Transaction (#3192)

* feat(tracing): Add context update methods to Span and Transaction

* ref: Return this from updateWithContext instead of Span/Transaction

Author: jennmueng

packages/tracing/src/span.ts

@@ -246,6 +246,44 @@ export class Span implements SpanInterface {
     return `${this.traceId}-${this.spanId}${sampledString}`;
   }
 
+  /**
+   * @inheritDoc
+   */
+  public toContext(): SpanContext {
+    return dropUndefinedKeys({
+      data: this.data,
+      description: this.description,
+      endTimestamp: this.endTimestamp,
+      op: this.op,
+      parentSpanId: this.parentSpanId,
+      sampled: this.sampled,
+      spanId: this.spanId,
+      startTimestamp: this.startTimestamp,
+      status: this.status,
+      tags: this.tags,
+      traceId: this.traceId,
+    });
+  }
+
+  /**
+   * @inheritDoc
+   */
+  public updateWithContext(spanContext: SpanContext): this {
+    this.data = spanContext.data ?? {};
+    this.description = spanContext.description;
+    this.endTimestamp = spanContext.endTimestamp;
+    this.op = spanContext.op;
+    this.parentSpanId = spanContext.parentSpanId;
+    this.sampled = spanContext.sampled;
+    this.spanId = spanContext.spanId ?? this.spanId;
+    this.startTimestamp = spanContext.startTimestamp ?? this.startTimestamp;
+    this.status = spanContext.status;
+    this.tags = spanContext.tags ?? {};
+    this.traceId = spanContext.traceId ?? this.traceId;
+
+    return this;
+  }
+
   /**
    * @inheritDoc
    */

packages/tracing/src/transaction.ts

@@ -1,6 +1,6 @@
 import { getCurrentHub, Hub } from '@sentry/hub';
 import { Event, Measurements, Transaction as TransactionInterface, TransactionContext } from '@sentry/types';
-import { isInstanceOf, logger } from '@sentry/utils';
+import { dropUndefinedKeys, isInstanceOf, logger } from '@sentry/utils';
 
 import { Span as SpanClass, SpanRecorder } from './span';
 
@@ -14,7 +14,7 @@ export class Transaction extends SpanClass implements TransactionInterface {
    */
   private readonly _hub: Hub = (getCurrentHub() as unknown) as Hub;
 
-  private readonly _trimEnd?: boolean;
+  private _trimEnd?: boolean;
 
   /**
    * This constructor should never be called manually. Those instrumenting tracing should use
@@ -119,4 +119,30 @@ export class Transaction extends SpanClass implements TransactionInterface {
 
     return this._hub.captureEvent(transaction);
   }
+
+  /**
+   * @inheritDoc
+   */
+  public toContext(): TransactionContext {
+    const spanContext = super.toContext();
+
+    return dropUndefinedKeys({
+      ...spanContext,
+      name: this.name,
+      trimEnd: this._trimEnd,
+    });
+  }
+
+  /**
+   * @inheritDoc
+   */
+  public updateWithContext(transactionContext: TransactionContext): this {
+    super.updateWithContext(transactionContext);
+
+    this.name = transactionContext.name ?? '';
+
+    this._trimEnd = transactionContext.trimEnd;
+
+    return this;
+  }
 }

packages/tracing/test/span.test.ts

@@ -287,4 +287,86 @@ describe('Span', () => {
       });
     });
   });
+
+  describe('toContext and updateWithContext', () => {
+    test('toContext should return correct context', () => {
+      const originalContext = { traceId: 'a', spanId: 'b', sampled: false, description: 'test', op: 'op' };
+      const span = new Span(originalContext);
+
+      const newContext = span.toContext();
+
+      expect(newContext).toStrictEqual({
+        ...originalContext,
+        data: {},
+        spanId: expect.any(String),
+        startTimestamp: expect.any(Number),
+        tags: {},
+        traceId: expect.any(String),
+      });
+    });
+
+    test('updateWithContext should completely change span properties', () => {
+      const originalContext = {
+        traceId: 'a',
+        spanId: 'b',
+        sampled: false,
+        description: 'test',
+        op: 'op',
+        tags: {
+          tag0: 'hello',
+        },
+      };
+      const span = new Span(originalContext);
+
+      span.updateWithContext({
+        traceId: 'c',
+        spanId: 'd',
+        sampled: true,
+      });
+
+      expect(span.traceId).toBe('c');
+      expect(span.spanId).toBe('d');
+      expect(span.sampled).toBe(true);
+      expect(span.description).toBe(undefined);
+      expect(span.op).toBe(undefined);
+      expect(span.tags).toStrictEqual({});
+    });
+
+    test('using toContext and updateWithContext together should update only changed properties', () => {
+      const originalContext = {
+        traceId: 'a',
+        spanId: 'b',
+        sampled: false,
+        description: 'test',
+        op: 'op',
+        tags: { tag0: 'hello' },
+        data: { data0: 'foo' },
+      };
+      const span = new Span(originalContext);
+
+      const newContext = {
+        ...span.toContext(),
+        description: 'new',
+        endTimestamp: 1,
+        op: 'new-op',
+        sampled: true,
+        tags: {
+          tag1: 'bye',
+        },
+      };
+
+      if (newContext.data) newContext.data.data1 = 'bar';
+
+      span.updateWithContext(newContext);
+
+      expect(span.traceId).toBe('a');
+      expect(span.spanId).toBe('b');
+      expect(span.description).toBe('new');
+      expect(span.endTimestamp).toBe(1);
+      expect(span.op).toBe('new-op');
+      expect(span.sampled).toBe(true);
+      expect(span.tags).toStrictEqual({ tag1: 'bye' });
+      expect(span.data).toStrictEqual({ data0: 'foo', data1: 'bar' });
+    });
+  });
 });

packages/types/src/span.ts

@@ -152,6 +152,12 @@ export interface Span extends SpanContext {
   /** Return a traceparent compatible header string */
   toTraceparent(): string;
 
+  /** Returns the current span properties as a `SpanContext` */
+  toContext(): SpanContext;
+
+  /** Updates the current span with a new `SpanContext` */
+  updateWithContext(spanContext: SpanContext): this;
+
   /** Convert the object to JSON for w. spans array info only */
   getTraceContext(): {
     data?: { [key: string]: any };

packages/types/src/transaction.ts

@@ -62,6 +62,12 @@ export interface Transaction extends TransactionContext, Span {
    * Set the name of the transaction
    */
   setName(name: string): void;
+
+  /** Returns the current transaction properties as a `TransactionContext` */
+  toContext(): TransactionContext;
+
+  /** Updates the current transaction with a new `TransactionContext` */
+  updateWithContext(transactionContext: TransactionContext): this;
 }
 
 /**