Allow to set SNS Message Attributes via context

pubsub/aws/aws.go

@@ -19,6 +19,15 @@ import (
 	"golang.org/x/net/context"
 )
 
+// The key type is unexported to prevent collisions with context keys defined in
+// other packages.
+type key int
+
+// msgAttrsKey is the context key for the SNS Message Attributes.  Its value of zero is
+// arbitrary. If this package defined other context keys, they would have
+// different integer values.
+const msgAttrsKey key = 0
+
 // publisher will accept AWS credentials and an SNS topic name
 // and it will emit any publish events to it.
 type publisher struct {
@@ -79,17 +88,28 @@ func (p *publisher) Publish(ctx context.Context, key string, m proto.Message) er
 
 // PublishRaw will emit the byte array to the SNS topic.
 // The key will be used as the SNS message subject.
-func (p *publisher) PublishRaw(_ context.Context, key string, m []byte) error {
+// You can use func WithMessageAttributes to set SNS message attributes for the message
+func (p *publisher) PublishRaw(ctx context.Context, key string, m []byte) error {
 	msg := &sns.PublishInput{
 		TopicArn: &p.topic,
 		Subject:  &key,
 		Message:  aws.String(base64.StdEncoding.EncodeToString(m)),
 	}
 
+	if v, ok := ctx.Value(msgAttrsKey).(map[string]*sns.MessageAttributeValue); ok {
+		msg.MessageAttributes = v
+	}
+
 	_, err := p.sns.Publish(msg)
 	return err
 }
 
+// WithMessaggeAttributes used to add SNS Message Attributes to the context
+// for further usage in publishing messages to sns with provided attributes
+func WithMessaggeAttributes(ctx context.Context, msgAttrs map[string]*sns.MessageAttributeValue) context.Context {
+	return context.WithValue(ctx, msgAttrsKey, msgAttrs)
+}
+
 var (
 	// defaultSQSMaxMessages is default the number of bulk messages
 	// the subscriber will attempt to fetch on each