Skip to content

Commit a9f51fc

Browse files
acomofcgacomofcg
authored
Fix: Propagate AWSTraceHeader (#1287)
Fixes #1262
1 parent a08ebc1 commit a9f51fc

File tree

4 files changed

+72
-20
lines changed

4 files changed

+72
-20
lines changed

‎spring-cloud-aws-sqs/src/main/java/io/awspring/cloud/sqs/operations/SqsTemplate.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -454,7 +454,7 @@ private Integer getDelaySeconds(Message message) {
454454
private Map<MessageSystemAttributeNameForSends, MessageSystemAttributeValue> mapMessageSystemAttributes(
455455
Message message) {
456456
return message.attributes().entrySet().stream().filter(Predicate.not(entry -> isSkipAttribute(entry.getKey())))
457-
.collect(Collectors.toMap(entry -> MessageSystemAttributeNameForSends.fromValue(entry.getKey().name()),
457+
.collect(Collectors.toMap(entry -> MessageSystemAttributeNameForSends.fromValue(entry.getKey().toString()),
458458
entry -> MessageSystemAttributeValue.builder().dataType(MessageAttributeDataTypes.STRING)
459459
.stringValue(entry.getValue()).build()));
460460
}

‎spring-cloud-aws-sqs/src/main/java/io/awspring/cloud/sqs/support/converter/SqsHeaderMapper.java

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,10 @@ public Message fromHeaders(MessageHeaders headers) {
7979
attributes.put(MessageSystemAttributeName.MESSAGE_DEDUPLICATION_ID,
8080
headers.get(SqsHeaders.MessageSystemAttributes.SQS_MESSAGE_DEDUPLICATION_ID_HEADER, String.class));
8181
}
82+
if (headers.containsKey(SqsHeaders.MessageSystemAttributes.SQS_AWS_TRACE_HEADER)) {
83+
attributes.put(MessageSystemAttributeName.AWS_TRACE_HEADER,
84+
headers.get(SqsHeaders.MessageSystemAttributes.SQS_AWS_TRACE_HEADER, String.class));
85+
}
8286
Map<String, MessageAttributeValue> messageAttributes = headers.entrySet().stream()
8387
.filter(entry -> !isSkipHeader(entry.getKey())).collect(Collectors.toMap(Map.Entry::getKey,
8488
entry -> getMessageAttributeValue(entry.getKey(), entry.getValue())));
@@ -110,6 +114,7 @@ else if (messageHeaderValue instanceof ByteBuffer) {
110114
private boolean isSkipHeader(String headerName) {
111115
return SqsHeaders.MessageSystemAttributes.SQS_MESSAGE_GROUP_ID_HEADER.equals(headerName)
112116
|| SqsHeaders.MessageSystemAttributes.SQS_MESSAGE_DEDUPLICATION_ID_HEADER.equals(headerName)
117+
|| SqsHeaders.MessageSystemAttributes.SQS_AWS_TRACE_HEADER.equals(headerName)
113118
|| SqsHeaders.SQS_DELAY_HEADER.equals(headerName) || MessageHeaders.ID.equals(headerName)
114119
|| MessageHeaders.TIMESTAMP.equals(headerName);
115120
}

‎spring-cloud-aws-sqs/src/test/java/io/awspring/cloud/sqs/operations/SqsTemplateTests.java

Lines changed: 33 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,7 @@
2020
import static org.mockito.ArgumentMatchers.any;
2121
import static org.mockito.BDDMockito.given;
2222
import static org.mockito.BDDMockito.then;
23-
import static org.mockito.Mockito.mock;
24-
import static org.mockito.Mockito.never;
23+
import static org.mockito.Mockito.*;
2524

2625
import com.fasterxml.jackson.databind.ObjectMapper;
2726
import io.awspring.cloud.sqs.QueueAttributesResolvingException;
@@ -47,23 +46,7 @@
4746
import org.springframework.messaging.MessageHeaders;
4847
import org.springframework.messaging.support.MessageBuilder;
4948
import software.amazon.awssdk.services.sqs.SqsAsyncClient;
50-
import software.amazon.awssdk.services.sqs.model.CreateQueueResponse;
51-
import software.amazon.awssdk.services.sqs.model.DeleteMessageBatchRequest;
52-
import software.amazon.awssdk.services.sqs.model.DeleteMessageBatchResponse;
53-
import software.amazon.awssdk.services.sqs.model.GetQueueAttributesRequest;
54-
import software.amazon.awssdk.services.sqs.model.GetQueueAttributesResponse;
55-
import software.amazon.awssdk.services.sqs.model.GetQueueUrlRequest;
56-
import software.amazon.awssdk.services.sqs.model.GetQueueUrlResponse;
57-
import software.amazon.awssdk.services.sqs.model.MessageSystemAttributeName;
58-
import software.amazon.awssdk.services.sqs.model.QueueAttributeName;
59-
import software.amazon.awssdk.services.sqs.model.QueueDoesNotExistException;
60-
import software.amazon.awssdk.services.sqs.model.ReceiveMessageRequest;
61-
import software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse;
62-
import software.amazon.awssdk.services.sqs.model.SendMessageBatchRequest;
63-
import software.amazon.awssdk.services.sqs.model.SendMessageBatchRequestEntry;
64-
import software.amazon.awssdk.services.sqs.model.SendMessageBatchResponse;
65-
import software.amazon.awssdk.services.sqs.model.SendMessageRequest;
66-
import software.amazon.awssdk.services.sqs.model.SendMessageResponse;
49+
import software.amazon.awssdk.services.sqs.model.*;
6750

6851
/**
6952
* @author Tomaz Fernandes
@@ -1208,4 +1191,35 @@ void shouldReceiveBatchFifo() {
12081191

12091192
}
12101193

1194+
@Test
1195+
void shouldPropagateTracingAsMessageSystemAttribute() {
1196+
String queue = "test-queue";
1197+
GetQueueUrlResponse urlResponse = GetQueueUrlResponse.builder().queueUrl(queue).build();
1198+
given(mockClient.getQueueUrl(any(GetQueueUrlRequest.class)))
1199+
.willReturn(CompletableFuture.completedFuture(urlResponse));
1200+
mockQueueAttributes(mockClient, Map.of());
1201+
SendMessageResponse response = SendMessageResponse.builder().messageId(UUID.randomUUID().toString())
1202+
.sequenceNumber("123").build();
1203+
given(mockClient.sendMessage(any(SendMessageRequest.class)))
1204+
.willReturn(CompletableFuture.completedFuture(response));
1205+
1206+
SqsOperations sqsOperations = SqsTemplate.newSyncTemplate(mockClient);
1207+
SendResult<Object> result = sqsOperations.send(options -> options
1208+
.queue(queue)
1209+
.header(SqsHeaders.MessageSystemAttributes.SQS_AWS_TRACE_HEADER, "abc")
1210+
.payload("test")
1211+
);
1212+
1213+
assertThat(result).isNotNull();
1214+
1215+
ArgumentCaptor<SendMessageRequest> captor = ArgumentCaptor.forClass(SendMessageRequest.class);
1216+
then(mockClient).should().sendMessage(captor.capture());
1217+
SendMessageRequest sendMessageRequest = captor.getValue();
1218+
1219+
assertThat(sendMessageRequest.messageSystemAttributes()).hasEntrySatisfying(
1220+
MessageSystemAttributeNameForSends.AWS_TRACE_HEADER,
1221+
value -> assertThat(value.stringValue()).isEqualTo("abc")
1222+
);
1223+
}
1224+
12111225
}

‎spring-cloud-aws-sqs/src/test/java/io/awspring/cloud/sqs/support/converter/SqsHeaderMapperTests.java

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@
2121
import java.util.Map;
2222
import java.util.UUID;
2323
import java.util.stream.Stream;
24+
25+
import io.awspring.cloud.sqs.listener.SqsHeaders;
2426
import org.junit.jupiter.api.Test;
2527
import org.junit.jupiter.params.ParameterizedTest;
2628
import org.junit.jupiter.params.provider.Arguments;
@@ -29,6 +31,7 @@
2931
import software.amazon.awssdk.core.SdkBytes;
3032
import software.amazon.awssdk.services.sqs.model.Message;
3133
import software.amazon.awssdk.services.sqs.model.MessageAttributeValue;
34+
import software.amazon.awssdk.services.sqs.model.MessageSystemAttributeName;
3235

3336
/**
3437
* Tests for {@link SqsHeaderMapper}.
@@ -139,6 +142,36 @@ void shouldAddNumberMessageAttributes() {
139142
assertThat(headers.get(headerName)).isEqualTo(headerValue);
140143
}
141144

145+
@Test
146+
void shouldCreateMessageWithSystemAttributesFromHeaders() {
147+
MessageHeaders headers = new MessageHeaders(
148+
Map.of(
149+
SqsHeaders.MessageSystemAttributes.SQS_MESSAGE_GROUP_ID_HEADER, "value1",
150+
SqsHeaders.MessageSystemAttributes.SQS_MESSAGE_DEDUPLICATION_ID_HEADER, "value2",
151+
SqsHeaders.MessageSystemAttributes.SQS_AWS_TRACE_HEADER, "value3",
152+
"customHeaderString", "customValueString",
153+
"customHeaderNumber", 42
154+
)
155+
);
156+
157+
SqsHeaderMapper mapper = new SqsHeaderMapper();
158+
Message message = mapper.fromHeaders(headers);
159+
160+
assertThat(message.attributes())
161+
.hasSize(3)
162+
.containsExactlyInAnyOrderEntriesOf(Map.of(
163+
MessageSystemAttributeName.MESSAGE_GROUP_ID, "value1",
164+
MessageSystemAttributeName.MESSAGE_DEDUPLICATION_ID, "value2",
165+
MessageSystemAttributeName.AWS_TRACE_HEADER, "value3"
166+
));
167+
assertThat(message.messageAttributes())
168+
.hasSize(2)
169+
.containsExactlyInAnyOrderEntriesOf(Map.of(
170+
"customHeaderString", MessageAttributeValue.builder().dataType(MessageAttributeDataTypes.STRING).stringValue("customValueString").build(),
171+
"customHeaderNumber", MessageAttributeValue.builder().dataType(MessageAttributeDataTypes.NUMBER + ".java.lang.Integer").stringValue("42").build()
172+
));
173+
}
174+
142175
@ParameterizedTest
143176
@MethodSource("validArguments")
144177
void createsMessageWithNumberHeader(String value, String type, Number expected) {

0 commit comments

Comments
 (0)