Merge pull request #51 from inloco/feature/add-key-to-json

Feature/add key to json

Author: João Assunção

README.md

@@ -39,6 +39,7 @@ To create new injectors for your topics, you should create a new kubernetes depl
 - `ES_TIME_SUFFIX` Indicates what time unit to append to index names on Elasticsearch. Supported values are `day` and `hour`. Default value is `day` **OPTIONAL**
 - `KAFKA_CONSUMER_RECORD_TYPE` Kafka record type. Should be set to "avro" or "json". Defaults to avro. **OPTIONAL**
 - `KAFKA_CONSUMER_METRICS_UPDATE_INTERVAL` The interval which the app updates the exported metrics in the format of golang's `time.ParseDuration`. Defaults to 30s. **OPTIONAL**
+- `KAFKA_CONSUMER_INCLUDE_KEY` Determines whether to include the Kafka key in the Elasticsearch message(as the "key" field). Defaults to false. **OPTIONAL**
 
 ### Important note about Elasticsearch mappings and types
 
@@ -159,7 +160,8 @@ To build the project from source, run from project root:
 go build cmd/injector.go
 ```
 
-To run tests, run `docker-compose up -d zookeeper kafka schema-registry elasticsearch` and run `make test`. 
+To run tests, first add kafka to your /etc/hosts file `echo "127.0.0.1 kafka" | sudo tee -a /etc/hosts` 
+and then run `docker-compose up -d zookeeper kafka schema-registry elasticsearch` and run `make test`. 
 
 ### Versioning
 

cmd/injector.go

@@ -45,6 +45,7 @@ func main() {
 		BufferSize:            os.Getenv("KAFKA_CONSUMER_BUFFER_SIZE"),
 		MetricsUpdateInterval: os.Getenv("KAFKA_CONSUMER_METRICS_UPDATE_INTERVAL"),
 		RecordType:            os.Getenv("KAFKA_CONSUMER_RECORD_TYPE"),
+		IncludeKey:            os.Getenv("KAFKA_CONSUMER_INCLUDE_KEY"),
 	}
 	metricsPublisher := metrics.NewMetricsPublisher()
 	service := injector.NewService(logger, metricsPublisher)

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,6 +37,15 @@ func MakeKafkaConsumer(endpoints Endpoints, logger log.Logger, schemaRegistry *s
 		SchemaRegistry: schemaRegistry,
 	}
 
+	includeKey, err := strconv.ParseBool(kafkaConfig.IncludeKey)
+	if err != nil {
+		err = level.Warn(logger).Log("err", err, "message", "failed to get consumer include key configuration flag")
+		if err != nil {
+			panic(err)
+		}
+		includeKey = false
+	}
+
 	return kafka.Consumer{
 		Topics:                kafkaConfig.Topics,
 		Group:                 kafkaConfig.ConsumerGroup,
@@ -47,5 +56,6 @@ func MakeKafkaConsumer(endpoints Endpoints, logger log.Logger, schemaRegistry *s
 		BatchSize:             batchSize,
 		MetricsUpdateInterval: metricsUpdateInterval,
 		BufferSize:            bufferSize,
+		IncludeKey:            includeKey,
 	}, nil
 }

src/kafka/config.go

@@ -13,4 +13,5 @@ type Config struct {
 	MetricsUpdateInterval string
 	BufferSize            string
 	RecordType            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,6 +43,7 @@ type Consumer struct {
 	BatchSize             int
 	MetricsUpdateInterval time.Duration
 	BufferSize            int
+	IncludeKey            bool
 }
 
 type topicPartitionOffset struct {
@@ -144,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)
+					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,19 +11,20 @@ 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/models"
 	"github.com/inloco/kafka-elasticsearch-injector/src/schema_registry"
+	"github.com/linkedin/goavro/v2"
 )
 
 // DecodeMessageFunc extracts a user-domain request object from an Kafka
 // message object. It's designed to be used in Kafka consumers.
 // One straightforward DecodeMessageFunc could be something that
 // Avro decodes the message body to the concrete response type.
-type DecodeMessageFunc func(context.Context, *sarama.ConsumerMessage) (record *models.Record, err error)
+type DecodeMessageFunc func(context.Context, *sarama.ConsumerMessage, bool) (record *models.Record, err error)
 
 const kafkaTimestampKey = "@timestamp"
+const keyField = "key"
 
 type Decoder struct {
 	SchemaRegistry *schema_registry.SchemaRegistry
@@ -38,32 +39,12 @@ func (d *Decoder) DeserializerFor(recordType string) DecodeMessageFunc {
 	}
 }
 
-func (d *Decoder) AvroMessageToRecord(context context.Context, msg *sarama.ConsumerMessage) (*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
 	}
 
-	schemaId := getSchemaId(msg)
-	avroRecord := msg.Value[5:]
-	schema, err := d.SchemaRegistry.GetSchema(schemaId)
-	if err != nil {
-		return nil, err
-	}
-	var codec *goavro.Codec
-	if codecI, ok := d.CodecCache.Load(schemaId); ok {
-		codec, ok = codecI.(*goavro.Codec)
-	}
-
-	if codec == nil {
-		codec, err = goavro.NewCodec(schema)
-		if err != nil {
-			return nil, err
-		}
-
-		d.CodecCache.Store(schemaId, codec)
-	}
-
-	native, _, err := codec.NativeFromBinary(avroRecord)
+	native, err := d.nativeFromBinary(msg.Value)
 	if err != nil {
 		return nil, err
 	}
@@ -82,6 +63,14 @@ func (d *Decoder) AvroMessageToRecord(context context.Context, msg *sarama.Consu
 
 	parsedNative[kafkaTimestampKey] = makeTimestamp(msg.Timestamp)
 
+	if includeKey && msg.Key != nil {
+		nativeKey, err := d.nativeFromBinary(msg.Key)
+		if err != nil {
+			return nil, err
+		}
+		parsedNative[keyField] = nativeKey
+	}
+
 	return &models.Record{
 		Topic:     msg.Topic,
 		Partition: msg.Partition,
@@ -95,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) (*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 {
@@ -105,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,
@@ -114,7 +112,36 @@ func (d *Decoder) JsonMessageToRecord(context context.Context, msg *sarama.Consu
 	}, nil
 }
 
-func getSchemaId(msg *sarama.ConsumerMessage) int32 {
-	schemaIdBytes := msg.Value[1:5]
+func (d *Decoder) nativeFromBinary(value []byte) (interface{}, error) {
+	schemaId := getSchemaId(value)
+	avroRecord := value[5:]
+	schema, err := d.SchemaRegistry.GetSchema(schemaId)
+	if err != nil {
+		return nil, err
+	}
+	var codec *goavro.Codec
+	if codecI, ok := d.CodecCache.Load(schemaId); ok {
+		codec, _ = codecI.(*goavro.Codec)
+	}
+
+	if codec == nil {
+		codec, err = goavro.NewCodec(schema)
+		if err != nil {
+			return nil, err
+		}
+
+		d.CodecCache.Store(schemaId, codec)
+	}
+
+	native, _, err := codec.NativeFromBinary(avroRecord)
+	if err != nil {
+		return nil, err
+	}
+
+	return native, nil
+}
+
+func getSchemaId(value []byte) int32 {
+	schemaIdBytes := value[1:5]
 	return int32(schemaIdBytes[0])<<24 | int32(schemaIdBytes[1])<<16 | int32(schemaIdBytes[2])<<8 | int32(schemaIdBytes[3])
 }

src/kafka/encode_decode_test.go

@@ -13,26 +13,36 @@ 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,
 		Topic:     "test",
 		Partition: 1,
 		Offset:    54,
 		Timestamp: time.Now(),
-	})
+	}, false)
 	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)
@@ -47,15 +57,71 @@ func TestDecoder_JsonMessageToRecord_MalformedJson(t *testing.T) {
 		Partition: 1,
 		Offset:    54,
 		Timestamp: time.Now(),
-	})
+	}, false)
 	assert.Nil(t, record)
 	assert.NotNil(t, err)
 }
 
 func TestDecoder_AvroMessageToRecord_NilMessageValue(t *testing.T) {
 	d := &Decoder{CodecCache: sync.Map{}}
-	record, err := d.AvroMessageToRecord(context.Background(), &sarama.ConsumerMessage{Value: nil, Topic: "test", Partition: 1, Offset: 54, Timestamp: time.Now()})
+	record, err := d.AvroMessageToRecord(context.Background(), &sarama.ConsumerMessage{
+		Value:     nil,
+		Topic:     "test",
+		Partition: 1, Offset: 54,
+		Timestamp: time.Now()},
+		false)
 	isErrNilMessage := errors.Is(err, e.ErrNilMessage)
 	assert.Nil(t, record)
 	assert.True(t, isErrNilMessage)
 }
+
+func TestDecoder_AvroMessageToRecord_NilMessageValue_IncludeKey(t *testing.T) {
+	d := &Decoder{CodecCache: sync.Map{}}
+	key := dummyKey{"marco"}
+	jsonBytesKey, _ := json.Marshal(key)
+	record, err := d.AvroMessageToRecord(context.Background(), &sarama.ConsumerMessage{
+		Value:     nil,
+		Key:       jsonBytesKey,
+		Topic:     "test",
+		Partition: 1, Offset: 54,
+		Timestamp: time.Now()},
+		true)
+	isErrNilMessage := errors.Is(err, e.ErrNilMessage)
+	assert.Nil(t, record)
+	assert.True(t, isErrNilMessage)
+}
+
+func TestDecoder_JsonMessageToRecord_IncludeKey(t *testing.T) {
+	d := &Decoder{CodecCache: sync.Map{}}
+
+	key := dummyKey{"marco"}
+	jsonBytesKey, _ := json.Marshal(key)
+
+	val := dummyValue{"pop", 60}
+	jsonBytesValue, _ := 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.Nil(t, err)
+	assert.Equal(t, expected, returnedKeyIncluded)
+}