From 17d4069262967ac7cc7c01d305ecbc58481344eb Mon Sep 17 00:00:00 2001
From: Francesco Novy <francesco.novy@sentry.io>
Date: Tue, 21 Nov 2023 10:10:21 +0100
Subject: [PATCH] feat(core): Allow to use `continueTrace` without callback

---
 packages/core/src/tracing/trace.ts           | 25 +++++++++++-
 packages/core/test/lib/tracing/trace.test.ts | 41 ++++++++++++++++++++
 2 files changed, 64 insertions(+), 2 deletions(-)

diff --git a/packages/core/src/tracing/trace.ts b/packages/core/src/tracing/trace.ts
index 4572eed79ee9..33e19ec7ecb4 100644
--- a/packages/core/src/tracing/trace.ts
+++ b/packages/core/src/tracing/trace.ts
@@ -203,6 +203,23 @@ export function getActiveSpan(): Span | undefined {
   return getCurrentHub().getScope().getSpan();
 }
 
+export function continueTrace({
+  sentryTrace,
+  baggage,
+}: {
+  sentryTrace: Parameters<typeof tracingContextFromHeaders>[0];
+  baggage: Parameters<typeof tracingContextFromHeaders>[1];
+}): Partial<TransactionContext>;
+export function continueTrace<V>(
+  {
+    sentryTrace,
+    baggage,
+  }: {
+    sentryTrace: Parameters<typeof tracingContextFromHeaders>[0];
+    baggage: Parameters<typeof tracingContextFromHeaders>[1];
+  },
+  callback: (transactionContext: Partial<TransactionContext>) => V,
+): V;
 /**
  * Continue a trace from `sentry-trace` and `baggage` values.
  * These values can be obtained from incoming request headers,
@@ -219,8 +236,8 @@ export function continueTrace<V>(
     sentryTrace: Parameters<typeof tracingContextFromHeaders>[0];
     baggage: Parameters<typeof tracingContextFromHeaders>[1];
   },
-  callback: (transactionContext: Partial<TransactionContext>) => V,
-): V {
+  callback?: (transactionContext: Partial<TransactionContext>) => V,
+): V | Partial<TransactionContext> {
   const hub = getCurrentHub();
   const currentScope = hub.getScope();
 
@@ -242,6 +259,10 @@ export function continueTrace<V>(
     }),
   };
 
+  if (!callback) {
+    return transactionContext;
+  }
+
   return callback(transactionContext);
 }
 
diff --git a/packages/core/test/lib/tracing/trace.test.ts b/packages/core/test/lib/tracing/trace.test.ts
index 144ec35f1f0e..b8f2d584434e 100644
--- a/packages/core/test/lib/tracing/trace.test.ts
+++ b/packages/core/test/lib/tracing/trace.test.ts
@@ -320,4 +320,45 @@ describe('continueTrace', () => {
 
     expect(scope['_sdkProcessingMetadata']).toEqual({});
   });
+
+  it('returns response of callback', () => {
+    const expectedContext = {
+      metadata: {
+        dynamicSamplingContext: {},
+      },
+      parentSampled: false,
+      parentSpanId: '1121201211212012',
+      traceId: '12312012123120121231201212312012',
+    };
+
+    const result = continueTrace(
+      {
+        sentryTrace: '12312012123120121231201212312012-1121201211212012-0',
+        baggage: undefined,
+      },
+      ctx => {
+        return { ctx };
+      },
+    );
+
+    expect(result).toEqual({ ctx: expectedContext });
+  });
+
+  it('works without a callback', () => {
+    const expectedContext = {
+      metadata: {
+        dynamicSamplingContext: {},
+      },
+      parentSampled: false,
+      parentSpanId: '1121201211212012',
+      traceId: '12312012123120121231201212312012',
+    };
+
+    const ctx = continueTrace({
+      sentryTrace: '12312012123120121231201212312012-1121201211212012-0',
+      baggage: undefined,
+    });
+
+    expect(ctx).toEqual(expectedContext);
+  });
 });