Merge pull request #214 from toninho09/add-header-on-message

Adding header on message
lovoo · Jan 6, 2020 · 1cf2c41 · 1cf2c41
2 parents 2554930 + 6c253cf
commit 1cf2c41
Show file tree

Hide file tree

Showing 9 changed files with 50 additions and 4 deletions.
diff --git a/context.go b/context.go
@@ -33,6 +33,9 @@ type Context interface {
 	// Value returns the value of the key in the group table.
 	Value() interface{}
 
+	// Headers returns the headers of the input message
+	Headers() map[string][]byte
+
 	// SetValue updates the value of the key in the group table.
 	SetValue(value interface{})
 
@@ -196,6 +199,10 @@ func (ctx *cbContext) Partition() int32 {
 	return ctx.msg.Partition
 }
 
+func (ctx *cbContext) Headers() map[string][]byte {
+	return ctx.msg.Header
+}
+
 func (ctx *cbContext) Join(topic Table) interface{} {
 	if ctx.pviews == nil {
 		ctx.Fail(fmt.Errorf("table %s not subscribed", topic))

diff --git a/kafka/confluent/confluent.go b/kafka/confluent/confluent.go
@@ -161,13 +161,19 @@ func (c *confluent) run() {
 					partition = e.TopicPartition.Partition
 				)
 
+				headers := make(map[string][]byte)
+				for _, header := range e.Headers {
+					headers[header.Key] = header.Value
+				}
+
 				c.events <- &kafka.Message{
 					Topic:     topic,
 					Partition: partition,
 					Offset:    int64(e.TopicPartition.Offset),
 					Key:       string(e.Key),
 					Value:     e.Value,
 					Timestamp: e.Timestamp,
+					Header:    headers,
 				}
 
 			case rdkafka.PartitionEOF:

diff --git a/kafka/event.go b/kafka/event.go
@@ -48,6 +48,7 @@ type Message struct {
 	Partition int32
 	Offset    int64
 	Timestamp time.Time
+	Header    map[string][]byte
 
 	Key   string
 	Value []byte

diff --git a/kafka/group_consumer.go b/kafka/group_consumer.go
@@ -211,6 +211,12 @@ func (c *groupConsumer) waitForMessages() bool {
 				return false
 			}
 		case msg := <-c.consumer.Messages():
+
+			headers := make(map[string][]byte)
+			for _, header := range msg.Headers {
+				headers[string(header.Key)] = header.Value
+			}
+
 			select {
 			case c.events <- &Message{
 				Topic:     msg.Topic,
@@ -219,6 +225,7 @@ func (c *groupConsumer) waitForMessages() bool {
 				Timestamp: msg.Timestamp,
 				Key:       string(msg.Key),
 				Value:     msg.Value,
+				Header:    headers,
 			}:
 			case <-c.stop:
 				return false

diff --git a/kafka/group_consumer_test.go b/kafka/group_consumer_test.go
@@ -255,6 +255,7 @@ func TestGroupConsumer_GroupConsumeMessages(t *testing.T) {
 		Offset:    0,
 		Key:       string(key),
 		Value:     value,
+		Header:    make(map[string][]byte),
 	})
 
 	ensure.DeepEqual(t, c.partitionMap, map[int32]bool{
@@ -458,6 +459,7 @@ func TestGroupConsumer_GroupNotificationsAfterMessages(t *testing.T) {
 		Offset:    0,
 		Key:       string(key),
 		Value:     value,
+		Header:    make(map[string][]byte),
 	})
 
 	ensure.DeepEqual(t, c.partitionMap, map[int32]bool{
@@ -564,6 +566,7 @@ func TestGroupConsumer_GroupEmptyNotifications(t *testing.T) {
 		Offset:    0,
 		Key:       string(key),
 		Value:     value,
+		Header:    make(map[string][]byte),
 	})
 
 	ensure.DeepEqual(t, c.partitionMap, map[int32]bool{

diff --git a/kafka/simple_consumer.go b/kafka/simple_consumer.go
@@ -139,6 +139,12 @@ func (c *simpleConsumer) run(pc sarama.PartitionConsumer, topic string, partitio
 				// drained.
 				continue
 			}
+
+			headers := make(map[string][]byte)
+			for _, header := range m.Headers {
+				headers[string(header.Key)] = header.Value
+			}
+
 			select {
 			case c.events <- &Message{
 				Topic:     m.Topic,
@@ -147,6 +153,7 @@ func (c *simpleConsumer) run(pc sarama.PartitionConsumer, topic string, partitio
 				Key:       string(m.Key),
 				Value:     m.Value,
 				Timestamp: m.Timestamp,
+				Header:    headers,
 			}:
 			case <-c.dying:
 				return

diff --git a/kafka/simple_consumer_test.go b/kafka/simple_consumer_test.go
@@ -276,7 +276,23 @@ func TestSimpleConsumer_ErrorBlocked(t *testing.T) {
 	pc.EXPECT().HighWaterMarkOffset().Return(int64(0))
 	mo, ok := (<-ch).(*Message)
 	ensure.True(t, ok)
-	ensure.DeepEqual(t, mo, &Message{Topic: "sometopic", Partition: 123, Key: "somekey", Value: []byte("somevalue")})
+	ensure.DeepEqual(t, mo, &Message{Topic: "sometopic", Partition: 123, Key: "somekey", Value: []byte("somevalue"), Header: map[string][]byte{}})
+
+	messages <- &sarama.ConsumerMessage{
+		Key:       []byte("somekey"),
+		Value:     []byte("somevalue"),
+		Topic:     "sometopic",
+		Partition: 123,
+		Headers: []*sarama.RecordHeader{{
+			Key:   []byte("someHeader"),
+			Value: []byte("someHeaderValue"),
+		}},
+	}
+
+	pc.EXPECT().HighWaterMarkOffset().Return(int64(0))
+	mo1, ok := (<-ch).(*Message)
+	ensure.True(t, ok)
+	ensure.DeepEqual(t, mo1, &Message{Topic: "sometopic", Partition: 123, Key: "somekey", Value: []byte("somevalue"), Header: map[string][]byte{"someHeader": []byte("someHeaderValue")}})
 
 	// we now write, but don't read events
 	messages <- &sarama.ConsumerMessage{

diff --git a/partition.go b/partition.go
@@ -134,6 +134,7 @@ func newMessage(ev *kafka.Message) *message {
 		Timestamp: ev.Timestamp,
 		Data:      ev.Value,
 		Key:       ev.Key,
+		Header:    ev.Header,
 	}
 }
 

diff --git a/processor.go b/processor.go
@@ -47,6 +47,7 @@ type message struct {
 	Partition int32
 	Offset    int64
 	Timestamp time.Time
+	Header    map[string][]byte
 }
 
 // ProcessCallback function is called for every message received by the
@@ -655,15 +656,12 @@ func (g *Processor) rebalance(errg *multierr.ErrGroup, ctx context.Context, part
 	errs := new(multierr.Errors)
 	g.opts.log.Printf("Processor: rebalancing: %+v", partitions)
 
-
 	// callback the new partition assignment
 	g.opts.rebalanceCallback(partitions)
 
-
 	g.m.Lock()
 	defer g.m.Unlock()
 
-
 	for id := range partitions {
 		// create partition views
 		if err := g.createPartitionViews(errg, ctx, id); err != nil {