From 2de9596ccc8e8f4ffbcc4e0ddbfd4a167df6b091 Mon Sep 17 00:00:00 2001
From: Cyril Tovena <cyril.tovena@gmail.com>
Date: Mon, 13 Apr 2020 13:11:49 -0400
Subject: [PATCH] Add relabeling to docker driver using `loki-relabel-config`
 log option.

Signed-off-by: Cyril Tovena <cyril.tovena@gmail.com>
---
 cmd/docker-driver/config.go      | 28 ++++++++++++++++--
 cmd/docker-driver/config_test.go | 51 ++++++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+), 2 deletions(-)
 create mode 100644 cmd/docker-driver/config_test.go

diff --git a/cmd/docker-driver/config.go b/cmd/docker-driver/config.go
index 67d93a421a77..2368cf25b09c 100644
--- a/cmd/docker-driver/config.go
+++ b/cmd/docker-driver/config.go
@@ -9,16 +9,20 @@ import (
 	"strings"
 	"time"
 
-	"github.com/cortexproject/cortex/pkg/util"
+	cortex_util "github.com/cortexproject/cortex/pkg/util"
 	"github.com/cortexproject/cortex/pkg/util/flagext"
 	"github.com/docker/docker/daemon/logger"
 	"github.com/docker/docker/daemon/logger/templates"
 	"github.com/prometheus/common/model"
+	"github.com/prometheus/prometheus/pkg/labels"
+	"github.com/prometheus/prometheus/pkg/relabel"
+	"gopkg.in/yaml.v2"
 
 	"github.com/grafana/loki/pkg/helpers"
 	"github.com/grafana/loki/pkg/logentry/stages"
 	"github.com/grafana/loki/pkg/promtail/client"
 	"github.com/grafana/loki/pkg/promtail/targets"
+	"github.com/grafana/loki/pkg/util"
 )
 
 const (
@@ -40,6 +44,7 @@ const (
 	cfgMaxRetriesKey     = "loki-retries"
 	cfgPipelineStagesKey = "loki-pipeline-stage-file"
 	cfgTenantIDKey       = "loki-tenant-id"
+	cfgRelabelKey        = "loki-relabel-config"
 
 	swarmServiceLabelKey = "com.docker.swarm.service.name"
 	swarmStackLabelKey   = "com.docker.stack.namespace"
@@ -61,7 +66,7 @@ var (
 	defaultClientConfig = client.Config{
 		BatchWait: 1 * time.Second,
 		BatchSize: 100 * 1024,
-		BackoffConfig: util.BackoffConfig{
+		BackoffConfig: cortex_util.BackoffConfig{
 			MinBackoff: 100 * time.Millisecond,
 			MaxBackoff: 10 * time.Second,
 			MaxRetries: 10,
@@ -101,6 +106,7 @@ func validateDriverOpt(loggerInfo logger.Info) error {
 		case cfgMaxRetriesKey:
 		case cfgPipelineStagesKey:
 		case cfgTenantIDKey:
+		case cfgRelabelKey:
 		case "labels":
 		case "env":
 		case "env-regex":
@@ -271,6 +277,15 @@ func parseConfig(logCtx logger.Info) (*config, error) {
 	}
 	labels[targets.FilenameLabel] = model.LabelValue(logCtx.LogPath)
 
+	// Process relabel configs.
+	if relabelString, ok := logCtx.Config[cfgRelabelKey]; ok && relabelString != "" {
+		relabeled, err := relabelConfig(relabelString, labels)
+		if err != nil {
+			return nil, fmt.Errorf("error applying relabel config: %s err:%s", relabelString, err)
+		}
+		labels = relabeled
+	}
+
 	// parse pipeline stages
 	var pipeline PipelineConfig
 	pipelineFile, ok := logCtx.Config[cfgPipelineStagesKey]
@@ -321,3 +336,12 @@ func parseInt(key string, logCtx logger.Info, set func(i int)) error {
 	}
 	return nil
 }
+
+func relabelConfig(config string, lbs model.LabelSet) (model.LabelSet, error) {
+	relabelConfig := make([]*relabel.Config, 0)
+	if err := yaml.UnmarshalStrict([]byte(config), &relabelConfig); err != nil {
+		return nil, err
+	}
+	relabed := relabel.Process(labels.FromMap(util.ModelLabelSetToMap(lbs)), relabelConfig...)
+	return model.LabelSet(util.LabelsToMetric(relabed)), nil
+}
diff --git a/cmd/docker-driver/config_test.go b/cmd/docker-driver/config_test.go
new file mode 100644
index 000000000000..966600aa29fb
--- /dev/null
+++ b/cmd/docker-driver/config_test.go
@@ -0,0 +1,51 @@
+package main
+
+import (
+	"testing"
+
+	"github.com/prometheus/common/model"
+	"github.com/stretchr/testify/require"
+)
+
+var jobRename = `
+- regex: (.*)
+  source_labels: [swarm_stack]
+  target_label: job
+- regex: ^swarm_stack$
+  action: labeldrop`
+
+func Test_relabelConfig(t *testing.T) {
+
+	tests := []struct {
+		name    string
+		config  string
+		in      model.LabelSet
+		out     model.LabelSet
+		wantErr bool
+	}{
+		{
+			"config err",
+			"foo",
+			nil,
+			nil,
+			true,
+		},
+		{
+			"config err",
+			jobRename,
+			model.LabelSet{"swarm_stack": "foo", "bar": "buzz"},
+			model.LabelSet{"job": "foo", "bar": "buzz"},
+			false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := relabelConfig(tt.config, tt.in)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("relabelConfig() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			require.True(t, got.Equal(tt.out))
+		})
+	}
+}