added OptFns to sns/sqs and extra logs

ThreeDotsLabs · Oct 15, 2024 · 8cf4b95 · 8cf4b95
1 parent bcd0160
commit 8cf4b95
Show file tree

Hide file tree

Showing 11 changed files with 233 additions and 37 deletions.
diff --git a/internal/endpoint.go b/internal/endpoint.go
diff --git a/sns/config.go b/sns/config.go
@@ -14,6 +14,9 @@ import (
 type PublisherConfig struct {
 	AWSConfig aws.Config
 
+	// OptFns are options for the SNS client.
+	OptFns []func(*sns.Options)
+
 	CreateTopicConfig           ConfigAttributes
 	DoNotCreateTopicIfNotExists bool
 
@@ -64,6 +67,9 @@ func GenerateCreateTopicInputDefault(ctx context.Context, topic TopicName, attrs
 type SubscriberConfig struct {
 	AWSConfig aws.Config
 
+	// OptFns are options for the SNS client.
+	OptFns []func(*sns.Options)
+
 	TopicResolver TopicResolver
 
 	GenerateSqsQueueName GenerateSqsQueueNameFn

diff --git a/sns/publisher.go b/sns/publisher.go
@@ -2,13 +2,13 @@ package sns
 
 import (
 	"context"
+	"errors"
 	"fmt"
 
 	"github.com/ThreeDotsLabs/watermill"
 	"github.com/ThreeDotsLabs/watermill/message"
 	"github.com/aws/aws-sdk-go-v2/service/sns"
 	"github.com/aws/aws-sdk-go-v2/service/sns/types"
-	"github.com/pkg/errors"
 )
 
 type Publisher struct {
@@ -25,7 +25,7 @@ func NewPublisher(config PublisherConfig, logger watermill.LoggerAdapter) (*Publ
 	}
 
 	return &Publisher{
-		sns:    sns.NewFromConfig(config.AWSConfig),
+		sns:    sns.NewFromConfig(config.AWSConfig, config.OptFns...),
 		config: config,
 		logger: logger,
 	}, nil

diff --git a/sns/pubsub_test.go b/sns/pubsub_test.go
@@ -3,11 +3,15 @@ package sns_test
 import (
 	"context"
 	"fmt"
+	"net/url"
 	"testing"
 
-	"github.com/ThreeDotsLabs/watermill-amazonsqs/internal"
+	amazonsns "github.com/aws/aws-sdk-go-v2/service/sns"
+	amazonsqs "github.com/aws/aws-sdk-go-v2/service/sqs"
+
 	"github.com/ThreeDotsLabs/watermill-amazonsqs/sqs"
 	awsconfig "github.com/aws/aws-sdk-go-v2/config"
+	transport "github.com/aws/smithy-go/endpoints"
 	"github.com/stretchr/testify/assert"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
@@ -55,18 +59,27 @@ func TestPubSub_arn_topic_resolver(t *testing.T) {
 			return createPubSubWithConfig(
 				t,
 				sns.PublisherConfig{
-					AWSConfig:         cfg,
+					AWSConfig: cfg,
+					OptFns: []func(*amazonsns.Options){
+						GetEndpointResolverSns(),
+					},
 					CreateTopicConfig: sns.ConfigAttributes{},
 					Marshaler:         sns.DefaultMarshalerUnmarshaler{},
 					TopicResolver:     sns.TransparentTopicResolver{},
 				},
 				sns.SubscriberConfig{
-					AWSConfig:            cfg,
+					AWSConfig: cfg,
+					OptFns: []func(*amazonsns.Options){
+						GetEndpointResolverSns(),
+					},
 					GenerateSqsQueueName: sns.GenerateSqsQueueNameEqualToTopicName,
 					TopicResolver:        sns.TransparentTopicResolver{},
 				},
 				sqs.SubscriberConfig{
 					AWSConfig: cfg,
+					OptFns: []func(*amazonsqs.Options){
+						GetEndpointResolverSqs(),
+					},
 					QueueConfigAttributes: sqs.QueueConfigAttributes{
 						// Default value is 30 seconds - need to be lower for tests
 						VisibilityTimeout: "1",
@@ -112,18 +125,27 @@ func createPubSub(t *testing.T) (message.Publisher, message.Subscriber) {
 	return createPubSubWithConfig(
 		t,
 		sns.PublisherConfig{
-			AWSConfig:         cfg,
+			AWSConfig: cfg,
+			OptFns: []func(*amazonsns.Options){
+				GetEndpointResolverSns(),
+			},
 			CreateTopicConfig: sns.ConfigAttributes{},
 			TopicResolver:     topicResolver,
 			Marshaler:         sns.DefaultMarshalerUnmarshaler{},
 		},
 		sns.SubscriberConfig{
-			AWSConfig:            cfg,
+			AWSConfig: cfg,
+			OptFns: []func(*amazonsns.Options){
+				GetEndpointResolverSns(),
+			},
 			TopicResolver:        topicResolver,
 			GenerateSqsQueueName: sns.GenerateSqsQueueNameEqualToTopicName,
 		},
 		sqs.SubscriberConfig{
 			AWSConfig: cfg,
+			OptFns: []func(*amazonsqs.Options){
+				GetEndpointResolverSqs(),
+			},
 			QueueConfigAttributes: sqs.QueueConfigAttributes{
 				// Default value is 30 seconds - need to be lower for tests
 				VisibilityTimeout: "1",
@@ -141,20 +163,29 @@ func createPubSubWithConsumerGroup(t *testing.T, consumerGroup string) (message.
 	return createPubSubWithConfig(
 		t,
 		sns.PublisherConfig{
-			AWSConfig:         cfg,
+			AWSConfig: cfg,
+			OptFns: []func(*amazonsns.Options){
+				GetEndpointResolverSns(),
+			},
 			CreateTopicConfig: sns.ConfigAttributes{},
 			Marshaler:         sns.DefaultMarshalerUnmarshaler{},
 			TopicResolver:     topicResolver,
 		},
 		sns.SubscriberConfig{
 			AWSConfig: cfg,
+			OptFns: []func(*amazonsns.Options){
+				GetEndpointResolverSns(),
+			},
 			GenerateSqsQueueName: func(ctx context.Context, sqsTopic sns.TopicArn) (string, error) {
 				return consumerGroup, nil
 			},
 			TopicResolver: topicResolver,
 		},
 		sqs.SubscriberConfig{
 			AWSConfig: cfg,
+			OptFns: []func(*amazonsqs.Options){
+				GetEndpointResolverSqs(),
+			},
 			QueueConfigAttributes: sqs.QueueConfigAttributes{
 				// Default value is 30 seconds - need to be lower for tests
 				VisibilityTimeout: "1",
@@ -185,10 +216,25 @@ func GetAWSConfig(t *testing.T) aws.Config {
 
 	cfg, err := awsconfig.LoadDefaultConfig(
 		context.Background(),
-		internal.SetEndPoint("http://localhost:4566"),
 		awsconfig.WithRegion("us-west-2"),
 	)
 	require.NoError(t, err)
 
 	return cfg
 }
+
+func GetEndpointResolverSns() func(*amazonsns.Options) {
+	return amazonsns.WithEndpointResolverV2(sns.OverrideEndpointResolver{
+		Endpoint: transport.Endpoint{
+			URI: url.URL{Scheme: "http", Host: "localhost:4566"},
+		},
+	})
+}
+
+func GetEndpointResolverSqs() func(*amazonsqs.Options) {
+	return amazonsqs.WithEndpointResolverV2(sqs.OverrideEndpointResolver{
+		Endpoint: transport.Endpoint{
+			URI: url.URL{Scheme: "http", Host: "localhost:4566"},
+		},
+	})
+}
diff --git a/sns/resolver.go b/sns/resolver.go
@@ -0,0 +1,39 @@
+package sns
+
+import (
+	"context"
+
+	"github.com/aws/aws-sdk-go-v2/service/sns"
+	"github.com/aws/smithy-go/endpoints"
+)
+
+// OverrideEndpointResolver is a custom endpoint resolver that always returns the same endpoint.
+// It can be used to use AWS emulators like Localstack.
+//
+// For example:
+// import (
+//
+//	"github.com/ThreeDotsLabs/watermill-amazonsns/sns"
+//	amazonsns "github.com/aws/aws-sdk-go-v2/service/sns"
+//	"github.com/aws/smithy-go/transport"
+//
+// )
+//
+//	pub, err := sns.NewPublisher(sns.PublisherConfig{
+//			AWSConfig: cfg,
+//			Marshaler: sns.DefaultMarshalerUnmarshaler{},
+//			OptFns: []func(*amazonsns.Options){
+//				amazonsns.WithEndpointResolverV2(sns.OverrideEndpointResolver{
+//					Endpoint: transport.Endpoint{
+//						URI: url.URL{Scheme: "http", Host: "localstack:4566"},
+//					},
+//				}),
+//			},
+//		}, logger)
+type OverrideEndpointResolver struct {
+	Endpoint transport.Endpoint
+}
+
+func (o OverrideEndpointResolver) ResolveEndpoint(ctx context.Context, params sns.EndpointParameters) (transport.Endpoint, error) {
+	return o.Endpoint, nil
+}
diff --git a/sns/subscriber.go b/sns/subscriber.go
@@ -48,8 +48,8 @@ func NewSubscriber(
 	return &Subscriber{
 		config:    config,
 		logger:    logger,
-		sns:       sns.NewFromConfig(config.AWSConfig),
-		sqsClient: awsSqs.NewFromConfig(config.AWSConfig),
+		sns:       sns.NewFromConfig(config.AWSConfig, config.OptFns...),
+		sqsClient: awsSqs.NewFromConfig(sqsConfig.AWSConfig, sqsConfig.OptFns...),
 		sqs:       sqs,
 	}, nil
 }
@@ -60,6 +60,12 @@ func (s *Subscriber) Subscribe(ctx context.Context, topic string) (<-chan *messa
 		return nil, err
 	}
 
+	logger := s.logger.With(watermill.LogFields{
+		"topic":     topic,
+		"topic_arn": snsTopicArn,
+	})
+	logger.Debug("Subscribing to topic", nil)
+
 	if !s.config.DoNotCreateSqsSubscription {
 		if err := s.SubscribeInitializeWithContext(ctx, topic); err != nil {
 			return nil, err
@@ -71,6 +77,8 @@ func (s *Subscriber) Subscribe(ctx context.Context, topic string) (<-chan *messa
 		return nil, fmt.Errorf("failed to generate SQS queue name: %w", err)
 	}
 
+	logger.Debug("Subscribing to SQS", watermill.LogFields{"sqs_topic": sqsTopic})
+
 	return s.sqs.Subscribe(ctx, sqsTopic)
 }
 
@@ -79,11 +87,18 @@ func (s *Subscriber) SubscribeInitialize(topic string) error {
 }
 
 func (s *Subscriber) SubscribeInitializeWithContext(ctx context.Context, topic string) error {
+	logger := s.logger.With(watermill.LogFields{
+		"topic": topic,
+	})
+	logger.Debug("Initializing SNS subscription", nil)
+
 	snsTopicArn, err := s.config.TopicResolver.ResolveTopic(ctx, topic)
 	if err != nil {
 		return err
 	}
 
+	logger.Debug("Resolved topic", watermill.LogFields{"topic_arn": snsTopicArn})
+
 	sqsTopic, err := s.config.GenerateSqsQueueName(ctx, snsTopicArn)
 	if err != nil {
 		return fmt.Errorf("failed to generate SQS queue name: %w", err)

diff --git a/sqs/config.go b/sqs/config.go
@@ -15,6 +15,9 @@ import (
 type SubscriberConfig struct {
 	AWSConfig aws.Config
 
+	// OptFns are options for the SQS client.
+	OptFns []func(*sqs.Options)
+
 	DoNotCreateQueueIfNotExists bool
 
 	QueueUrlResolver QueueUrlResolver
@@ -83,6 +86,9 @@ func (c SubscriberConfig) Validate() error {
 type PublisherConfig struct {
 	AWSConfig aws.Config
 
+	// OptFns are options for the SQS client.
+	OptFns []func(*sqs.Options)
+
 	CreateQueueConfig QueueConfigAttributes
 	DoNotCacheQueues  bool
 

diff --git a/sqs/publisher.go b/sqs/publisher.go
@@ -2,13 +2,13 @@ package sqs
 
 import (
 	"context"
+	"errors"
 	"fmt"
 
 	"github.com/ThreeDotsLabs/watermill"
 	"github.com/ThreeDotsLabs/watermill/message"
 	"github.com/aws/aws-sdk-go-v2/service/sqs"
 	"github.com/aws/aws-sdk-go-v2/service/sqs/types"
-	"github.com/pkg/errors"
 )
 
 type Publisher struct {
@@ -25,7 +25,7 @@ func NewPublisher(config PublisherConfig, logger watermill.LoggerAdapter) (*Publ
 	}
 
 	return &Publisher{
-		sqs:    sqs.NewFromConfig(config.AWSConfig),
+		sqs:    sqs.NewFromConfig(config.AWSConfig, config.OptFns...),
 		config: config,
 		logger: logger,
 	}, nil