transmit key from input topic to output topic (#3578)

SeldonIO · Sep 15, 2021 · ab1868b · ab1868b
1 parent 8988f5c
commit ab1868b
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/executor/api/kafka/server.go b/executor/api/kafka/server.go
@@ -231,6 +231,7 @@ func (ks *SeldonKafkaServer) Serve() error {
 
 				job := KafkaJob{
 					headers:    headers,
+					reqKey:     e.Key,
 					reqPayload: reqPayload,
 				}
 				// enqueue a job

diff --git a/executor/api/kafka/worker.go b/executor/api/kafka/worker.go
@@ -12,6 +12,7 @@ import (
 
 type KafkaJob struct {
 	headers    map[string][]string
+	reqKey     []byte
 	reqPayload payload.SeldonPayload
 }
 
@@ -62,6 +63,7 @@ func (ks *SeldonKafkaServer) processKafkaRequest(job *KafkaJob) {
 
 	err = ks.Producer.Produce(&kafka.Message{
 		TopicPartition: kafka.TopicPartition{Topic: &ks.TopicOut, Partition: kafka.PartitionAny},
+		Key:            job.reqKey,
 		Value:          resBytes,
 		Headers:        kafkaHeaders,
 	}, nil)