Fix SQS message attributes mapping with custom data types (#746)

Co-authored-by: Normand Poirier <[email protected]>

Author: npoirier

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

@@ -181,11 +181,16 @@ private Object getValue(Map.Entry<String, MessageAttributeValue> entry) {
 		MessageAttributeValue value = entry.getValue();
 		String dataType = value.dataType();
 		Assert.notNull(dataType, "dataType must not be null");
-		return MessageAttributeDataTypes.STRING.equals(dataType)
-			? value.stringValue()
-			: MessageAttributeDataTypes.BINARY.equals(dataType)
-				? value.binaryValue()
-				: getNumberValue(value);
+
+		if (dataType.contains(".")) {
+			dataType = dataType.substring(0, dataType.indexOf('.'));
+		}
+
+		return switch (dataType) {
+			case MessageAttributeDataTypes.NUMBER -> getNumberValue(value);
+			case MessageAttributeDataTypes.BINARY -> value.binaryValue();
+			default -> value.stringValue();
+		};
 	}
 
 	private Map<String, String> getMessageSystemAttributesAsHeaders(Message source) {

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

@@ -54,13 +54,43 @@ void shouldAddExtraHeader() {
 
 	@Test
 	void shouldAddStringMessageAttributes() {
+		SqsHeaderMapper mapper = new SqsHeaderMapper();
+		String headerName = "stringAttribute";
+		String headerValue = "myString";
+		Message message = Message.builder().body("payload")
+			.messageAttributes(
+				Map.of(headerName,
+					MessageAttributeValue.builder().dataType(MessageAttributeDataTypes.STRING)
+						.stringValue(headerValue).build()))
+			.messageId(UUID.randomUUID().toString()).build();
+		MessageHeaders headers = mapper.toHeaders(message);
+		assertThat(headers.get(headerName)).isEqualTo(headerValue);
+	}
+
+	@Test
+	void shouldAddStringCustomMessageAttributes() {
+		SqsHeaderMapper mapper = new SqsHeaderMapper();
+		String headerName = "stringAttribute";
+		String headerValue = "myString";
+		Message message = Message.builder().body("payload")
+			.messageAttributes(
+				Map.of(headerName,
+					MessageAttributeValue.builder().dataType(MessageAttributeDataTypes.STRING + ".Array")
+						.stringValue(headerValue).build()))
+			.messageId(UUID.randomUUID().toString()).build();
+		MessageHeaders headers = mapper.toHeaders(message);
+		assertThat(headers.get(headerName)).isEqualTo(headerValue);
+	}
+
+	@Test
+	void shouldDefaultToStringIfDataTypeUnknownMessageAttributes() {
 		SqsHeaderMapper mapper = new SqsHeaderMapper();
 		String headerName = "stringAttribute";
 		String headerValue = "myString";
 		Message message = Message.builder().body("payload")
 				.messageAttributes(
 						Map.of(headerName,
-								MessageAttributeValue.builder().dataType(MessageAttributeDataTypes.STRING)
+								MessageAttributeValue.builder().dataType("invalid data type")
 										.stringValue(headerValue).build()))
 				.messageId(UUID.randomUUID().toString()).build();
 		MessageHeaders headers = mapper.toHeaders(message);
@@ -69,13 +99,28 @@ void shouldAddStringMessageAttributes() {
 
 	@Test
 	void shouldAddBinaryMessageAttributes() {
+		SqsHeaderMapper mapper = new SqsHeaderMapper();
+		String headerName = "stringAttribute";
+		SdkBytes headerValue = SdkBytes.fromUtf8String("myString");
+		Message message = Message.builder().body("payload")
+			.messageAttributes(
+				Map.of(headerName,
+					MessageAttributeValue.builder().dataType(MessageAttributeDataTypes.BINARY)
+						.binaryValue(headerValue).build()))
+			.messageId(UUID.randomUUID().toString()).build();
+		MessageHeaders headers = mapper.toHeaders(message);
+		assertThat(headers.get(headerName)).isEqualTo(headerValue);
+	}
+
+	@Test
+	void shouldAddBinaryCustomMessageAttributes() {
 		SqsHeaderMapper mapper = new SqsHeaderMapper();
 		String headerName = "stringAttribute";
 		SdkBytes headerValue = SdkBytes.fromUtf8String("myString");
 		Message message = Message.builder().body("payload")
 				.messageAttributes(
 						Map.of(headerName,
-								MessageAttributeValue.builder().dataType(MessageAttributeDataTypes.BINARY)
+								MessageAttributeValue.builder().dataType(MessageAttributeDataTypes.BINARY + ".protobuf")
 										.binaryValue(headerValue).build()))
 				.messageId(UUID.randomUUID().toString()).build();
 		MessageHeaders headers = mapper.toHeaders(message);