Repository: beatlabs/patron
Pull request #468 (2022-02-18 15:33:13)

Kafka consumer option: new session callback
## Which problem is this PR solving?

https://github.com/beatlabs/patron/issues/467

## Short description of the changes

Add a new option that set a specific callback.



Patch:
diff --git a/component/kafka/group/component.go b/component/kafka/group/component.go
index 0f6f7a3e3..2c8fa5821 100644
--- a/component/kafka/group/component.go
+++ b/component/kafka/group/component.go
@@ -169,6 +169,7 @@ type Component struct {
 	retries                   uint
 	retryWait                 time.Duration
 	commitSync                bool
+	sessionCallback           func(sarama.ConsumerGroupSession) error
 }
 
 // Run starts the consumer processing loop to process messages from Kafka.
@@ -184,7 +185,7 @@ func (c *Component) processing(ctx context.Context) error {
 	retries := int(c.retries)
 	for i := 0; i <= retries; i++ {
 		handler := newConsumerHandler(ctx, c.name, c.group, c.proc, c.failStrategy, c.batchSize,
-			c.batchTimeout, c.commitSync, c.batchMessageDeduplication)
+			c.batchTimeout, c.commitSync, c.batchMessageDeduplication, c.sessionCallback)
 
 		client, err := sarama.NewConsumerGroup(c.brokers, c.group, c.saramaConfig)
 		componentError = err
@@ -284,10 +285,12 @@ type consumerHandler struct {
 
 	// whether the handler has processed any messages
 	processedMessages bool
+	sessionCallback   func(sarama.ConsumerGroupSession) error
 }
 
 func newConsumerHandler(ctx context.Context, name, group string, processorFunc kafka.BatchProcessorFunc,
-	fs kafka.FailStrategy, batchSize uint, batchTimeout time.Duration, commitSync bool, batchMessageDeduplication bool) *consumerHandler {
+	fs kafka.FailStrategy, batchSize uint, batchTimeout time.Duration, commitSync, batchMessageDeduplication bool,
+	sessionCallback func(sarama.ConsumerGroupSession) error) *consumerHandler {
 	return &consumerHandler{
 		ctx:                       ctx,
 		name:                      name,
@@ -300,12 +303,16 @@ func newConsumerHandler(ctx context.Context, name, group string, processorFunc k
 		proc:                      processorFunc,
 		failStrategy:              fs,
 		commitSync:                commitSync,
+		sessionCallback:           sessionCallback,
 	}
 }
 
 // Setup is run at the beginning of a new session, before ConsumeClaim
-func (c *consumerHandler) Setup(sarama.ConsumerGroupSession) error {
-	return nil
+func (c *consumerHandler) Setup(cgs sarama.ConsumerGroupSession) error {
+	if c.sessionCallback == nil {
+		return nil
+	}
+	return c.sessionCallback(cgs)
 }
 
 // Cleanup is run at the end of a session, once all ConsumeClaim goroutines have exited
diff --git a/component/kafka/group/option.go b/component/kafka/group/option.go
index 5007c07e1..19eca8b87 100644
--- a/component/kafka/group/option.go
+++ b/component/kafka/group/option.go
@@ -124,3 +124,15 @@ func CommitSync() OptionFunc {
 		return nil
 	}
 }
+
+// NewSessionCallback adds a callback when a new consumer group session is created (e.g., rebalancing).
+func NewSessionCallback(sessionCallback func(sarama.ConsumerGroupSession) error) OptionFunc {
+	return func(c *Component) error {
+		if sessionCallback == nil {
+			return errors.New("nil session callback")
+		}
+
+		c.sessionCallback = sessionCallback
+		return nil
+	}
+}


Test patch:
diff --git a/component/kafka/group/component_test.go b/component/kafka/group/component_test.go
index 1b74b2826..0912c4fb9 100644
--- a/component/kafka/group/component_test.go
+++ b/component/kafka/group/component_test.go
@@ -310,7 +310,7 @@ func TestHandler_ConsumeClaim(t *testing.T) {
 			ctx, cancel := context.WithCancel(context.Background())
 			defer cancel()
 			h := newConsumerHandler(ctx, tt.name, "grp", tt.proc.Process, tt.failStrategy, tt.batchSize,
-				10*time.Millisecond, true, tt.batchMessageDeduplication)
+				10*time.Millisecond, true, tt.batchMessageDeduplication, nil)
 
 			ch := make(chan *sarama.ConsumerMessage, len(tt.msgs))
 			for _, m := range tt.msgs {
diff --git a/component/kafka/group/option_test.go b/component/kafka/group/option_test.go
index a4d332c44..2c342127d 100644
--- a/component/kafka/group/option_test.go
+++ b/component/kafka/group/option_test.go
@@ -4,6 +4,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/Shopify/sarama"
 	"github.com/beatlabs/patron/component/kafka"
 	"github.com/stretchr/testify/assert"
 )
@@ -150,6 +151,42 @@ func TestBatchTimeout(t *testing.T) {
 	}
 }
 
+func TestNewSessionCallback(t *testing.T) {
+	t.Parallel()
+	type args struct {
+		sessionCallback func(sarama.ConsumerGroupSession) error
+	}
+	tests := map[string]struct {
+		args        args
+		expectedErr string
+	}{
+		"success": {
+			args: args{sessionCallback: func(session sarama.ConsumerGroupSession) error {
+				return nil
+			}},
+		},
+		"nil session callback": {
+			args:        args{},
+			expectedErr: "nil session callback",
+		},
+	}
+	for name, tt := range tests {
+		tt := tt
+		t.Run(name, func(t *testing.T) {
+			t.Parallel()
+			c := &Component{}
+			err := NewSessionCallback(tt.args.sessionCallback)(c)
+			if tt.expectedErr != "" {
+				assert.EqualError(t, err, tt.expectedErr)
+				assert.Nil(t, c.sessionCallback)
+			} else {
+				assert.NoError(t, err)
+				assert.NotNil(t, c.sessionCallback)
+			}
+		})
+	}
+}
+
 func TestBatchMessageDeduplication(t *testing.T) {
 	c := &Component{}
 	err := BatchMessageDeduplication()(c)