Merge pull request #3 from asecurityteam/update-interface

upgrade interface version

compression.go

@@ -15,7 +15,7 @@ import (
 // runsqs.SQSProducer to produce the compressed message.
 // We must base64 encode the compressed message because SQS
 // supports only character data -- valid utf8 -- not raw binary data
-type CompressionSQSProducer struct {
+type CompressionSQSProducer struct { // nolint
 	Wrapped runsqs.SQSProducer
 }
 
@@ -41,3 +41,23 @@ func (producer *CompressionSQSProducer) ProduceMessage(ctx context.Context, mess
 	messageInput.MessageBody = aws.String(string(encodedBytes))
 	return producer.Wrapped.ProduceMessage(ctx, messageInput)
 }
+
+// BatchProduceMessage produces compressed, base64 encoded messages to the configured sqs queue.
+func (producer *CompressionSQSProducer) BatchProduceMessage(ctx context.Context, messageBatchInput *sqs.SendMessageBatchInput) (*sqs.SendMessageBatchOutput, error) {
+	for i := range messageBatchInput.Entries {
+		var bytes bytes.Buffer
+		gz := gzip.NewWriter(&bytes)
+		if _, err := gz.Write([]byte(*messageBatchInput.Entries[i].MessageBody)); err != nil {
+			return nil, err
+		}
+		if err := gz.Close(); err != nil {
+			return nil, err
+		}
+		encodedBytes := make([]byte, base64.StdEncoding.EncodedLen(len(bytes.Bytes())))
+		base64.StdEncoding.Encode(encodedBytes, bytes.Bytes())
+
+		messageBatchInput.QueueUrl = aws.String(producer.Wrapped.QueueURL())
+		messageBatchInput.Entries[i].MessageBody = aws.String(string(encodedBytes))
+	}
+	return producer.Wrapped.BatchProduceMessage(ctx, messageBatchInput)
+}

compression_test.go

@@ -10,7 +10,13 @@ import (
 	"github.com/golang/mock/gomock"
 )
 
-func TestCompression(t *testing.T) {
+const (
+	helloWorld               = `hello world!`
+	encodedCompressedMessage = `H4sIAAAAAAAA/8pIzcnJVyjPL8pJUQQEAAD//23CtAMMAAAA`
+	exampleURL               = `www.example.com`
+)
+
+func TestCompressionProduceMessage(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()
 
@@ -20,9 +26,9 @@ func TestCompression(t *testing.T) {
 	compressionProducer := CompressionSQSProducer{
 		Wrapped: mockMessageProducer,
 	}
-	queueURL := "www.example.com"
-	originalMessage := "hello world!"
-	encodedCompressedMessage := "H4sIAAAAAAAA/8pIzcnJVyjPL8pJUQQEAAD//23CtAMMAAAA"
+	queueURL := exampleURL
+	originalMessage := helloWorld
+	encodedCompressedMessage := encodedCompressedMessage
 
 	mockMessageProducer.EXPECT().QueueURL().Return(queueURL)
 	mockMessageProducer.EXPECT().ProduceMessage(mockContext, &sqs.SendMessageInput{
@@ -36,3 +42,35 @@ func TestCompression(t *testing.T) {
 	assert.Nil(t, e)
 
 }
+
+func TestCompressionBatchProduceMessage(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	mockContext := context.Background()
+	mockMessageProducer := NewMockSQSProducer(ctrl)
+
+	compressionProducer := CompressionSQSProducer{
+		Wrapped: mockMessageProducer,
+	}
+	queueURL := exampleURL
+	originalMessage := helloWorld
+	encodedCompressedMessage := encodedCompressedMessage
+
+	mockMessageProducer.EXPECT().QueueURL().Return(queueURL)
+	mockMessageProducer.EXPECT().BatchProduceMessage(mockContext, &sqs.SendMessageBatchInput{
+		QueueUrl: &queueURL,
+		Entries: []*sqs.SendMessageBatchRequestEntry{{
+			MessageBody: &encodedCompressedMessage,
+		}},
+	}).Return(nil, nil)
+
+	_, e := compressionProducer.BatchProduceMessage(mockContext, &sqs.SendMessageBatchInput{
+		QueueUrl: &queueURL,
+		Entries: []*sqs.SendMessageBatchRequestEntry{{
+			MessageBody: &originalMessage,
+		}},
+	})
+	assert.Nil(t, e)
+
+}

go.mod

@@ -3,7 +3,7 @@ module github.com/asecurityteam/runsqs-compressor
 go 1.13
 
 require (
-	github.com/asecurityteam/runsqs/v2 v2.0.0
+	github.com/asecurityteam/runsqs/v2 v2.5.0
 	github.com/aws/aws-sdk-go v1.43.43
 	github.com/golang/mock v1.6.0
 	github.com/rs/zerolog v1.26.1 // indirect

go.sum

@@ -1,7 +1,7 @@
 github.com/asecurityteam/logevent v1.6.1 h1:D/V11UxgZMBrktTgL3ugRO4hKdBWkz/xLBlpqDq4PoM=
 github.com/asecurityteam/logevent v1.6.1/go.mod h1:wtopE0Pe+072poxy+lGBMah604Uu6/BVwFWRwKwF3TA=
-github.com/asecurityteam/runsqs/v2 v2.0.0 h1:HhIGlmYS5KyjthZrFnntJsln3RM7uhvwXv04suQRCP8=
-github.com/asecurityteam/runsqs/v2 v2.0.0/go.mod h1:G11mzsloYaO7IJLy/1BiqFF4T+2v1jLZhijZSyPu4t8=
+github.com/asecurityteam/runsqs/v2 v2.5.0 h1:p/cG4nrTiIsH6cjLtKuqaCCs0Vh7ZCh6v35MaO1+XUs=
+github.com/asecurityteam/runsqs/v2 v2.5.0/go.mod h1:G11mzsloYaO7IJLy/1BiqFF4T+2v1jLZhijZSyPu4t8=
 github.com/aws/aws-sdk-go v1.35.0/go.mod h1:H7NKnBqNVzoTJpGfLrQkkD+ytBA93eiDYi/+8rV9s48=
 github.com/aws/aws-sdk-go v1.43.43 h1:1L06qzQvl4aC3Skfh5rV7xVhGHjIZoHcqy16NoyQ1o4=
 github.com/aws/aws-sdk-go v1.43.43/go.mod h1:y4AeaBuwd2Lk+GepC1E9v0qOiTws0MIWAX4oIKwKHZo=