propogation fix

baselime · Mar 28, 2024 · 6e6278d · 6e6278d
1 parent 3a6c771
commit 6e6278d
Show file tree

Hide file tree

Showing 5 changed files with 12 additions and 12 deletions.
diff --git a/examples/state-machine/packages/functions/src/lambda.ts b/examples/state-machine/packages/functions/src/lambda.ts
@@ -1,4 +1,4 @@
-// import { withOpenTelemetry } from "./tracing";
+import { withOpenTelemetry } from "./tracing";
 import { ApiHandler } from "sst/node/api";
 
 export const handler = withOpenTelemetry(ApiHandler(async (_evt) => {

diff --git a/examples/state-machine/packages/functions/src/task-one.ts b/examples/state-machine/packages/functions/src/task-one.ts
@@ -1,8 +1,8 @@
 import { withOpenTelemetry } from "./tracing";
 
-export const handler = withOpenTelemetry(async (_,__, callback) => {
+export const handler = async (_,__, callback) => {
   callback(null, {
     statusCode: 200,
     body: `Hello world. The time is ${new Date().toISOString()}`,
   });
-});
+};
diff --git a/examples/state-machine/packages/functions/src/task-two.ts b/examples/state-machine/packages/functions/src/task-two.ts
@@ -1,8 +1,8 @@
-import { ApiHandler } from "sst/node/api";
 
-export const handler = ApiHandler(async (_evt) => {
+export const handler = async (_evt) => {
+  console.log(_evt)
   return {
     statusCode: 200,
     body: `Hello world. The time is ${new Date().toISOString()}`,
   };
-});
+};
diff --git a/examples/state-machine/stacks/MyStack.ts b/examples/state-machine/stacks/MyStack.ts
@@ -1,6 +1,6 @@
 import { StackContext, Function } from "sst/constructs";
-import { LambdaInvoke } from "aws-cdk-lib/aws-stepfunctions-tasks";
-import { Chain, Parallel, StateMachine, TaskInput, State} from "aws-cdk-lib/aws-stepfunctions";
+import { LambdaInvoke, SnsPublish } from "aws-cdk-lib/aws-stepfunctions-tasks";
+import { Chain, Parallel, StateMachine, TaskInput} from "aws-cdk-lib/aws-stepfunctions";
 export function API({ stack }: StackContext) {
 
   stack.addDefaultFunctionEnv({
@@ -19,10 +19,10 @@ export function API({ stack }: StackContext) {
     }),
     payload: TaskInput.fromObject({
       code: TaskInput.fromJsonPathAt("$.Payload.statusCode").value,
-      _baselime: TaskInput.fromJsonPathAt("$.Payload._baselime").value
+      _baselime: TaskInput.fromJsonPathAt("$.Payload.['_baselime', 'null']").value
     })
   })
-
+  
   const taskTwoB = new LambdaInvoke(stack, "TaskTwoB", {
     lambdaFunction: new Function(stack, "task-two-b", {
       handler: "packages/functions/src/task-two.handler",

diff --git a/src/lambda/propation.ts b/src/lambda/propation.ts
@@ -37,7 +37,7 @@ export function extractContext(service: string, event: any): { parent?: Context,
             if (Array.isArray(event)) {
                 return {
                     links: event.map((parent) => {
-                        const traceparent = parent._baselime?.traceparent || parent.Payload?._baselime?.traceparent;
+                        const traceparent = parent._baselime?.traceparent || parent.Payload?._baselime?.traceparent || parent._baselime?._baselime?.traceparent;
                         if (!traceparent) {
                             return
                         }
@@ -51,7 +51,7 @@ export function extractContext(service: string, event: any): { parent?: Context,
                     }).filter(el => el)
                 }
             }
-            const traceparent = event._baselime?.traceparent || event.Payload?._baselime?.traceparent;
+            const traceparent = event._baselime?.traceparent || event.Payload?._baselime?.traceparent || event._baselime?._baselime?.traceparent;
 
             return {
                 parent: propagation.extract(context.active(), { traceparent }, headerGetter)