Merge branch 'hotfix/check-msg-with-nil-values'

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/elasticsearch/elasticsearch.go

@@ -89,6 +89,7 @@ func (d recordDatabase) Insert(records []*models.ElasticRecord) (*InsertResponse
 		return &InsertResponse{AlreadyExists: nil, Retry: records, Backoff: true}, nil
 	}
 	if err != nil {
+		_ = level.Debug(d.logger).Log("message", "something went wrong with elasticsearch", "err", err)
 		return nil, err
 	}
 	if res.Errors {
@@ -112,10 +113,12 @@ func (d recordDatabase) Insert(records []*models.ElasticRecord) (*InsertResponse
 			}
 			for _, f := range failed {
 				if f.Status == http.StatusBadRequest {
+					_ = level.Debug(d.logger).Log("message", "elasticsearch bad requests", "err", f)
 					d.metricsPublisher.ElasticsearchBadRequests(1)
 					continue
 				}
 				if f.Status == http.StatusConflict {
+					_ = level.Debug(d.logger).Log("message", "elasticsearch conflicts", "err", f)
 					d.metricsPublisher.ElasticsearchConflicts(1)
 					continue
 				}

src/errors/ErrNilMessage.go

@@ -0,0 +1,5 @@
+package errors
+
+import "errors"
+
+var ErrNilMessage = errors.New("message value is nil")

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/injector/store/store.go

@@ -21,6 +21,10 @@ type basicStore struct {
 }
 
 func (s basicStore) Insert(records []*models.Record) error {
+	if len(records) == 0 {
+		return nil
+	}
+	
 	elasticRecords, err := s.codec.EncodeElasticRecords(records)
 	if err != nil {
 		return err

src/kafka/config.go

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

src/kafka/consumer.go

@@ -3,16 +3,18 @@ package kafka
 import (
 	"context"
 	"os"
+	"errors"
 
 	"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"
 	"github.com/inloco/kafka-elasticsearch-injector/src/metrics"
 	"github.com/inloco/kafka-elasticsearch-injector/src/models"
+	e "github.com/inloco/kafka-elasticsearch-injector/src/errors"
 )
 
 type Notification int32
@@ -41,6 +43,7 @@ type Consumer struct {
 	BatchSize             int
 	MetricsUpdateInterval time.Duration
 	BufferSize            int
+	IncludeKey            bool
 }
 
 type topicPartitionOffset struct {
@@ -142,8 +145,12 @@ 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
+						}
+
 						level.Error(k.consumer.Logger).Log(
 							"message", "Error decoding message",
 							"err", err.Error(),

src/kafka/consumer_test.go

@@ -30,6 +30,10 @@ type fixtureService struct {
 }
 
 func (s fixtureService) Insert(records []*models.Record) error {
+	if len(records) == 0 {
+		return nil
+	}
+
 	elasticRecords, err := s.codec.EncodeElasticRecords(records)
 	if err != nil {
 		return err
@@ -63,8 +67,11 @@ var (
 	service   = fixtureService{db, codec}
 	endpoints = &fixtureEndpoints{
 		func(ctx context.Context, request interface{}) (response interface{}, err error) {
-			records := request.([]*models.Record)
+			if request == nil {
+				return nil, nil
+			}
 
+			records := request.([]*models.Record)
 			return nil, service.Insert(records)
 		},
 	}
@@ -111,7 +118,6 @@ func TestKafka_Start(t *testing.T) {
 	rec := fixtures.NewFixtureRecord()
 	var msg *sarama.ProducerMessage
 	if assert.NoError(t, err) {
-
 		err = producer.Publish(rec)
 		if assert.NoError(t, err) {
 			msg = <-producer.GetSuccesses()

src/kafka/encode_decode.go

@@ -11,18 +11,20 @@ import (
 	"sync"
 
 	"github.com/Shopify/sarama"
-	"github.com/linkedin/goavro/v2"
+	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
@@ -37,28 +39,12 @@ func (d *Decoder) DeserializerFor(recordType string) DecodeMessageFunc {
 	}
 }
 
-func (d *Decoder) AvroMessageToRecord(context context.Context, msg *sarama.ConsumerMessage) (*models.Record, error) {
-	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)
+func (d *Decoder) AvroMessageToRecord(context context.Context, msg *sarama.ConsumerMessage, includeKey bool) (*models.Record, error) {
+	if msg.Value == nil {
+		return nil, e.ErrNilMessage
 	}
 
-	native, _, err := codec.NativeFromBinary(avroRecord)
+	native, err := d.nativeFromBinary(msg.Value)
 	if err != nil {
 		return nil, err
 	}
@@ -77,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,
@@ -90,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 {
@@ -100,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,
@@ -109,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])
 }