[cfg] Allow specifying max request size. (#22)

artie-labs · Jan 5, 2024 · 42e52ff · 42e52ff
1 parent 475dc1e
commit 42e52ff
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 7 deletions.
diff --git a/config/config.go b/config/config.go
@@ -14,7 +14,8 @@ type Kafka struct {
 	BootstrapServers string `yaml:"bootstrapServers"`
 	TopicPrefix      string `yaml:"topicPrefix"`
 	AwsEnabled       bool   `yaml:"awsEnabled"`
-	PublishSize      int    `yaml:"publishSize"`
+	PublishSize      int    `yaml:"publishSize,omitempty"`
+	MaxRequestSize   int64  `yaml:"maxRequestSize,omitempty"`
 }
 
 func (k *Kafka) GenerateDefault() {

diff --git a/lib/kafkalib/kafka.go b/lib/kafkalib/kafka.go
@@ -3,17 +3,16 @@ package kafkalib
 import (
 	"context"
 	"crypto/tls"
-	"github.com/artie-labs/reader/config"
-	"github.com/artie-labs/reader/constants"
-	"github.com/artie-labs/reader/lib/logger"
 	awsCfg "github.com/aws/aws-sdk-go-v2/config"
 	"github.com/segmentio/kafka-go"
 	"github.com/segmentio/kafka-go/sasl/aws_msk_iam_v2"
 	"strings"
 	"time"
-)
 
-const maxKafkaItemBytes = 4000000
+	"github.com/artie-labs/reader/config"
+	"github.com/artie-labs/reader/constants"
+	"github.com/artie-labs/reader/lib/logger"
+)
 
 func FromContext(ctx context.Context) *kafka.Writer {
 	log := logger.FromContext(ctx)
@@ -42,13 +41,16 @@ func InjectIntoContext(ctx context.Context) context.Context {
 
 	writer := &kafka.Writer{
 		Addr:                   kafka.TCP(strings.Split(cfg.Kafka.BootstrapServers, ",")...),
-		BatchBytes:             maxKafkaItemBytes,
 		Compression:            kafka.Gzip,
 		Balancer:               &kafka.LeastBytes{},
 		WriteTimeout:           5 * time.Second,
 		AllowAutoTopicCreation: true,
 	}
 
+	if cfg.Kafka.MaxRequestSize > 0 {
+		writer.BatchBytes = cfg.Kafka.MaxRequestSize
+	}
+
 	if cfg.Kafka.AwsEnabled {
 		saslCfg, err := awsCfg.LoadDefaultConfig(ctx)
 		if err != nil {