Implement more format on Loki writer

- Implement same formats as stdout writer (printf, json, fields)
- Allow reordering json keys
- Refactor and add more tests

Author: jotak

pkg/api/write_loki.go

@@ -44,6 +44,8 @@ type WriteLoki struct {
 	// scales of '1ms' (one millisecond) or just '1' (one nanosecond)
 	// Default value is '1s'
 	TimestampScale string `yaml:"timestampScale,omitempty" json:"timestampScale,omitempty" doc:"timestamp units scale (e.g. for UNIX = 1s)"`
+	Format         string `yaml:"format,omitempty" json:"format,omitempty" doc:"the format of each line: printf (writes using golang's default map printing), fields (writes one key and value field per line) or json (default)"`
+	Reorder        bool   `yaml:"reorder,omitempty" json:"reorder,omitempty" doc:"reorder json map keys"`
 }
 
 func (w *WriteLoki) SetDefaults() {
@@ -71,6 +73,9 @@ func (w *WriteLoki) SetDefaults() {
 	if w.TimestampScale == "" {
 		w.TimestampScale = "1s"
 	}
+	if w.Format == "" {
+		w.Format = "json"
+	}
 }
 
 func (w *WriteLoki) Validate() error {

pkg/pipeline/write/write_loki.go

@@ -30,16 +30,13 @@ import (
 	"github.com/netobserv/flowlogs-pipeline/pkg/utils"
 
 	logAdapter "github.com/go-kit/kit/log/logrus"
-	jsonIter "github.com/json-iterator/go"
 	"github.com/netobserv/loki-client-go/loki"
 	"github.com/netobserv/loki-client-go/pkg/backoff"
 	"github.com/netobserv/loki-client-go/pkg/urlutil"
 	"github.com/prometheus/common/model"
 	"github.com/sirupsen/logrus"
 )
 
-var jsonEncodingConfig = jsonIter.Config{}.Froze()
-
 var (
 	keyReplacer = strings.NewReplacer("/", "_", ".", "_", "-", "_")
 )
@@ -60,6 +57,7 @@ type Loki struct {
 	timeNow        func() time.Time
 	exitChan       <-chan struct{}
 	metrics        *metrics
+	formatter      func(config.GenericMap) string
 }
 
 func buildLokiConfig(c *api.WriteLoki) (loki.Config, error) {
@@ -124,13 +122,10 @@ func (l *Loki) ProcessRecord(in config.GenericMap) error {
 		delete(out, label)
 	}
 
-	js, err := jsonEncodingConfig.Marshal(out)
-	if err != nil {
-		return err
-	}
+	logline := l.formatter(out)
 
 	timestamp := l.extractTimestamp(out)
-	err = l.client.Handle(labels, timestamp, string(js))
+	err := l.client.Handle(labels, timestamp, logline)
 	if err == nil {
 		l.metrics.recordsWritten.Inc()
 	}
@@ -255,6 +250,7 @@ func NewWriteLoki(opMetrics *operational.Metrics, params config.StageParam) (*Lo
 		}
 	}
 
+	f := formatter(lokiConfigIn.Format, lokiConfigIn.Reorder)
 	l := &Loki{
 		lokiConfig:     lokiConfig,
 		apiConfig:      lokiConfigIn,
@@ -264,6 +260,7 @@ func NewWriteLoki(opMetrics *operational.Metrics, params config.StageParam) (*Lo
 		timeNow:        time.Now,
 		exitChan:       pUtils.ExitChannel(),
 		metrics:        newMetrics(opMetrics, params.Name),
+		formatter:      f,
 	}
 
 	return l, nil

pkg/pipeline/write/write_loki_test.go

@@ -41,19 +41,24 @@ const timeout = 5 * time.Second
 
 type fakeEmitter struct {
 	mock.Mock
+	reorderJSON bool
 }
 
 func (f *fakeEmitter) Handle(labels model.LabelSet, timestamp time.Time, record string) error {
-	// sort alphabetically records just for simplifying testing verification with JSON strings
-	recordMap := map[string]interface{}{}
-	if err := json.Unmarshal([]byte(record), &recordMap); err != nil {
-		panic("expected JSON: " + err.Error())
-	}
-	recordBytes, err := json.Marshal(recordMap)
-	if err != nil {
-		panic("error unmarshaling: " + err.Error())
+	if f.reorderJSON {
+		// sort alphabetically records just for simplifying testing verification with JSON strings
+		recordMap := map[string]interface{}{}
+		if err := json.Unmarshal([]byte(record), &recordMap); err != nil {
+			panic("expected JSON: " + err.Error())
+		}
+		recordBytes, err := json.Marshal(recordMap)
+		if err != nil {
+			panic("error unmarshaling: " + err.Error())
+		}
+		a := f.Mock.Called(labels, timestamp, string(recordBytes))
+		return a.Error(0)
 	}
-	a := f.Mock.Called(labels, timestamp, string(recordBytes))
+	a := f.Mock.Called(labels, timestamp, record)
 	return a.Error(0)
 }
 
@@ -122,32 +127,17 @@ parameters:
 }
 
 func TestLoki_ProcessRecord(t *testing.T) {
-	var yamlConfig = `
-log-level: debug
-pipeline:
-  - name: write1
-parameters:
-  - name: write1
-    write:
-      type: loki
-      loki:
-        url: http://loki:3100/
-        timestampLabel: ts
-        ignoreList:
-        - ignored
-        staticLabels:
-          static: label
-        labels:
-          - foo
-          - bar
-`
-	v, cfg := test.InitConfig(t, yamlConfig)
-	require.NotNil(t, v)
-
-	loki, err := NewWriteLoki(operational.NewMetrics(&config.MetricsSettings{}), cfg.Parameters[0])
+	params := api.WriteLoki{
+		URL:            "http://loki:3100/",
+		TimestampLabel: "ts",
+		IgnoreList:     []string{"ignored"},
+		StaticLabels:   model.LabelSet{"static": "label"},
+		Labels:         []string{"foo", "bar"},
+	}
+	loki, err := NewWriteLoki(operational.NewMetrics(&config.MetricsSettings{}), config.StageParam{Write: &config.Write{Loki: &params}})
 	require.NoError(t, err)
 
-	fe := fakeEmitter{}
+	fe := fakeEmitter{reorderJSON: true}
 	fe.On("Handle", mock.Anything, mock.Anything, mock.Anything).Return(nil)
 	loki.client = &fe
 
@@ -171,6 +161,49 @@ parameters:
 	}, time.Unix(124567, 0), `{"other":"val","ts":124567,"value":5678}`)
 }
 
+func TestLoki_ProcessRecordOrdered(t *testing.T) {
+	params := api.WriteLoki{
+		URL:            "http://loki:3100/",
+		TimestampLabel: "ts",
+		Reorder:        true,
+	}
+	loki, err := NewWriteLoki(operational.NewMetrics(&config.MetricsSettings{}), config.StageParam{Write: &config.Write{Loki: &params}})
+	require.NoError(t, err)
+
+	fe := fakeEmitter{reorderJSON: false}
+	fe.On("Handle", mock.Anything, mock.Anything, mock.Anything).Return(nil)
+	loki.client = &fe
+
+	require.NoError(t, loki.ProcessRecord(map[string]interface{}{"ts": 123456, "c": "c", "e": "e", "a": "a", "d": "d", "b": "b"}))
+
+	fe.AssertCalled(t, "Handle", model.LabelSet{}, time.Unix(123456, 0), `{"a":"a","b":"b","c":"c","d":"d","e":"e","ts":123456}`)
+}
+
+func TestLoki_ProcessRecordGoFormat(t *testing.T) {
+	params := api.WriteLoki{
+		URL:            "http://loki:3100/",
+		TimestampLabel: "ts",
+		Labels:         []string{"foo"},
+		Format:         "printf",
+	}
+	loki, err := NewWriteLoki(operational.NewMetrics(&config.MetricsSettings{}), config.StageParam{Write: &config.Write{Loki: &params}})
+	require.NoError(t, err)
+
+	fe := fakeEmitter{}
+	fe.On("Handle", mock.Anything, mock.Anything, mock.Anything).Return(nil)
+	loki.client = &fe
+
+	// WHEN it processes input records
+	require.NoError(t, loki.ProcessRecord(map[string]interface{}{
+		"ts": 123456, "foo": "fooLabel", "bar": "barLabel", "value": 1234}))
+	require.NoError(t, loki.ProcessRecord(map[string]interface{}{
+		"ts": 124567, "foo": "fooLabel2", "bar": "barLabel2", "value": 5678, "other": "val"}))
+
+	// THEN it forwards the records extracting the timestamp and labels from the configuration
+	fe.AssertCalled(t, "Handle", model.LabelSet{"foo": "fooLabel"}, time.Unix(123456, 0), `map[bar:barLabel ts:123456 value:1234]`)
+	fe.AssertCalled(t, "Handle", model.LabelSet{"foo": "fooLabel2"}, time.Unix(124567, 0), `map[bar:barLabel2 other:val ts:124567 value:5678]`)
+}
+
 func TestTimestampScale(t *testing.T) {
 	// verifies that the unix residual time (below 1-second precision) is properly
 	// incorporated into the timestamp whichever scale it is
@@ -262,26 +295,12 @@ func TestTimestampExtraction_LocalTime(t *testing.T) {
 // Tests that labels are sanitized before being sent to loki.
 // Labels that are invalid even if sanitized are ignored
 func TestSanitizedLabels(t *testing.T) {
-	var yamlConfig = `
-log-level: debug
-pipeline:
-  - name: write1
-parameters:
-  - name: write1
-    write:
-      type: loki
-      loki:
-        url: http://loki:3100/
-        labels:
-          - "fo.o"
-          - "ba-r"
-          - "ba/z"
-          - "ignored?"
-`
-	v, cfg := test.InitConfig(t, yamlConfig)
-	require.NotNil(t, v)
-
-	loki, err := NewWriteLoki(operational.NewMetrics(&config.MetricsSettings{}), cfg.Parameters[0])
+	params := api.WriteLoki{
+		URL:            "http://loki:3100/",
+		TimestampLabel: "ts",
+		Labels:         []string{"fo.o", "ba-r", "ba/z", "ignored?"},
+	}
+	loki, err := NewWriteLoki(operational.NewMetrics(&config.MetricsSettings{}), config.StageParam{Write: &config.Write{Loki: &params}})
 	require.NoError(t, err)
 
 	fe := fakeEmitter{}

pkg/pipeline/write/write_stdout.go

@@ -18,50 +18,61 @@
 package write
 
 import (
-	"encoding/json"
 	"fmt"
-	"os"
 	"sort"
+	"strings"
 	"text/tabwriter"
 	"time"
 
+	jsonIter "github.com/json-iterator/go"
 	"github.com/netobserv/flowlogs-pipeline/pkg/config"
 	"github.com/sirupsen/logrus"
 )
 
 type writeStdout struct {
-	format string
+	formatter func(config.GenericMap) string
 }
 
 // Write writes a flow before being stored
 func (t *writeStdout) Write(v config.GenericMap) {
 	logrus.Tracef("entering writeStdout Write")
-	if t.format == "json" {
-		txt, _ := json.Marshal(v)
-		fmt.Println(string(txt))
-	} else if t.format == "fields" {
-		var order sort.StringSlice
-		for fieldName := range v {
-			order = append(order, fieldName)
+	fmt.Println(t.formatter(v))
+}
+
+func formatter(format string, reorder bool) func(config.GenericMap) string {
+	if format == "json" {
+		jconf := jsonIter.Config{
+			SortMapKeys: reorder,
+		}.Froze()
+		return func(v config.GenericMap) string {
+			b, _ := jconf.Marshal(v)
+			return string(b)
 		}
-		order.Sort()
-		w := tabwriter.NewWriter(os.Stdout, 0, 0, 1, ' ', 0)
-		fmt.Fprintf(w, "\n\nFlow record at %s:\n", time.Now().Format(time.StampMilli))
-		for _, field := range order {
-			fmt.Fprintf(w, "%v\t=\t%v\n", field, v[field])
+	} else if format == "fields" {
+		return func(v config.GenericMap) string {
+			var sb strings.Builder
+			var order sort.StringSlice
+			for fieldName := range v {
+				order = append(order, fieldName)
+			}
+			order.Sort()
+			w := tabwriter.NewWriter(&sb, 0, 0, 1, ' ', 0)
+			fmt.Fprintf(w, "\n\nFlow record at %s:\n", time.Now().Format(time.StampMilli))
+			for _, field := range order {
+				fmt.Fprintf(w, "%v\t=\t%v\n", field, v[field])
+			}
+			w.Flush()
+			return sb.String()
 		}
-		w.Flush()
-	} else {
-		fmt.Printf("%s: %v\n", time.Now().Format(time.StampMilli), v)
+	}
+	return func(v config.GenericMap) string {
+		return fmt.Sprintf("%v", v)
 	}
 }
 
 // NewWriteStdout create a new write
 func NewWriteStdout(params config.StageParam) (Writer, error) {
 	logrus.Debugf("entering NewWriteStdout")
-	writeStdout := &writeStdout{}
-	if params.Write != nil && params.Write.Stdout != nil {
-		writeStdout.format = params.Write.Stdout.Format
-	}
-	return writeStdout, nil
+	f := formatter(params.Write.Stdout.Format, false)
+	return &writeStdout{formatter: f}, nil
 }

pkg/pipeline/write/write_stdout_test.go

@@ -20,17 +20,13 @@ package write
 import (
 	"testing"
 
+	"github.com/netobserv/flowlogs-pipeline/pkg/api"
 	"github.com/netobserv/flowlogs-pipeline/pkg/config"
 	"github.com/stretchr/testify/require"
 )
 
-func Test_WriteStdout(_ *testing.T) {
-	ws := writeStdout{}
-	ws.Write(config.GenericMap{"key": "test"})
-}
-
 func Test_NewWriteStdout(t *testing.T) {
-	writer, err := NewWriteStdout(config.StageParam{})
-	require.Nil(t, err)
-	require.Equal(t, writer, &writeStdout{})
+	writer, err := NewWriteStdout(config.StageParam{Write: &config.Write{Stdout: &api.WriteStdout{}}})
+	require.NoError(t, err)
+	writer.Write(config.GenericMap{"key": "test"})
 }