Add CustomFirestoreClient to subscriber/publisher config

czeslavo · May 14, 2020 · 78d4cb0 · 78d4cb0
1 parent 04be856
commit 78d4cb0
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 7 deletions.
diff --git a/pkg/firestore/publisher.go b/pkg/firestore/publisher.go
@@ -45,6 +45,9 @@ type PublisherConfig struct {
 
 	// Marshaler marshals message from Watermill to Firestore format and vice versa.
 	Marshaler Marshaler
+
+	// CustomFirestoreClient can be used to override a default client.
+	CustomFirestoreClient *firestore.Client
 }
 
 func (c *PublisherConfig) setDefaults() {
@@ -80,9 +83,15 @@ type subscriptionsCacheEntry struct {
 func NewPublisher(config PublisherConfig, logger watermill.LoggerAdapter) (*Publisher, error) {
 	config.setDefaults()
 
-	client, err := firestore.NewClient(context.Background(), config.ProjectID, config.GoogleClientOpts...)
-	if err != nil {
-		return nil, err
+	var client *firestore.Client
+	if config.CustomFirestoreClient != nil {
+		client = config.CustomFirestoreClient
+	} else {
+		var err error
+		client, err = firestore.NewClient(context.Background(), config.ProjectID, config.GoogleClientOpts...)
+		if err != nil {
+			return nil, errors.Wrap(err, "cannot create default firestore client")
+		}
 	}
 
 	return &Publisher{

diff --git a/pkg/firestore/subscriber.go b/pkg/firestore/subscriber.go
@@ -2,13 +2,13 @@ package firestore
 
 import (
 	"context"
-	"errors"
 	"sync"
 	"time"
 
 	"cloud.google.com/go/firestore"
 	"github.com/ThreeDotsLabs/watermill"
 	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/pkg/errors"
 	"google.golang.org/api/option"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -45,6 +45,9 @@ type SubscriberConfig struct {
 
 	// Marshaler marshals message from Watermill to Firestore format and vice versa.
 	Marshaler Marshaler
+
+	// CustomFirestoreClient can be used to override a default client.
+	CustomFirestoreClient *firestore.Client
 }
 
 func (c *SubscriberConfig) setDefaults() {
@@ -84,9 +87,15 @@ type Subscriber struct {
 func NewSubscriber(config SubscriberConfig, logger watermill.LoggerAdapter) (*Subscriber, error) {
 	config.setDefaults()
 
-	client, err := firestore.NewClient(context.Background(), config.ProjectID, config.GoogleClientOpts...)
-	if err != nil {
-		return nil, err
+	var client *firestore.Client
+	if config.CustomFirestoreClient != nil {
+		client = config.CustomFirestoreClient
+	} else {
+		var err error
+		client, err = firestore.NewClient(context.Background(), config.ProjectID, config.GoogleClientOpts...)
+		if err != nil {
+			return nil, errors.Wrap(err, "cannot create default firestore client")
+		}
 	}
 
 	return &Subscriber{