Add manual partitioning to Kakfa output component (redpanda-data#857)

README.md

@@ -154,8 +154,6 @@ docker run --rm \
 	benthos -c /config.yaml
 ```
 
-There are a [few examples here][compose-examples] that show you some ways of setting up Benthos containers using `docker-compose`.
-
 ### ZMQ4 Support
 
 Benthos supports ZMQ4 for both data input and output. To add this you need to install libzmq4 and use the compile time flag when building Benthos:
@@ -181,7 +179,6 @@ Contributions are welcome, please [read the guidelines](CONTRIBUTING.md), come a
 [tracers]: https://www.benthos.dev/docs/components/tracers/about
 [metrics-config]: config/metrics
 [config-interp]: https://www.benthos.dev/docs/configuration/interpolation
-[compose-examples]: resources/docker/compose_examples
 [streams-api]: https://www.benthos.dev/docs/guides/streams_mode/streams_api
 [streams-mode]: https://www.benthos.dev/docs/guides/streams_mode/about
 [general-docs]: https://www.benthos.dev/docs/about

config/kafka.yaml

@@ -73,6 +73,7 @@ output:
     client_id: benthos_kafka_output
     key: ""
     partitioner: fnv1a_hash
+    partition: ""
     compression: none
     static_headers: {}
     metadata:

lib/output/kafka.go

@@ -46,7 +46,8 @@ However, this also means that manual intervention will eventually be required in
 			docs.FieldCommon("topic", "The topic to publish messages to.").IsInterpolated(),
 			docs.FieldCommon("client_id", "An identifier for the client connection."),
 			docs.FieldCommon("key", "The key to publish messages with.").IsInterpolated(),
-			docs.FieldCommon("partitioner", "The partitioning algorithm to use.").HasOptions("fnv1a_hash", "murmur2_hash", "random", "round_robin"),
+			docs.FieldCommon("partitioner", "The partitioning algorithm to use.").HasOptions("fnv1a_hash", "murmur2_hash", "random", "round_robin", "manual"),
+			docs.FieldCommon("partition", "The manually-specified partition to publish messages to.  Must be able to parse as a 32-bit integer.").IsInterpolated(),
 			docs.FieldCommon("compression", "The compression algorithm to use.").HasOptions("none", "snappy", "lz4", "gzip"),
 			docs.FieldString("static_headers", "An optional map of static headers that should be added to messages in addition to metadata.", map[string]string{"first-static-header": "value-1", "second-static-header": "value-2"}).Map(),
 			docs.FieldCommon("metadata", "Specify criteria for which metadata values are sent with messages as headers.").WithChildren(output.MetadataFields()...),

lib/output/writer/kafka.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"crypto/tls"
 	"fmt"
+	"strconv"
 	"strings"
 	"sync"
 	"time"
@@ -32,6 +33,7 @@ type KafkaConfig struct {
 	ClientID         string      `json:"client_id" yaml:"client_id"`
 	Key              string      `json:"key" yaml:"key"`
 	Partitioner      string      `json:"partitioner" yaml:"partitioner"`
+	Partition        string      `json:"partition" yaml:"partition"`
 	Topic            string      `json:"topic" yaml:"topic"`
 	Compression      string      `json:"compression" yaml:"compression"`
 	MaxMsgBytes      int         `json:"max_msg_bytes" yaml:"max_msg_bytes"`
@@ -65,6 +67,7 @@ func NewKafkaConfig() KafkaConfig {
 		Key:                  "",
 		RoundRobinPartitions: false,
 		Partitioner:          "fnv1a_hash",
+		Partition:            "",
 		Topic:                "benthos_stream",
 		Compression:          "none",
 		MaxMsgBytes:          1000000,
@@ -99,8 +102,9 @@ type Kafka struct {
 	version   sarama.KafkaVersion
 	conf      KafkaConfig
 
-	key   *field.Expression
-	topic *field.Expression
+	key       *field.Expression
+	topic     *field.Expression
+	partition *field.Expression
 
 	producer    sarama.SyncProducer
 	compression sarama.CompressionCodec
@@ -124,6 +128,13 @@ func NewKafka(conf KafkaConfig, mgr types.Manager, log log.Modular, stats metric
 		conf.Partitioner = "round_robin"
 		log.Warnln("The field 'round_robin_partitions' is deprecated, please use the 'partitioner' field (set to 'round_robin') instead.")
 	}
+
+	if conf.Partition == "" && conf.Partitioner == "manual" {
+		return nil, fmt.Errorf("partition field required for 'manual' partitioner")
+	} else if len(conf.Partition) > 0 && conf.Partitioner != "manual" {
+		return nil, fmt.Errorf("partition field can only be specified for 'manual' partitioner")
+	}
+
 	partitioner, err := strToPartitioner(conf.Partitioner)
 	if err != nil {
 		return nil, err
@@ -150,6 +161,9 @@ func NewKafka(conf KafkaConfig, mgr types.Manager, log log.Modular, stats metric
 	if k.topic, err = bloblang.NewField(conf.Topic); err != nil {
 		return nil, fmt.Errorf("failed to parse topic expression: %v", err)
 	}
+	if k.partition, err = bloblang.NewField(conf.Partition); err != nil {
+		return nil, fmt.Errorf("failed to parse parition expression: %v", err)
+	}
 	if k.backoffCtor, err = conf.Config.GetCtor(); err != nil {
 		return nil, err
 	}
@@ -214,6 +228,8 @@ func strToPartitioner(str string) (sarama.PartitionerConstructor, error) {
 		return sarama.NewRandomPartitioner, nil
 	case "round_robin":
 		return sarama.NewRoundRobinPartitioner, nil
+	case "manual":
+		return sarama.NewManualPartitioner, nil
 	default:
 	}
 	return nil, fmt.Errorf("partitioner not recognised: %v", str)
@@ -329,7 +345,8 @@ func (k *Kafka) WriteWithContext(ctx context.Context, msg types.Message) error {
 
 	userDefinedHeaders := k.buildUserDefinedHeaders(k.staticHeaders)
 	msgs := []*sarama.ProducerMessage{}
-	msg.Iter(func(i int, p types.Part) error {
+
+	err := msg.Iter(func(i int, p types.Part) error {
 		key := k.key.Bytes(i, msg)
 		nextMsg := &sarama.ProducerMessage{
 			Topic:    k.topic.String(i, msg),
@@ -340,11 +357,36 @@ func (k *Kafka) WriteWithContext(ctx context.Context, msg types.Message) error {
 		if len(key) > 0 {
 			nextMsg.Key = sarama.ByteEncoder(key)
 		}
+
+		// Only parse and set the partition if we are configured for manual
+		// partitioner.  Although samara will (currently) ignore the partition
+		// field when not using a manual partitioner, we should only set it when
+		// we explicitly want that.
+		if k.conf.Partitioner == "manual" {
+			partitionString := k.partition.String(i, msg)
+			if partitionString == "" {
+				return fmt.Errorf("partition expression failed to produce a value")
+			}
+
+			partitionInt, err := strconv.Atoi(partitionString)
+			if err != nil {
+				return fmt.Errorf("failed to parse valid integer from partition expression: %w", err)
+			}
+			if partitionInt < 0 {
+				return fmt.Errorf("invalid partition parsed from expression, must be >= 0, got %v", partitionInt)
+			}
+			// samara requires a 32-bit integer for the partition field
+			nextMsg.Partition = int32(partitionInt)
+		}
 		msgs = append(msgs, nextMsg)
 		return nil
 	})
 
-	err := producer.SendMessages(msgs)
+	if err != nil {
+		return err
+	}
+
+	err = producer.SendMessages(msgs)
 	for err != nil {
 		if pErrs, ok := err.(sarama.ProducerErrors); !k.conf.RetryAsBatch && ok {
 			if len(pErrs) == 0 {

lib/output/writer/tests/kafka_test.go

@@ -0,0 +1,35 @@
+package tests
+
+import (
+	"testing"
+
+	"github.com/Jeffail/benthos/v3/lib/log"
+	"github.com/Jeffail/benthos/v3/lib/metrics"
+	"github.com/Jeffail/benthos/v3/lib/output/writer"
+	"github.com/Jeffail/benthos/v3/lib/types"
+	"github.com/stretchr/testify/require"
+
+	_ "github.com/Jeffail/benthos/v3/public/components/all"
+)
+
+func TestKafkaConfigurationManualPartitioner(t *testing.T) {
+	conf := writer.NewKafkaConfig()
+	require.NoError(t, createKafkaWriter(conf), "Expected no error with default configuration")
+
+	conf.Partitioner = "manual"
+	require.Error(t, createKafkaWriter(conf), "Expected error with manual partitioner set and partition unset")
+
+	conf.Partition = "test"
+	require.NoError(t, createKafkaWriter(conf), "Expected no error with manual partitioner set and partition set")
+
+	conf.Partitioner = "random"
+	require.Error(t, createKafkaWriter(conf), "Expected error with non-manual partitioner set and partition set")
+
+	conf.Partition = ""
+	require.NoError(t, createKafkaWriter(conf), "Expected no error with non-manual partitioner set and partition unset")
+}
+
+func createKafkaWriter(conf writer.KafkaConfig) error {
+	_, err := writer.NewKafka(conf, types.NoopMgr(), log.Noop(), metrics.Noop())
+	return err
+}

lib/test/integration/kafka_test.go

@@ -228,6 +228,46 @@ input:
 			testOptVarThree("false"),
 		)
 	})
+
+	templateManualPartitioner := `
+output:
+  kafka:
+    addresses: [ localhost:$PORT ]
+    topic: topic-$ID
+    max_in_flight: $MAX_IN_FLIGHT
+    retry_as_batch: $VAR3
+    metadata:
+      exclude_prefixes: [ $OUTPUT_META_EXCLUDE_PREFIX ]
+    batching:
+      count: $OUTPUT_BATCH_COUNT
+    partitioner: manual
+    partition: '${! random_int() % 4 }'
+
+input:
+  kafka:
+    addresses: [ localhost:$PORT ]
+    topics: [ topic-$ID$VAR1 ]
+    consumer_group: "$VAR4"
+    checkpoint_limit: $VAR2
+    start_from_oldest: true
+    batching:
+      count: $INPUT_BATCH_COUNT
+`
+
+	t.Run("manual_partitioner", func(t *testing.T) {
+		t.Parallel()
+		suite.Run(
+			t, templateManualPartitioner,
+			testOptPreTest(func(t testing.TB, env *testEnvironment) {
+				env.configVars.var4 = "group" + env.configVars.id
+				require.NoError(t, createKafkaTopic("localhost:"+kafkaPortStr, env.configVars.id, 4))
+			}),
+			testOptPort(kafkaPortStr),
+			testOptVarTwo("1"),
+			testOptVarThree("false"),
+		)
+	})
+
 })
 
 func createKafkaTopic(address, id string, partitions int32) error {

website/docs/components/outputs/kafka.md

@@ -37,6 +37,7 @@ output:
     client_id: benthos_kafka_output
     key: ""
     partitioner: fnv1a_hash
+    partition: ""
     compression: none
     static_headers: {}
     metadata:
@@ -77,6 +78,7 @@ output:
     client_id: benthos_kafka_output
     key: ""
     partitioner: fnv1a_hash
+    partition: ""
     compression: none
     static_headers: {}
     metadata:
@@ -375,7 +377,16 @@ The partitioning algorithm to use.
 
 Type: `string`  
 Default: `"fnv1a_hash"`  
-Options: `fnv1a_hash`, `murmur2_hash`, `random`, `round_robin`.
+Options: `fnv1a_hash`, `murmur2_hash`, `random`, `round_robin`, `manual`.
+
+### `partition`
+
+The manually-specified partition to publish messages to.  Must be able to parse as a 32-bit integer.
+This field supports [interpolation functions](/docs/configuration/interpolation#bloblang-queries).
+
+
+Type: `string`  
+Default: `""`  
 
 ### `compression`