fix(sqs): UUID lost via marshalling

sqs/marshaler.go

@@ -15,65 +15,60 @@ type UnMarshaler interface {
 	Unmarshal(msg *sqs.Message) (*message.Message, error)
 }
 
+// UUIDHeaderKey is the key of the Pub/Sub attribute that carries Waterfall UUID.
+const UUIDHeaderKey = "_watermill_message_uuid"
+
 type DefaultMarshalerUnmarshaler struct{}
 
 func (d DefaultMarshalerUnmarshaler) Marshal(msg *message.Message) (*sqs.Message, error) {
+	attributes := make(map[string]*sqs.MessageAttributeValue, len(msg.Metadata)+1)
+	for k, v := range msg.Metadata {
+		attributes[k] = &sqs.MessageAttributeValue{
+			StringValue: aws.String(v),
+			DataType:    aws.String(AWSStringDataType),
+		}
+	}
+	attributes[UUIDHeaderKey] = &sqs.MessageAttributeValue{
+		StringValue: aws.String(msg.UUID),
+		DataType:    aws.String(AWSStringDataType),
+	}
+
 	return &sqs.Message{
-		MessageAttributes: metadataToAttributes(msg.Metadata),
+		MessageAttributes: attributes,
 		Body:              aws.String(string(msg.Payload)),
-		MessageId:         aws.String(msg.UUID),
 	}, nil
 }
 
 func (d DefaultMarshalerUnmarshaler) Unmarshal(msg *sqs.Message) (*message.Message, error) {
-	var uuid, payload string
-
-	if msg.MessageId != nil {
-		uuid = *msg.MessageId
-	}
-
+	var payload string
 	if msg.Body != nil {
 		payload = *msg.Body
 	}
 
-	wmsg := message.NewMessage(uuid, []byte(payload))
-	wmsg.Metadata = attributesToMetadata(msg.MessageAttributes)
-
-	return wmsg, nil
-}
-
-func metadataToAttributes(meta message.Metadata) map[string]*sqs.MessageAttributeValue {
-	attributes := make(map[string]*sqs.MessageAttributeValue)
+	wmsg := message.NewMessage("", []byte(payload))
+	wmsg.Metadata = make(message.Metadata, len(msg.MessageAttributes))
 
-	for k, v := range meta {
-		attributes[k] = &sqs.MessageAttributeValue{
-			StringValue: aws.String(v),
-			DataType:    aws.String(AWSStringDataType),
+	for k, v := range msg.MessageAttributes {
+		if v.DataType == nil {
+			continue
 		}
-	}
-
-	return attributes
-}
 
-func attributesToMetadata(attributes map[string]*sqs.MessageAttributeValue) message.Metadata {
-	meta := make(message.Metadata)
-
-	for k, v := range attributes {
-		if v.DataType == nil {
+		if k == UUIDHeaderKey && *v.DataType == AWSStringDataType {
+			wmsg.UUID = *v.StringValue
 			continue
 		}
 
 		switch *v.DataType {
 		case AWSStringDataType, AWSNumberDataType:
 			if v.StringValue != nil {
-				meta[k] = *v.StringValue
+				wmsg.Metadata[k] = *v.StringValue
 			}
 		case AWSBinaryDataType:
-			meta[k] = string(v.BinaryValue)
+			wmsg.Metadata[k] = string(v.BinaryValue)
 		}
 	}
 
-	return meta
+	return wmsg, nil
 }
 
 const (

sqs/marshaler_test.go

@@ -4,20 +4,12 @@ import (
 	"testing"
 
 	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/service/sqs"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
-func TestMetadaConverter(t *testing.T) {
-	metadata := message.Metadata{
-		"name": "some-event-name",
-		"test": "other-value",
-	}
-
-	metadataOut := attributesToMetadata(metadataToAttributes(metadata))
-
-	require.Equal(t, metadata, metadataOut, "metadata after transformations should be the same")
-}
-
 func TestMarshaler(t *testing.T) {
 	msg := message.NewMessage("some-id", []byte("payload"))
 	msg.Metadata = message.Metadata{
@@ -29,11 +21,68 @@ func TestMarshaler(t *testing.T) {
 
 	awsMsg, err := dum.Marshal(msg)
 	require.NoError(t, err)
+	assert.Equal(t, &sqs.Message{
+		Body: aws.String("payload"),
+		MessageAttributes: map[string]*sqs.MessageAttributeValue{
+			"_watermill_message_uuid": {
+				DataType:    aws.String(AWSStringDataType),
+				StringValue: aws.String("some-id"),
+			},
+			"name": {
+				DataType:    aws.String(AWSStringDataType),
+				StringValue: aws.String("some-event-name"),
+			},
+			"test": {
+				DataType:    aws.String(AWSStringDataType),
+				StringValue: aws.String("other-value"),
+			},
+		},
+	}, awsMsg)
+}
 
-	msgOut, err := dum.Unmarshal(awsMsg)
-	require.NoError(t, err)
+func TestUnmarshaler(t *testing.T) {
+	m := &DefaultMarshalerUnmarshaler{}
 
-	require.Equal(t, msg.UUID, msgOut.UUID)
-	require.Equal(t, msg.Payload, msgOut.Payload)
-	require.Equal(t, msg.Metadata, msgOut.Metadata)
+	msg, err := m.Unmarshal(&sqs.Message{
+		Body: aws.String("payload"),
+		MessageAttributes: map[string]*sqs.MessageAttributeValue{
+			"_watermill_message_uuid": {
+				DataType:    aws.String(AWSStringDataType),
+				StringValue: aws.String("some-id"),
+			},
+			"name": {
+				DataType:    aws.String(AWSStringDataType),
+				StringValue: aws.String("some-event-name"),
+			},
+			"test": {
+				DataType:    aws.String(AWSStringDataType),
+				StringValue: aws.String("other-value"),
+			},
+			"nil datatype": {
+				DataType:    nil,
+				StringValue: nil,
+			},
+			"nil stringvalue": {
+				DataType:    aws.String(AWSStringDataType),
+				StringValue: nil,
+			},
+			"numeric": {
+				DataType:    aws.String(AWSNumberDataType),
+				StringValue: aws.String("1234"),
+			},
+			"binary": {
+				DataType:    aws.String(AWSBinaryDataType),
+				BinaryValue: []byte("binary-data"),
+			},
+		},
+	})
+	require.NoError(t, err)
+	assert.Equal(t, "some-id", msg.UUID)
+	assert.Equal(t, message.Payload("payload"), msg.Payload)
+	assert.Equal(t, message.Metadata{
+		"name":    "some-event-name",
+		"test":    "other-value",
+		"numeric": "1234",
+		"binary":  "binary-data",
+	}, msg.Metadata)
 }