Merge pull request #5 from czeslavo/transactional-publisher

Add transactional publisher

pkg/firestore/publisher.go

@@ -2,10 +2,11 @@ package firestore
 
 import (
 	"context"
-	"github.com/pkg/errors"
 	"sync"
 	"time"
 
+	"github.com/pkg/errors"
+
 	"cloud.google.com/go/firestore"
 	"github.com/ThreeDotsLabs/watermill"
 	"github.com/ThreeDotsLabs/watermill/message"
@@ -164,7 +165,6 @@ func (p *Publisher) PublishInTransaction(topic string, t *firestore.Transaction,
 			})
 
 			logger.Trace("Added message to transaction", nil)
-			continue
 		}
 	}
 
@@ -277,7 +277,6 @@ func (p *Publisher) publishInBatches(
 			logger.Error("Failed to commit messages batch", err, nil)
 			return err
 		}
-		logger.Trace("Published message", nil)
 	}
 
 	return nil

pkg/firestore/transactional.go

@@ -0,0 +1,25 @@
+package firestore
+
+import (
+	"cloud.google.com/go/firestore"
+	"github.com/ThreeDotsLabs/watermill"
+	"github.com/ThreeDotsLabs/watermill/message"
+)
+
+type TransactionalPublisher struct {
+	publisher *Publisher
+	tx        *firestore.Transaction
+}
+
+func NewTransactionalPublisher(config PublisherConfig, tx *firestore.Transaction, logger watermill.LoggerAdapter) (*TransactionalPublisher, error) {
+	publisher, err := NewPublisher(config, logger)
+	if err != nil {
+		return nil, err
+	}
+
+	return &TransactionalPublisher{publisher, tx}, nil
+}
+
+func (p *TransactionalPublisher) Publish(topic string, messages ...*message.Message) error {
+	return p.publisher.PublishInTransaction(topic, p.tx, messages...)
+}

pkg/firestore/transactional_test.go

@@ -0,0 +1,56 @@
+package firestore_test
+
+import (
+	"context"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+
+	stdFirestore "cloud.google.com/go/firestore"
+	"github.com/ThreeDotsLabs/watermill"
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/czeslavo/watermill-firestore/pkg/firestore"
+	"github.com/stretchr/testify/require"
+)
+
+func TestTransactionalPublisher(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
+	defer cancel()
+
+	projectID := os.Getenv("FIRESTORE_PROJECT_ID")
+	client, err := stdFirestore.NewClient(ctx, projectID)
+	require.NoError(t, err)
+
+	logger := watermill.NewStdLogger(true, true)
+	topic := "transactional_publisher_test_" + watermill.NewShortUUID()
+
+	subscriber, err := firestore.NewSubscriber(firestore.SubscriberConfig{
+		ProjectID: projectID,
+	}, logger)
+	msgs, err := subscriber.Subscribe(ctx, topic)
+	require.NoError(t, err)
+
+	msg := message.NewMessage(watermill.NewUUID(), message.Payload("payload"))
+	msg.Metadata = message.Metadata{"key": "value"}
+
+	err = client.RunTransaction(ctx, func(ctx context.Context, tx *stdFirestore.Transaction) error {
+		publisher, err := firestore.NewTransactionalPublisher(firestore.PublisherConfig{
+			ProjectID: projectID,
+		}, tx, logger)
+		require.NoError(t, err)
+
+		err = publisher.Publish(topic, msg)
+		require.NoError(t, err)
+
+		return nil
+	})
+	require.NoError(t, err)
+
+	receivedMsg := <-msgs
+	require.NotNil(t, receivedMsg)
+	assert.Equal(t, msg.UUID, receivedMsg.UUID)
+	assert.Equal(t, msg.Payload, receivedMsg.Payload)
+	assert.Equal(t, msg.Metadata, receivedMsg.Metadata)
+}