feat: support reading trace context from AppSync event (#229)

* feat: support reading trace context from appsync event

* update isAppSyncResolverEvent to a more selective check on info.selectionSetGraphQL

* lint

Co-authored-by: Maxime David <[email protected]>

Author: bigkraig

src/trace/context.spec.ts

@@ -200,6 +200,26 @@ describe("readTraceFromEvent", () => {
       source: Source.Event,
     });
   });
+  it("can read from appsync source", () => {
+    const result = readTraceFromEvent({
+      info: {
+        selectionSetGraphQL: "{ items }",
+      },
+      request: {
+        headers: {
+          "x-datadog-parent-id": "797643193680388254",
+          "x-datadog-sampling-priority": "2",
+          "x-datadog-trace-id": "4110911582297405557",
+        },
+      },
+    });
+    expect(result).toEqual({
+      parentID: "797643193680388254",
+      sampleMode: SampleMode.USER_KEEP,
+      traceID: "4110911582297405557",
+      source: Source.Event,
+    });
+  });
   it("can read from sqs source", () => {
     const result = readTraceFromEvent({
       Records: [

src/trace/context.ts

@@ -3,7 +3,7 @@ import { BigNumber } from "bignumber.js";
 import { randomBytes } from "crypto";
 import { createSocket, Socket } from "dgram";
 import { logDebug, logError } from "../utils";
-import { isSQSEvent } from "../utils/event-type-guards";
+import { isAppSyncResolverEvent, isSQSEvent } from "../utils/event-type-guards";
 import {
   awsXrayDaemonAddressEnvVar,
   parentIDHeader,
@@ -183,6 +183,11 @@ export function sendXraySubsegment(segment: string) {
   }
 }
 
+export function readTraceFromAppSyncEvent(event: any): TraceContext | undefined {
+  event.headers = event.request.headers;
+  return readTraceFromHTTPEvent(event);
+}
+
 export function readTraceFromSQSEvent(event: SQSEvent): TraceContext | undefined {
   if (
     event.Records[0].messageAttributes &&
@@ -322,6 +327,10 @@ export function readTraceFromEvent(event: any): TraceContext | undefined {
     return readTraceFromHTTPEvent(event);
   }
 
+  if (isAppSyncResolverEvent(event)) {
+    return readTraceFromAppSyncEvent(event);
+  }
+
   if (isSQSEvent(event)) {
     return readTraceFromSQSEvent(event);
   }

src/utils/event-type-guards.ts

@@ -1,6 +1,7 @@
 import {
   APIGatewayEvent,
   APIGatewayProxyEventV2,
+  AppSyncResolverEvent,
   ALBEvent,
   CloudWatchLogsEvent,
   ScheduledEvent,
@@ -58,3 +59,7 @@ export function isSNSEvent(event: any): event is SNSEvent {
 export function isSQSEvent(event: any): event is SQSEvent {
   return Array.isArray(event.Records) && event.Records.length > 0 && event.Records[0].eventSource === "aws:sqs";
 }
+
+export function isAppSyncResolverEvent(event: any): event is AppSyncResolverEvent<any> {
+  return event.info !== undefined && event.info.selectionSetGraphQL !== undefined;
+}