RHINENG-22325: extend KafkaMessage with Headers

Author: MichaelMraka

base/mqueue/message.go

@@ -1,12 +1,16 @@
 package mqueue
 
-import "github.com/bytedance/sonic"
+import (
+	"github.com/bytedance/sonic"
+	"github.com/segmentio/kafka-go"
+)
 
-func MessageFromJSON(k string, v interface{}) (KafkaMessage, error) {
+func MessageFromJSON(k string, v interface{}, h []kafka.Header) (KafkaMessage, error) {
 	var m KafkaMessage
 	var err error
 
 	m.Key = []byte(k)
+	m.Headers = h
 	m.Value, err = sonic.Marshal(v)
 	return m, err
 }

base/mqueue/mqueue.go

@@ -12,6 +12,7 @@ import (
 	"sync"
 
 	"github.com/lestrrat-go/backoff/v2"
+	"github.com/segmentio/kafka-go"
 )
 
 const errContextCanceled = "context canceled"
@@ -42,8 +43,9 @@ func createLoggerFunc(counter Counter) func(fmt string, args ...interface{}) {
 }
 
 type KafkaMessage struct {
-	Key   []byte
-	Value []byte
+	Key     []byte
+	Value   []byte
+	Headers []kafka.Header
 }
 
 type MessageHandler func(message KafkaMessage) error

base/mqueue/mqueue_impl_gokafka.go

@@ -39,7 +39,7 @@ func (t *kafkaGoReaderImpl) HandleMessages(handler MessageHandler) {
 			}
 		}
 		// At this level, all errors are fatal
-		kafkaMessage := KafkaMessage{Key: m.Key, Value: m.Value}
+		kafkaMessage := KafkaMessage{Key: m.Key, Value: m.Value, Headers: m.Headers}
 		if err = handler(kafkaMessage); err != nil {
 			utils.LogPanic("err", err.Error(), "Handler failed")
 		}
@@ -61,7 +61,7 @@ type kafkaGoWriterImpl struct {
 func (t *kafkaGoWriterImpl) WriteMessages(ctx context.Context, msgs ...KafkaMessage) error {
 	kafkaGoMessages := make([]kafka.Message, len(msgs))
 	for i, m := range msgs {
-		kafkaGoMessages[i] = kafka.Message{Key: m.Key, Value: m.Value}
+		kafkaGoMessages[i] = kafka.Message{Key: m.Key, Value: m.Value, Headers: m.Headers}
 	}
 	err := t.Writer.WriteMessages(ctx, kafkaGoMessages...)
 	return err

evaluator/notifications.go

@@ -120,7 +120,7 @@ func publishNewAdvisoriesNotification(tx *gorm.DB, system *models.SystemPlatform
 		return errors.Wrap(err, "creating notification failed")
 	}
 
-	msg, err := mqueue.MessageFromJSON(system.InventoryID, notif)
+	msg, err := mqueue.MessageFromJSON(system.InventoryID, notif, nil)
 	if err != nil {
 		return errors.Wrap(err, "creating message from notification failed")
 	}

evaluator/notifications_test.go

@@ -114,7 +114,7 @@ func TestAdvisoriesNotificationMessage(t *testing.T) {
 	assert.Equal(t, displayName, notification.Context.DisplayName)
 	assert.Equal(t, tags, notification.Context.Tags)
 
-	msg, err := mqueue.MessageFromJSON(inventoryID, notification)
+	msg, err := mqueue.MessageFromJSON(inventoryID, notification, nil)
 	assert.Nil(t, err)
 	assert.Equal(t, inventoryID, string(msg.Key))
 

evaluator/remediations.go

@@ -87,7 +87,7 @@ func publishRemediationsState(system *models.SystemPlatform, response *vmaas.Upd
 	}
 
 	state := createRemediationsStateMsg(system.InventoryID, response)
-	msg, err := mqueue.MessageFromJSON(system.InventoryID, state)
+	msg, err := mqueue.MessageFromJSON(system.InventoryID, state, nil)
 	if err != nil {
 		return errors.Wrap(err, "formatting message")
 	}