Skip to content

Commit

Permalink
Inject trace context into SNS message attributes as binary
Browse files Browse the repository at this point in the history
  • Loading branch information
nhulston committed Oct 10, 2024
1 parent b71cdd5 commit 80be4e3
Show file tree
Hide file tree
Showing 3 changed files with 17 additions and 15 deletions.
6 changes: 3 additions & 3 deletions contrib/aws/aws-sdk-go-v2/aws/aws_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -508,12 +508,12 @@ func TestAppendMiddlewareSnsPublish(t *testing.T) {
assert.NotNil(t, tt.publishInput.MessageAttributes)
assert.Contains(t, tt.publishInput.MessageAttributes, "_datadog")
ddAttr := tt.publishInput.MessageAttributes["_datadog"]
assert.Equal(t, "String", *ddAttr.DataType)
assert.NotEmpty(t, *ddAttr.StringValue)
assert.Equal(t, "Binary", *ddAttr.DataType)
assert.NotEmpty(t, ddAttr.BinaryValue)

// Decode and verify the injected trace context
var traceContext map[string]string
err := json.Unmarshal([]byte(*ddAttr.StringValue), &traceContext)
err := json.Unmarshal(ddAttr.BinaryValue, &traceContext)
assert.NoError(t, err)
assert.Contains(t, traceContext, "x-datadog-trace-id")
assert.Contains(t, traceContext, "x-datadog-parent-id")
Expand Down
6 changes: 4 additions & 2 deletions contrib/aws/internal/sns/sns.go
Original file line number Diff line number Diff line change
Expand Up @@ -87,8 +87,10 @@ func injectTraceContext(ctx context.Context, messageAttributes map[string]types.
return
}

// Use Binary since SNS subscription filter policies fail silently with JSON
// strings. https://github.com/DataDog/datadog-lambda-js/pull/269
messageAttributes[datadogKey] = types.MessageAttributeValue{
DataType: aws.String("String"),
StringValue: aws.String(string(jsonBytes)),
DataType: aws.String("Binary"),
BinaryValue: jsonBytes,
}
}
20 changes: 10 additions & 10 deletions contrib/aws/internal/sns/sns_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -49,9 +49,9 @@ func TestEnrichOperation(t *testing.T) {
require.NotNil(t, params.MessageAttributes)
assert.Contains(t, params.MessageAttributes, datadogKey)
assert.NotNil(t, params.MessageAttributes[datadogKey].DataType)
assert.Equal(t, "String", *params.MessageAttributes[datadogKey].DataType)
assert.NotNil(t, params.MessageAttributes[datadogKey].StringValue)
assert.NotEmpty(t, *params.MessageAttributes[datadogKey].StringValue)
assert.Equal(t, "Binary", *params.MessageAttributes[datadogKey].DataType)
assert.NotNil(t, params.MessageAttributes[datadogKey].BinaryValue)
assert.NotEmpty(t, params.MessageAttributes[datadogKey].BinaryValue)
},
},
{
Expand Down Expand Up @@ -87,9 +87,9 @@ func TestEnrichOperation(t *testing.T) {
require.NotNil(t, entry.MessageAttributes)
assert.Contains(t, entry.MessageAttributes, datadogKey)
assert.NotNil(t, entry.MessageAttributes[datadogKey].DataType)
assert.Equal(t, "String", *entry.MessageAttributes[datadogKey].DataType)
assert.NotNil(t, entry.MessageAttributes[datadogKey].StringValue)
assert.NotEmpty(t, *entry.MessageAttributes[datadogKey].StringValue)
assert.Equal(t, "Binary", *entry.MessageAttributes[datadogKey].DataType)
assert.NotNil(t, entry.MessageAttributes[datadogKey].BinaryValue)
assert.NotEmpty(t, entry.MessageAttributes[datadogKey].BinaryValue)
}
},
},
Expand Down Expand Up @@ -158,12 +158,12 @@ func TestInjectTraceContext(t *testing.T) {
if tt.expectInjection {
assert.Contains(t, messageAttributes, datadogKey)
assert.NotNil(t, messageAttributes[datadogKey].DataType)
assert.Equal(t, "String", *messageAttributes[datadogKey].DataType)
assert.NotNil(t, messageAttributes[datadogKey].StringValue)
assert.NotEmpty(t, *messageAttributes[datadogKey].StringValue)
assert.Equal(t, "Binary", *messageAttributes[datadogKey].DataType)
assert.NotNil(t, messageAttributes[datadogKey].BinaryValue)
assert.NotEmpty(t, messageAttributes[datadogKey].BinaryValue)

carrier := tracer.TextMapCarrier{}
err := json.Unmarshal([]byte(*messageAttributes[datadogKey].StringValue), &carrier)
err := json.Unmarshal(messageAttributes[datadogKey].BinaryValue, &carrier)
assert.NoError(t, err)

extractedSpanContext, err := tracer.Extract(carrier)
Expand Down

0 comments on commit 80be4e3

Please sign in to comment.