Kafka: Add message key and other metadata as metadata in consumer (dapr#3289)

Signed-off-by: Patrick Assuied <[email protected]>

Author: passuied

common/component/kafka/consumer.go

@@ -141,12 +141,7 @@ func (consumer *consumer) doBulkCallback(session sarama.ConsumerGroupSession,
 
 	for i, message := range messages {
 		if message != nil {
-			metadata := make(map[string]string, len(message.Headers))
-			if message.Headers != nil {
-				for _, t := range message.Headers {
-					metadata[string(t.Key)] = string(t.Value)
-				}
-			}
+			metadata := GetEventMetadata(message)
 			childMessage := KafkaBulkMessageEntry{
 				EntryId:  strconv.Itoa(i),
 				Event:    message.Value,
@@ -193,20 +188,33 @@ func (consumer *consumer) doCallback(session sarama.ConsumerGroupSession, messag
 		Topic: message.Topic,
 		Data:  message.Value,
 	}
-	// This is true only when headers are set (Kafka > 0.11)
-	if len(message.Headers) > 0 {
-		event.Metadata = make(map[string]string, len(message.Headers))
-		for _, header := range message.Headers {
-			event.Metadata[string(header.Key)] = string(header.Value)
-		}
-	}
+	event.Metadata = GetEventMetadata(message)
+
 	err = handlerConfig.Handler(session.Context(), &event)
 	if err == nil {
 		session.MarkMessage(message, "")
 	}
 	return err
 }
 
+func GetEventMetadata(message *sarama.ConsumerMessage) map[string]string {
+	if message != nil {
+		metadata := make(map[string]string, len(message.Headers)+5)
+		if message.Key != nil {
+			metadata[keyMetadataKey] = string(message.Key)
+		}
+		metadata[offsetMetadataKey] = strconv.FormatInt(message.Offset, 10)
+		metadata[topicMetadataKey] = message.Topic
+		metadata[timestampMetadataKey] = strconv.FormatInt(message.Timestamp.UnixMilli(), 10)
+		metadata[partitionMetadataKey] = strconv.FormatInt(int64(message.Partition), 10)
+		for _, header := range message.Headers {
+			metadata[string(header.Key)] = string(header.Value)
+		}
+		return metadata
+	}
+	return nil
+}
+
 func (consumer *consumer) Cleanup(sarama.ConsumerGroupSession) error {
 	return nil
 }

common/component/kafka/metadata.go

@@ -28,6 +28,11 @@ import (
 
 const (
 	key                  = "partitionKey"
+	keyMetadataKey       = "__key"
+	timestampMetadataKey = "__timestamp"
+	offsetMetadataKey    = "__offset"
+	partitionMetadataKey = "__partition"
+	topicMetadataKey     = "__topic"
 	skipVerify           = "skipVerify"
 	caCert               = "caCert"
 	certificateAuthType  = "certificate"

common/component/kafka/metadata_test.go

@@ -15,6 +15,7 @@ package kafka
 
 import (
 	"fmt"
+	"strconv"
 	"testing"
 	"time"
 
@@ -389,3 +390,56 @@ func TestMetadataChannelBufferSize(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, 128, meta.channelBufferSize)
 }
+
+func TestGetEventMetadata(t *testing.T) {
+	ts := time.Now()
+
+	t.Run("no headers", func(t *testing.T) {
+		m := sarama.ConsumerMessage{
+			Headers: nil, Timestamp: ts, Key: []byte("MyKey"), Value: []byte("MyValue"), Partition: 0, Offset: 123, Topic: "TestTopic",
+		}
+		act := GetEventMetadata(&m)
+		require.Len(t, act, 5)
+		require.Equal(t, strconv.FormatInt(ts.UnixMilli(), 10), act["__timestamp"])
+		require.Equal(t, "MyKey", act["__key"])
+		require.Equal(t, "0", act["__partition"])
+		require.Equal(t, "123", act["__offset"])
+		require.Equal(t, "TestTopic", act["__topic"])
+	})
+
+	t.Run("with headers", func(t *testing.T) {
+		headers := []*sarama.RecordHeader{
+			{Key: []byte("key1"), Value: []byte("value1")},
+			{Key: []byte("key2"), Value: []byte("value2")},
+		}
+		m := sarama.ConsumerMessage{
+			Headers: headers, Timestamp: ts, Key: []byte("MyKey"), Value: []byte("MyValue"), Partition: 0, Offset: 123, Topic: "TestTopic",
+		}
+		act := GetEventMetadata(&m)
+		require.Len(t, act, 7)
+		require.Equal(t, strconv.FormatInt(ts.UnixMilli(), 10), act["__timestamp"])
+		require.Equal(t, "MyKey", act["__key"])
+		require.Equal(t, "0", act["__partition"])
+		require.Equal(t, "123", act["__offset"])
+		require.Equal(t, "TestTopic", act["__topic"])
+		require.Equal(t, "value1", act["key1"])
+		require.Equal(t, "value2", act["key2"])
+	})
+
+	t.Run("no key", func(t *testing.T) {
+		m := sarama.ConsumerMessage{
+			Headers: nil, Timestamp: ts, Key: nil, Value: []byte("MyValue"), Partition: 0, Offset: 123, Topic: "TestTopic",
+		}
+		act := GetEventMetadata(&m)
+		require.Len(t, act, 4)
+		require.Equal(t, strconv.FormatInt(ts.UnixMilli(), 10), act["__timestamp"])
+		require.Equal(t, "0", act["__partition"])
+		require.Equal(t, "123", act["__offset"])
+		require.Equal(t, "TestTopic", act["__topic"])
+	})
+
+	t.Run("null message", func(t *testing.T) {
+		act := GetEventMetadata(nil)
+		require.Nil(t, act)
+	})
+}