Repository: redpanda-data/connect (language: go)

Issue:
Possibility to set the number of partitions for new topics created by kafka output
I would like to add the possibility to set the number of partitions for topics created by kafka output.

Let me explain my use case:
Our application runs on Kubernetes, where we have a Redpanda instance alongside multiple Benthos instances and other stuff. The problem with running Redpanda in k8s with Helm, is that one cannot set a default number of partitions for new topics (as seen in [this](https://github.com/redpanda-data/helm-charts/blob/af6c466e3d5e1a4d1fc366430bc97b0d445ae4c7/charts/redpanda/values.yaml#L940C20-L940C20) line).
In other components of our application we use the Sarama library to generate and manage consumers and producers, and it's possible to configure the number of partitions when creating a new kafka topic. Unfortunately, the data processed through Benthos does not come from or goes to any of those other components, therefore when producing new kafka topics from Benthos, Redpanda uses it's default value of 1 partition per topic.

It would be great to implement a `partitions` parameter that, if set, creates the new topic with the set number of partition, using the [`ClusterAdmin`](https://pkg.go.dev/github.com/Shopify/sarama@v1.38.1#ClusterAdmin) interface of sarama. It would also be possible to cache the topic names for increased speed. An example implementation can be seen in our sarama-wrapper, where in [this function](https://github.com/united-manufacturing-hub/Sarama-Kafka-Wrapper/blob/8f7f860c384b1afac4434d394b98f10af9a108f3/pkg/kafka/client.go#L254) we check if a topic exists and eventually create it before enqueueing the message (we use an async producer).
Of course, if the parameter is not set (or set to 0 maybe), then it would just not set the `NumPartitions` filed in the `sarama.TopicDetail` struct before creating the topic, thus using whatever default value is set in the kafka broker.

It would be nice to have some feedback about this functionality, and if it would be well received by the users, so that I could start working on the implementation.


Discussion / hints:
Hey @DaruZero, I think that'd be generally nice to have. The topic is potentially dynamic as it's determined by interpolation functions, which means we can't determine the topic once at start up. That means whichever mechanism we use for customising topic creation needs to have a minimal impact on traffic.
Using a `syncmap` as cache we can check if the topic has already been created, and skip the topic creation process if that's the case.
Imho there are two options:
- use `ClusterAdmin.ListTopics()` at startup to prefetch all the existing topics and store them in cache. This can possibly slowing the startup if the broker has already many topics, but we can be sure that any topic not present in the cache is new and has to be created first
- only cache the topics coming in, being careful to check if `ClusterAdmin.CreateTopic()` returns a `TOPIC_ALREADY_EXISTS` error, in which case the topic can be cached and the message safely sent
Actually, after further analysis it turns out that calling `ClusterAdmin.ListTopics()` every time we hit the cache (like in the function that I linked above), shouldn't have a huge impact on traffic or performance, as kafka is not meant to have more than a couple thousands topics

Fix patch:
diff --git a/internal/component/output/config_kafka.go b/internal/component/output/config_kafka.go
index ab55e891ce..2dcfd7b9a0 100644
--- a/internal/component/output/config_kafka.go
+++ b/internal/component/output/config_kafka.go
@@ -10,27 +10,34 @@ import (
 
 // KafkaConfig contains configuration fields for the Kafka output type.
 type KafkaConfig struct {
-	Addresses        []string    `json:"addresses" yaml:"addresses"`
-	ClientID         string      `json:"client_id" yaml:"client_id"`
-	RackID           string      `json:"rack_id" yaml:"rack_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"`
-	Timeout          string      `json:"timeout" yaml:"timeout"`
-	AckReplicas      bool        `json:"ack_replicas" yaml:"ack_replicas"`
-	TargetVersion    string      `json:"target_version" yaml:"target_version"`
-	TLS              btls.Config `json:"tls" yaml:"tls"`
-	SASL             sasl.Config `json:"sasl" yaml:"sasl"`
-	MaxInFlight      int         `json:"max_in_flight" yaml:"max_in_flight"`
-	retries.Config   `json:",inline" yaml:",inline"`
-	RetryAsBatch     bool                         `json:"retry_as_batch" yaml:"retry_as_batch"`
-	Batching         batchconfig.Config           `json:"batching" yaml:"batching"`
-	StaticHeaders    map[string]string            `json:"static_headers" yaml:"static_headers"`
-	Metadata         metadata.ExcludeFilterConfig `json:"metadata" yaml:"metadata"`
-	InjectTracingMap string                       `json:"inject_tracing_map" yaml:"inject_tracing_map"`
+	Addresses           []string                  `json:"addresses" yaml:"addresses"`
+	ClientID            string                    `json:"client_id" yaml:"client_id"`
+	RackID              string                    `json:"rack_id" yaml:"rack_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"`
+	CustomTopicCreation customTopicCreationConfig `json:"custom_topic_creation" yaml:"custom_topic_creation"`
+	Compression         string                    `json:"compression" yaml:"compression"`
+	MaxMsgBytes         int                       `json:"max_msg_bytes" yaml:"max_msg_bytes"`
+	Timeout             string                    `json:"timeout" yaml:"timeout"`
+	AckReplicas         bool                      `json:"ack_replicas" yaml:"ack_replicas"`
+	TargetVersion       string                    `json:"target_version" yaml:"target_version"`
+	TLS                 btls.Config               `json:"tls" yaml:"tls"`
+	SASL                sasl.Config               `json:"sasl" yaml:"sasl"`
+	MaxInFlight         int                       `json:"max_in_flight" yaml:"max_in_flight"`
+	retries.Config      `json:",inline" yaml:",inline"`
+	RetryAsBatch        bool                         `json:"retry_as_batch" yaml:"retry_as_batch"`
+	Batching            batchconfig.Config           `json:"batching" yaml:"batching"`
+	StaticHeaders       map[string]string            `json:"static_headers" yaml:"static_headers"`
+	Metadata            metadata.ExcludeFilterConfig `json:"metadata" yaml:"metadata"`
+	InjectTracingMap    string                       `json:"inject_tracing_map" yaml:"inject_tracing_map"`
+}
+
+type customTopicCreationConfig struct {
+	Enabled           bool `json:"enabled" yaml:"enabled"`
+	Partitions        int  `json:"partitions" yaml:"partitions"`
+	ReplicationFactor int  `json:"replication_factor" yaml:"replication_factor"`
 }
 
 // NewKafkaConfig creates a new KafkaConfig with default values.
@@ -41,13 +48,18 @@ func NewKafkaConfig() KafkaConfig {
 	rConf.Backoff.MaxElapsedTime = "30s"
 
 	return KafkaConfig{
-		Addresses:     []string{},
-		ClientID:      "benthos",
-		RackID:        "",
-		Key:           "",
-		Partitioner:   "fnv1a_hash",
-		Partition:     "",
-		Topic:         "",
+		Addresses:   []string{},
+		ClientID:    "benthos",
+		RackID:      "",
+		Key:         "",
+		Partitioner: "fnv1a_hash",
+		Partition:   "",
+		Topic:       "",
+		CustomTopicCreation: customTopicCreationConfig{
+			Enabled:           false,
+			Partitions:        -1,
+			ReplicationFactor: -1,
+		},
 		Compression:   "none",
 		MaxMsgBytes:   1000000,
 		Timeout:       "5s",
diff --git a/internal/impl/kafka/output_sarama_kafka.go b/internal/impl/kafka/output_sarama_kafka.go
index dd80632d66..a55f2d5a59 100644
--- a/internal/impl/kafka/output_sarama_kafka.go
+++ b/internal/impl/kafka/output_sarama_kafka.go
@@ -12,6 +12,7 @@ import (
 
 	"github.com/IBM/sarama"
 	"github.com/cenkalti/backoff/v4"
+	"golang.org/x/sync/syncmap"
 
 	batchInternal "github.com/benthosdev/benthos/v4/internal/batch"
 	"github.com/benthosdev/benthos/v4/internal/batch/policy"
@@ -68,6 +69,11 @@ Unfortunately this error message will appear for a wide range of connection prob
 			docs.FieldString("key", "The key to publish messages with.").IsInterpolated(),
 			docs.FieldString("partitioner", "The partitioning algorithm to use.").HasOptions("fnv1a_hash", "murmur2_hash", "random", "round_robin", "manual"),
 			docs.FieldString("partition", "The manually-specified partition to publish messages to, relevant only when the field `partitioner` is set to `manual`. Must be able to parse as a 32-bit integer.").IsInterpolated().Advanced(),
+			docs.FieldObject("custom_topic_creation", "If enabled, topics will be created with the specified number of partitions and replication factor if they do not already exist.").WithChildren(
+				docs.FieldBool("enabled", "Whether to enable custom topic creation.").HasDefault(false),
+				docs.FieldInt("partitions", "The number of partitions to create for new topics. Must be greater than 1").HasDefault(-1),
+				docs.FieldInt("replication_factor", "The replication factor to use for new topics. Must be an odd number, and less then or equal to the number of brokers").HasDefault(-1),
+			).Advanced(),
 			docs.FieldString("compression", "The compression algorithm to use.").HasOptions("none", "snappy", "lz4", "gzip", "zstd"),
 			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.FieldObject("metadata", "Specify criteria for which metadata values are sent with messages as headers.").WithChildren(metadata.ExcludeFilterFields()...),
@@ -131,6 +137,7 @@ type kafkaWriter struct {
 	topic     *field.Expression
 	partition *field.Expression
 
+	admin       sarama.ClusterAdmin
 	producer    sarama.SyncProducer
 	compression sarama.CompressionCodec
 	partitioner sarama.PartitionerConstructor
@@ -138,7 +145,8 @@ type kafkaWriter struct {
 	staticHeaders map[string]string
 	metaFilter    *metadata.ExcludeFilter
 
-	connMut sync.RWMutex
+	connMut    sync.RWMutex
+	topicCache syncmap.Map
 }
 
 // NewKafkaWriter returns a kafka writer.
@@ -186,6 +194,10 @@ func NewKafkaWriter(conf output.KafkaConfig, mgr bundle.NewManagement) (output.A
 		return nil, err
 	}
 
+	if k.admin, err = sarama.NewClusterAdmin(conf.Addresses, sarama.NewConfig()); err != nil {
+		return nil, err
+	}
+
 	if tout := conf.Timeout; len(tout) > 0 {
 		var err error
 		if k.timeout, err = time.ParseDuration(tout); err != nil {
@@ -296,6 +308,13 @@ func (k *kafkaWriter) buildUserDefinedHeaders(staticHeaders map[string]string) [
 
 //------------------------------------------------------------------------------
 
+// topicCache serves as a cache for topics that have been created during the
+// lifetime of this writer. This is to avoid unnecessary calls to the Kafka
+// cluster to create topics that already exist.
+//
+// The map is keyed by topic name and the value is a boolean indicating whether
+// the topic exists or not.
+
 // Connect attempts to establish a connection to a Kafka broker.
 func (k *kafkaWriter) Connect(ctx context.Context) error {
 	k.connMut.Lock()
@@ -365,6 +384,18 @@ func (k *kafkaWriter) WriteBatch(ctx context.Context, msg message.Batch) error {
 		if err != nil {
 			return fmt.Errorf("topic interpolation error: %w", err)
 		}
+		if k.conf.CustomTopicCreation.Enabled {
+			if k.conf.CustomTopicCreation.Partitions < 2 {
+				return fmt.Errorf("topic_partitions must be greater than one, got %v", k.conf.CustomTopicCreation.Partitions)
+			}
+			if k.conf.CustomTopicCreation.ReplicationFactor%2 == 0 {
+				return fmt.Errorf("topic_replication_factor must be an odd number, got %v", k.conf.CustomTopicCreation.ReplicationFactor)
+			}
+
+			if err = k.createTopic(topic); err != nil {
+				return fmt.Errorf("failed to create topic '%v': %w", topic, err)
+			}
+		}
 		nextMsg := &sarama.ProducerMessage{
 			Topic:    topic,
 			Value:    sarama.ByteEncoder(p.AsBytes()),
@@ -564,3 +595,46 @@ func (mur *murmur2) Sum32() uint32 {
 	mur.cached = &cached
 	return cached
 }
+
+//------------------------------------------------------------------------------
+
+// createTopic creates a topic in the Kafka cluster if it does not already
+// exist.
+//
+// If k.conf.PartitionsPerNewTopic is set to a value greater than 0, then the
+// topic will be created with that number of partitions.
+func (k *kafkaWriter) createTopic(topic string) (err error) {
+	var exists bool
+	if exists, err = k.checkIfTopicExists(topic); err != nil {
+		return err
+	}
+	if exists {
+		return nil
+	}
+
+	k.topicCache.Store(topic, false)
+
+	topicDetail := sarama.TopicDetail{
+		NumPartitions:     int32(k.conf.CustomTopicCreation.Partitions),
+		ReplicationFactor: int16(k.conf.CustomTopicCreation.ReplicationFactor),
+	}
+
+	err = k.admin.CreateTopic(topic, &topicDetail, false)
+	return err
+}
+
+// checkIfTopicExists checks if a topic exists in the Kafka cluster.
+func (k *kafkaWriter) checkIfTopicExists(topic string) (exists bool, err error) {
+	initialized, ok := k.topicCache.Load(topic)
+	if !ok || !initialized.(bool) {
+		var topics map[string]sarama.TopicDetail
+		if topics, err = k.admin.ListTopics(); err != nil {
+			return false, err
+		}
+		_, exists = topics[topic]
+		k.topicCache.Store(topic, exists)
+		return exists, nil
+	}
+
+	return initialized.(bool), nil
+}
diff --git a/website/docs/components/outputs/kafka.md b/website/docs/components/outputs/kafka.md
index 2eeee9f0e8..3bc6ff873f 100644
--- a/website/docs/components/outputs/kafka.md
+++ b/website/docs/components/outputs/kafka.md
@@ -76,6 +76,10 @@ output:
     key: ""
     partitioner: fnv1a_hash
     partition: ""
+    custom_topic_creation:
+      enabled: false
+      partitions: -1
+      replication_factor: -1
     compression: none
     static_headers: {}
     metadata:
@@ -437,6 +441,37 @@ This field supports [interpolation functions](/docs/configuration/interpolation#
 Type: `string`  
 Default: `""`  
 
+### `custom_topic_creation`
+
+If enabled, topics will be created with the specified number of partitions and replication factor if they do not already exist.
+
+
+Type: `object`  
+
+### `custom_topic_creation.enabled`
+
+Whether to enable custom topic creation.
+
+
+Type: `bool`  
+Default: `false`  
+
+### `custom_topic_creation.partitions`
+
+The number of partitions to create for new topics. Must be greater than 1
+
+
+Type: `int`  
+Default: `-1`  
+
+### `custom_topic_creation.replication_factor`
+
+The replication factor to use for new topics. Must be an odd number, and less then or equal to the number of brokers
+
+
+Type: `int`  
+Default: `-1`  
+
 ### `compression`
 
 The compression algorithm to use.


Test patch: