add sqs attributes to message

wish · Oct 29, 2021 · a4f921d · a4f921d
1 parent ac81e2a
commit a4f921d
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 3 deletions.
diff --git a/backends/sqs/backend.go b/backends/sqs/backend.go
@@ -371,7 +371,13 @@ func (s *Backend) GetMessages(ctx context.Context, in *rpc.GetMessagesRequest) (
 		return nil, err
 	}
 
+	var sqsAttributeNames []*string
+	if in.RequireQueueSystemAttributes {
+		sqsAttributeNames = []*string{aws.String("All")}
+	}
+
 	output, err := s.sqs.ReceiveMessageWithContext(ctx, &sqs.ReceiveMessageInput{
+		AttributeNames:        sqsAttributeNames,
 		MessageAttributeNames: []*string{aws.String("All")},
 		QueueUrl:              &url,
 		WaitTimeSeconds:       &in.LongPollSeconds,
@@ -389,9 +395,17 @@ func (s *Backend) GetMessages(ctx context.Context, in *rpc.GetMessagesRequest) (
 		for key, valueStruct := range message.MessageAttributes {
 			attributes[key] = *valueStruct.StringValue
 		}
+
+		sqsAttributes := make(map[string]string)
+		if in.RequireQueueSystemAttributes {
+			for key, value := range message.Attributes {
+				sqsAttributes[key] = *value
+			}
+		}
 		messages = append(messages, &rpc.Message{
-			Data:       *message.Body,
-			Attributes: attributes,
+			Data:                  *message.Body,
+			Attributes:            attributes,
+			QueueSystemAttributes: sqsAttributes,
 			Receipt: &rpc.MessageReceipt{
 				Id: *message.ReceiptHandle,
 			},

diff --git a/rpc/qproxy.proto b/rpc/qproxy.proto
@@ -148,7 +148,8 @@ message FailedPublish {
 
 message Message {
   string Data = 1;
-  map<string, string> Attributes = 2;  
+  map<string, string> Attributes = 2;
+  map<string, string> QueueSystemAttributes = 3;
 
   MessageReceipt Receipt = 100;
 }
@@ -170,6 +171,7 @@ message GetMessagesRequest {
   int64 LongPollSeconds = 2;
   int64 MaxMessages = 3;
   int64 AckDeadlineSeconds = 4;
+  bool RequireQueueSystemAttributes = 5;
   int64 RPCTimeout = 100;
 }