github.com/wormhole-foundation/wormhole-explorer/common@v0.0.0-20240604151348-09585b5b97c5/client/sqs/consumer.go (about) 1 package sqs 2 3 import ( 4 "context" 5 "time" 6 7 "github.com/aws/aws-sdk-go-v2/aws" 8 aws_sqs "github.com/aws/aws-sdk-go-v2/service/sqs" 9 aws_sqs_types "github.com/aws/aws-sdk-go-v2/service/sqs/types" 10 ) 11 12 // ConsumerOption represents a consumer option function. 13 type ConsumerOption func(*Consumer) 14 15 // Consumer represents SQS consumer. 16 type Consumer struct { 17 api *aws_sqs.Client 18 url string 19 maxMessages int32 20 visibilityTimeout int32 21 waitTimeSeconds int32 22 } 23 24 // New instances of a Consumer to consume SQS messages. 25 func NewConsumer(awsConfig aws.Config, url string, opts ...ConsumerOption) (*Consumer, error) { 26 consumer := &Consumer{ 27 api: aws_sqs.NewFromConfig(awsConfig), 28 url: url, 29 maxMessages: 10, 30 visibilityTimeout: 60, 31 waitTimeSeconds: 20, 32 } 33 34 for _, opt := range opts { 35 opt(consumer) 36 } 37 38 return consumer, nil 39 } 40 41 // WithMaxMessages allows to specify an maximum number of messages to return when setting a value. 42 func WithMaxMessages(v int32) ConsumerOption { 43 return func(c *Consumer) { 44 c.maxMessages = v 45 } 46 } 47 48 // WithVisibilityTimeout allows to specify a visibility timeout when setting a value. 49 func WithVisibilityTimeout(v int32) ConsumerOption { 50 return func(c *Consumer) { 51 c.visibilityTimeout = v 52 } 53 } 54 55 // WithWaitTimeSeconds allows to specify a wait time when setting a value. 56 func WithWaitTimeSeconds(v int32) ConsumerOption { 57 return func(c *Consumer) { 58 c.waitTimeSeconds = v 59 } 60 } 61 62 // GetMessages retrieves messages from SQS. 63 func (c *Consumer) GetMessages(ctx context.Context) ([]aws_sqs_types.Message, error) { 64 params := &aws_sqs.ReceiveMessageInput{ 65 QueueUrl: aws.String(c.url), 66 MaxNumberOfMessages: c.maxMessages, 67 AttributeNames: []aws_sqs_types.QueueAttributeName{ 68 aws_sqs_types.QueueAttributeNameAll, 69 }, 70 MessageAttributeNames: []string{ 71 string(aws_sqs_types.QueueAttributeNameAll), 72 }, 73 WaitTimeSeconds: c.waitTimeSeconds, 74 VisibilityTimeout: c.visibilityTimeout, 75 } 76 77 res, err := c.api.ReceiveMessage(ctx, params) 78 if err != nil { 79 return nil, err 80 } 81 82 return res.Messages, nil 83 } 84 85 // DeleteMessage deletes messages from SQS. 86 func (c *Consumer) DeleteMessage(ctx context.Context, id *string) error { 87 params := &aws_sqs.DeleteMessageInput{ 88 QueueUrl: aws.String(c.url), 89 ReceiptHandle: id, 90 } 91 _, err := c.api.DeleteMessage(ctx, params) 92 93 return err 94 } 95 96 // GetVisibilityTimeout returns visibility timeout. 97 func (c *Consumer) GetVisibilityTimeout() time.Duration { 98 return time.Duration(int64(c.visibilityTimeout) * int64(time.Second)) 99 } 100 101 // GetQueueAttributes get queue attributes. 102 func (c *Consumer) GetQueueAttributes(ctx context.Context) (*aws_sqs.GetQueueAttributesOutput, error) { 103 params := &aws_sqs.GetQueueAttributesInput{ 104 QueueUrl: aws.String(c.url), 105 AttributeNames: []aws_sqs_types.QueueAttributeName{ 106 aws_sqs_types.QueueAttributeNameCreatedTimestamp, 107 }, 108 } 109 return c.api.GetQueueAttributes(ctx, params) 110 } 111 112 // GetQueueUrl returns queue url. 113 func (c *Consumer) GetQueueUrl() string { 114 return c.url 115 }