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  }