refactor(all): changes withKeyAndValue to includeKey also adds a test for when this parameter is true

Author: João Assunção

src/injector/injector.go

@@ -5,10 +5,10 @@ import (
 
 	"time"
 
-	"github.com/inloco/kafka-elasticsearch-injector/src/kafka"
-	"github.com/inloco/kafka-elasticsearch-injector/src/schema_registry"
 	"github.com/go-kit/kit/log"
 	"github.com/go-kit/kit/log/level"
+	"github.com/inloco/kafka-elasticsearch-injector/src/kafka"
+	"github.com/inloco/kafka-elasticsearch-injector/src/schema_registry"
 )
 
 func MakeKafkaConsumer(endpoints Endpoints, logger log.Logger, schemaRegistry *schema_registry.SchemaRegistry, kafkaConfig *kafka.Config) (kafka.Consumer, error) {
@@ -37,10 +37,10 @@ func MakeKafkaConsumer(endpoints Endpoints, logger log.Logger, schemaRegistry *s
 		SchemaRegistry: schemaRegistry,
 	}
 
-	withKeyAndValue, err := strconv.ParseBool(kafkaConfig.WithKeyAndValue)
+	includeKey, err := strconv.ParseBool(kafkaConfig.IncludeKey)
 	if err != nil {
 		level.Warn(logger).Log("err", err, "message", "failed to get consumer with key and value")
-		withKeyAndValue = false
+		includeKey = false
 	}
 
 	return kafka.Consumer{
@@ -53,6 +53,6 @@ func MakeKafkaConsumer(endpoints Endpoints, logger log.Logger, schemaRegistry *s
 		BatchSize:             batchSize,
 		MetricsUpdateInterval: metricsUpdateInterval,
 		BufferSize:            bufferSize,
-		WithKeyAndValue:       withKeyAndValue,
+		IncludeKey:            includeKey,
 	}, nil
 }

src/kafka/config.go

@@ -13,5 +13,5 @@ type Config struct {
 	MetricsUpdateInterval string
 	BufferSize            string
 	RecordType            string
-	WithKeyAndValue       string
+	IncludeKey            string
 }

src/kafka/consumer.go

@@ -8,7 +8,7 @@ import (
 	"time"
 
 	"github.com/Shopify/sarama"
-	"github.com/bsm/sarama-cluster"
+	cluster "github.com/bsm/sarama-cluster"
 	"github.com/go-kit/kit/endpoint"
 	"github.com/go-kit/kit/log"
 	"github.com/go-kit/kit/log/level"
@@ -43,7 +43,7 @@ type Consumer struct {
 	BatchSize             int
 	MetricsUpdateInterval time.Duration
 	BufferSize            int
-	WithKeyAndValue       bool
+	IncludeKey            bool
 }
 
 type topicPartitionOffset struct {
@@ -145,7 +145,7 @@ func (k *kafka) worker(consumer *cluster.Consumer, buffSize int, notifications c
 		for idx == buffSize {
 			if decoded == nil {
 				for _, msg := range buf {
-					req, err := k.consumer.Decoder(nil, msg, k.consumer.WithKeyAndValue)
+					req, err := k.consumer.Decoder(nil, msg, k.consumer.IncludeKey)
 					if err != nil {
 						if errors.Is(err, e.ErrNilMessage) {
 							continue

src/kafka/encode_decode.go

@@ -11,10 +11,10 @@ import (
 	"sync"
 
 	"github.com/Shopify/sarama"
-	"github.com/linkedin/goavro/v2"
 	"github.com/inloco/kafka-elasticsearch-injector/src/models"
 	e "github.com/inloco/kafka-elasticsearch-injector/src/errors"
 	"github.com/inloco/kafka-elasticsearch-injector/src/schema_registry"
+	"github.com/linkedin/goavro/v2"
 )
 
 // DecodeMessageFunc extracts a user-domain request object from an Kafka
@@ -39,7 +39,7 @@ func (d *Decoder) DeserializerFor(recordType string) DecodeMessageFunc {
 	}
 }
 
-func (d *Decoder) AvroMessageToRecord(context context.Context, msg *sarama.ConsumerMessage, withKeyAndValue bool) (*models.Record, error) {
+func (d *Decoder) AvroMessageToRecord(context context.Context, msg *sarama.ConsumerMessage, includeKey bool) (*models.Record, error) {
 	if msg.Value == nil {
 		return nil, e.ErrNilMessage
 	}
@@ -63,7 +63,7 @@ func (d *Decoder) AvroMessageToRecord(context context.Context, msg *sarama.Consu
 
 	parsedNative[kafkaTimestampKey] = makeTimestamp(msg.Timestamp)
 
-	if withKeyAndValue && msg.Key != nil {
+	if includeKey && msg.Key != nil {
 		nativeKey, err := d.nativeFromBinary(msg.Key)
 		if err != nil {
 			return nil, err
@@ -84,8 +84,9 @@ func makeTimestamp(timestamp time.Time) int64 {
 	return timestamp.UnixNano() / int64(time.Millisecond)
 }
 
-func (d *Decoder) JsonMessageToRecord(context context.Context, msg *sarama.ConsumerMessage, withKeyAndValue bool) (*models.Record, error) {
+func (d *Decoder) JsonMessageToRecord(context context.Context, msg *sarama.ConsumerMessage, includeKey bool) (*models.Record, error) {
 	var jsonValue map[string]interface{}
+	var jsonKey map[string]interface{}
 	err := json.Unmarshal(msg.Value, &jsonValue)
 
 	if err != nil {
@@ -94,6 +95,14 @@ func (d *Decoder) JsonMessageToRecord(context context.Context, msg *sarama.Consu
 
 	jsonValue[kafkaTimestampKey] = makeTimestamp(msg.Timestamp)
 
+	if includeKey && msg.Key != nil {
+		err := json.Unmarshal(msg.Key, &jsonKey)
+		if err != nil {
+			return nil, err
+		}
+		jsonValue[keyField] = jsonKey
+	}
+
 	return &models.Record{
 		Topic:     msg.Topic,
 		Partition: msg.Partition,

src/kafka/encode_decode_test.go

@@ -13,14 +13,24 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-type dummy struct {
+type dummyValue struct {
 	Id        string `json:"id"`
 	Timestamp int64  `json:"timestamp"`
 }
 
+type dummyKey struct {
+	Id string `json:"id"`
+}
+
+type dummyIncludeKey struct {
+	Id        string   `json:"id"`
+	Timestamp int64    `json:"timestamp"`
+	Key       dummyKey `json:"key"`
+}
+
 func TestDecoder_JsonMessageToRecord(t *testing.T) {
 	d := &Decoder{CodecCache: sync.Map{}}
-	val := dummy{"alo", 60}
+	val := dummyValue{"alo", 60}
 	jsonBytes, err := json.Marshal(val)
 	record, err := d.JsonMessageToRecord(context.Background(), &sarama.ConsumerMessage{
 		Value:     jsonBytes,
@@ -32,7 +42,7 @@ func TestDecoder_JsonMessageToRecord(t *testing.T) {
 	assert.Nil(t, err)
 	returnedJsonBytes, err := json.Marshal(record.Json)
 	assert.Nil(t, err)
-	var returnedVal dummy
+	var returnedVal dummyValue
 	err = json.Unmarshal(returnedJsonBytes, &returnedVal)
 	assert.Nil(t, err)
 	assert.Equal(t, val, returnedVal)
@@ -59,3 +69,30 @@ func TestDecoder_AvroMessageToRecord_NilMessageValue(t *testing.T) {
 	assert.Nil(t, record)
 	assert.True(t, isErrNilMessage)
 }
+
+func TestDecoder_JsonMessageToRecord_IncludeKey(t *testing.T) {
+	d := &Decoder{CodecCache: sync.Map{}}
+	key := dummyKey{"marco"}
+	jsonBytesKey, err := json.Marshal(key)
+	val := dummyValue{"pop", 60}
+	jsonBytesValue, err := json.Marshal(val)
+	expected := dummyIncludeKey{"pop", 60, dummyKey{"marco"}}
+	record, err := d.JsonMessageToRecord(context.Background(), &sarama.ConsumerMessage{
+		Key:       jsonBytesKey,
+		Value:     jsonBytesValue,
+		Topic:     "test",
+		Partition: 1,
+		Offset:    54,
+		Timestamp: time.Now(),
+	}, true)
+	assert.Nil(t, err)
+	returnedJsonBytes, err := json.Marshal(record.Json)
+	assert.Nil(t, err)
+	var returnedVal dummyValue
+	err = json.Unmarshal(returnedJsonBytes, &returnedVal)
+	assert.Nil(t, err)
+	assert.Equal(t, val, returnedVal)
+	var returnedKeyIncluded dummyIncludeKey
+	err = json.Unmarshal(returnedJsonBytes, &returnedKeyIncluded)
+	assert.Equal(t, expected, returnedKeyIncluded)
+}