allowing custom action for elastic format (#796)

* allowing custom action for elastic format

Author: i3149

cmd/ktranslate/main.go

@@ -450,6 +450,9 @@ func applyFlags(cfg *ktranslate.Config) error {
 				cfg.OtelFormat.ClientKey = val
 			case "otel.root_ca":
 				cfg.OtelFormat.RootCA = val
+			// pkg/formats/elasticsearch
+			case "elastic.action":
+				cfg.ElasticFormat.Action = val
 			// pkg/formats/snmp
 			case "snmp.format.conf":
 				cfg.SnmpFormat.ConfigFile = val

config.go

@@ -35,6 +35,11 @@ type PrometheusFormatConfig struct {
 	FlowsNeeded          int
 }
 
+// ElasticFormatConfig is the config for the elastic format
+type ElasticFormatConfig struct {
+	Action string
+}
+
 // OtelFormatConfig is the config for the otel format
 type OtelFormatConfig struct {
 	Endpoint   string
@@ -291,6 +296,8 @@ type Config struct {
 	OtelFormat *OtelFormatConfig
 	// pkg/formats/snmp
 	SnmpFormat *SnmpFormatConfig
+	// pkg/formats/elasticsearch
+	ElasticFormat *ElasticFormatConfig
 
 	// pkg/sinks/prom
 	PrometheusSink *PrometheusSinkConfig
@@ -389,6 +396,9 @@ func DefaultConfig() *Config {
 			ClientCert: "",
 			RootCA:     "",
 		},
+		ElasticFormat: &ElasticFormatConfig{
+			Action: "index",
+		},
 		SnmpFormat: &SnmpFormatConfig{
 			ConfigFile: "",
 		},

pkg/formats/elasticsearch/elasticsearch.go

@@ -4,11 +4,13 @@ import (
 	"bufio"
 	"bytes"
 	"compress/gzip"
+	"flag"
 	"fmt"
 	"io"
 	"strings"
 	"time"
 
+	"github.com/kentik/ktranslate"
 	"github.com/kentik/ktranslate/pkg/formats/util"
 	"github.com/kentik/ktranslate/pkg/kt"
 	"github.com/kentik/ktranslate/pkg/rollup"
@@ -18,21 +20,31 @@ import (
 )
 
 const (
-	actionEntry = `{"index":{}}`
+	actionEntryFormat = `{"%s":{}}`
 )
 
+var (
+	actionEntrySet string
+)
+
+func init() {
+	flag.StringVar(&actionEntrySet, "elastic.action", "index", "Use this action when sending to elastic.")
+}
+
 var json = jsoniter.ConfigFastest
 
 type ElasticsearchFormat struct {
 	logger.ContextL
 	compression kt.Compression
 	useGzip     bool
+	action      string
 }
 
-func NewFormat(log logger.Underlying, compression kt.Compression) (*ElasticsearchFormat, error) {
+func NewFormat(log logger.Underlying, compression kt.Compression, cfg *ktranslate.ElasticFormatConfig) (*ElasticsearchFormat, error) {
 	ef := &ElasticsearchFormat{
 		ContextL:    logger.NewContextLFromUnderlying(logger.SContext{S: "elasticsearchFormat"}, log),
 		compression: compression,
+		action:      fmt.Sprintf(actionEntryFormat, cfg.Action),
 	}
 
 	switch compression {
@@ -44,6 +56,8 @@ func NewFormat(log logger.Underlying, compression kt.Compression) (*Elasticsearc
 		return nil, fmt.Errorf("Invalid compression (%s): format json only supports none|gzip", compression)
 	}
 
+	ef.Infof("Using action %s", ef.action)
+
 	return ef, nil
 }
 
@@ -62,7 +76,7 @@ func (f *ElasticsearchFormat) To(msgs []*kt.JCHF, serBuf []byte) (*kt.Output, er
 
 	esBulkData := []string{}
 	for _, m := range msgsNew {
-		data, err := serialize(m)
+		data, err := serialize(m, f.action)
 		if err != nil {
 			return nil, err
 		}
@@ -108,7 +122,7 @@ func (f *ElasticsearchFormat) From(raw *kt.Output) ([]map[string]interface{}, er
 	for sc.Scan() {
 		// check for ES action and ignore
 		v := sc.Text()
-		if v == actionEntry {
+		if v == f.action {
 			continue
 		}
 		msg := &kt.JCHF{}
@@ -131,7 +145,7 @@ func (f *ElasticsearchFormat) Rollup(rolls []rollup.Rollup) (*kt.Output, error)
 	// serialize rolls
 	esBulkData := []string{}
 	for _, m := range rolls {
-		data, err := serialize(m)
+		data, err := serialize(m, f.action)
 		if err != nil {
 			return nil, err
 		}
@@ -166,8 +180,8 @@ func (f *ElasticsearchFormat) handleSynth(in *kt.JCHF) map[string]interface{} {
 	return attr
 }
 
-func serialize(o interface{}) (string, error) {
-	s := actionEntry + "\n"
+func serialize(o interface{}, action string) (string, error) {
+	s := action + "\n"
 	data, err := json.Marshal(o)
 	if err != nil {
 		return "", err

pkg/formats/elasticsearch/elasticsearch_test.go

@@ -3,15 +3,19 @@ package elasticsearch
 import (
 	"testing"
 
+	"github.com/kentik/ktranslate"
+	"github.com/kentik/ktranslate/pkg/eggs/logger"
+	lt "github.com/kentik/ktranslate/pkg/eggs/logger/testing"
 	"github.com/kentik/ktranslate/pkg/kt"
 	"github.com/stretchr/testify/assert"
 )
 
 func TestSerializeElasticsearch(t *testing.T) {
 	serBuf := make([]byte, 0)
 	assert := assert.New(t)
+	l := lt.NewTestContextL(logger.NilContext, t).GetLogger().GetUnderlyingLogger()
 
-	f, err := NewFormat(nil, kt.CompressionNone)
+	f, err := NewFormat(l, kt.CompressionNone, &ktranslate.ElasticFormatConfig{Action: "index"})
 	assert.NoError(err)
 
 	res, err := f.To(kt.InputTesting, serBuf)
@@ -29,7 +33,8 @@ func TestSerializeElasticsearch(t *testing.T) {
 func TestSerializeElasticsearchGzip(t *testing.T) {
 	serBuf := make([]byte, 0)
 	assert := assert.New(t)
-	f, err := NewFormat(nil, kt.CompressionGzip)
+	l := lt.NewTestContextL(logger.NilContext, t).GetLogger().GetUnderlyingLogger()
+	f, err := NewFormat(l, kt.CompressionGzip, &ktranslate.ElasticFormatConfig{Action: "index"})
 	assert.NoError(err)
 	res, err := f.To(kt.InputTesting, serBuf)
 	assert.NoError(err)

pkg/formats/format.go

@@ -60,7 +60,7 @@ func NewFormat(ctx context.Context, format Format, log logger.Underlying, regist
 	case FORMAT_AVRO:
 		return avro.NewFormat(log, compression)
 	case FORMAT_ELASTICSEARCH:
-		return elasticsearch.NewFormat(log, compression)
+		return elasticsearch.NewFormat(log, compression, cfg.ElasticFormat)
 	case FORMAT_JSON:
 		return json.NewFormat(log, compression, false)
 	case FORMAT_NETFLOW: