OB-35469: fix: add subscription request parent span as child of disco…

…very request

pkg/handler/subscriber/wrap_sqs_context.go

@@ -0,0 +1,48 @@
+package subscriber
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+
+	"github.com/aws/aws-lambda-go/events"
+
+	"github.com/aws/aws-lambda-go/lambda"
+)
+
+type SQSCheckerWrappedHandler struct {
+	handler lambda.Handler
+}
+
+// Compile time check our Handler implements lambda.Handler.
+var _ lambda.Handler = SQSCheckerWrappedHandler{}
+
+// Invoke checks if the incoming payload is from a SQS event and, if so,
+// extracts the context from the SQS message and injects it into the context.
+// This means that the spans created with this context will appear as children
+// of a span in the discover request that created the SQS event
+func (h SQSCheckerWrappedHandler) Invoke(ctx context.Context, payload []byte) ([]byte, error) {
+
+	var event events.SQSEvent
+	dec := json.NewDecoder(bytes.NewReader(payload))
+	dec.DisallowUnknownFields()
+	if err := dec.Decode(&event); err == nil {
+		for _, record := range event.Records {
+			ctx = NewSQSCarrier().Extract(ctx, record.MessageAttributes)
+			break
+		}
+	}
+
+	response, err := h.handler.Invoke(ctx, payload)
+	if err != nil {
+		return nil, err
+	}
+
+	return response, nil
+}
+
+// WrapHandlerSQSCheck Provides a Handler which wraps an existing Handler while
+// injecting the SQS context into the context if the payload is a SQS event.
+func WrapHandlerSQSCheck(handler lambda.Handler) lambda.Handler {
+	return SQSCheckerWrappedHandler{handler: handler}
+}

pkg/tracing/mux.go

@@ -2,6 +2,7 @@ package tracing
 
 import (
 	"context"
+	"github.com/observeinc/aws-sam-apps/pkg/handler/subscriber"
 
 	"github.com/aws/aws-lambda-go/lambda"
 	"go.opentelemetry.io/contrib/instrumentation/github.com/aws/aws-lambda-go/otellambda"
@@ -14,11 +15,11 @@ import (
 var payloadKey = attribute.Key("payload")
 
 func NewLambdaHandler(handler lambda.Handler, tp *sdktrace.TracerProvider) lambda.Handler {
-	return otellambda.WrapHandler(
+	return subscriber.WrapHandlerSQSCheck(otellambda.WrapHandler(
 		&LambdaHandler{handler},
 		otellambda.WithTracerProvider(tp),
 		otellambda.WithFlusher(tp),
-	)
+	))
 }
 
 type LambdaHandler struct {