From fb140eda33ccbe8fdea366c90cbcd0fef145b121 Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@gmail.com>
Date: Thu, 23 May 2019 16:14:01 +0300
Subject: [PATCH] app/vmselect/promql: add `label_transform(q, label, regexp,
 replacement)` function for replacing all the occurences of regexp with
 replacement in the given label for q

---
 app/vmselect/promql/exec_test.go    | 33 ++++++++++++++++++++++++++
 app/vmselect/promql/regexp_cache.go |  9 +++++---
 app/vmselect/promql/transform.go    | 36 +++++++++++++++++++++++++----
 3 files changed, 71 insertions(+), 7 deletions(-)

diff --git a/app/vmselect/promql/exec_test.go b/app/vmselect/promql/exec_test.go
index e683fcf02..24265ee71 100644
--- a/app/vmselect/promql/exec_test.go
+++ b/app/vmselect/promql/exec_test.go
@@ -1266,6 +1266,34 @@ func TestExecSuccess(t *testing.T) {
 		resultExpected := []netstorage.Result{r}
 		f(q, resultExpected)
 	})
+	t.Run(`label_transform(mismatch)`, func(t *testing.T) {
+		t.Parallel()
+		q := `label_transform(time(), "__name__", "foobar", "xx")`
+		r := netstorage.Result{
+			MetricName: metricNameExpected,
+			Values:     []float64{1000, 1200, 1400, 1600, 1800, 2000},
+			Timestamps: timestampsExpected,
+		}
+		resultExpected := []netstorage.Result{r}
+		f(q, resultExpected)
+	})
+	t.Run(`label_transform(match)`, func(t *testing.T) {
+		t.Parallel()
+		q := `label_transform(
+			label_set(time(), "foo", "a.bar.baz"),
+			"foo", "\\.", "-")`
+		r := netstorage.Result{
+			MetricName: metricNameExpected,
+			Values:     []float64{1000, 1200, 1400, 1600, 1800, 2000},
+			Timestamps: timestampsExpected,
+		}
+		r.MetricName.Tags = []storage.Tag{{
+			Key:   []byte("foo"),
+			Value: []byte("a-bar-baz"),
+		}}
+		resultExpected := []netstorage.Result{r}
+		f(q, resultExpected)
+	})
 	t.Run(`label_replace(mismatch)`, func(t *testing.T) {
 		t.Parallel()
 		q := `label_replace(time(), "__name__", "x${1}y", "foo", ".+")`
@@ -3461,6 +3489,7 @@ func TestExecError(t *testing.T) {
 	f(`hour(1,2)`)
 	f(`label_join()`)
 	f(`label_replace(1)`)
+	f(`label_transform(1)`)
 	f(`label_set()`)
 	f(`label_set(1, "foo")`)
 	f(`label_del()`)
@@ -3535,6 +3564,10 @@ func TestExecError(t *testing.T) {
 	f(`label_replace(1, "foo", "bar", 4, 5)`)
 	f(`label_replace(1, "foo", "bar", "baz", 5)`)
 	f(`label_replace(1, "foo", "bar", "baz", "invalid(regexp")`)
+	f(`label_transform(1, 2, 3, 4)`)
+	f(`label_transform(1, "foo", 3, 4)`)
+	f(`label_transform(1, "foo", "bar", 4)`)
+	f(`label_transform(1, "foo", "invalid(regexp", "baz`)
 
 	// Duplicate timeseries
 	f(`(label_set(1, "foo", "bar") or label_set(2, "foo", "baz"))
diff --git a/app/vmselect/promql/regexp_cache.go b/app/vmselect/promql/regexp_cache.go
index 51b16ba39..4d0ff4d93 100644
--- a/app/vmselect/promql/regexp_cache.go
+++ b/app/vmselect/promql/regexp_cache.go
@@ -1,7 +1,6 @@
 package promql
 
 import (
-	"fmt"
 	"regexp"
 	"sync"
 	"sync/atomic"
@@ -10,12 +9,16 @@ import (
 )
 
 func compileRegexpAnchored(re string) (*regexp.Regexp, error) {
+	reAnchored := "^(?:" + re + ")$"
+	return compileRegexp(reAnchored)
+}
+
+func compileRegexp(re string) (*regexp.Regexp, error) {
 	rcv := regexpCacheV.Get(re)
 	if rcv != nil {
 		return rcv.r, rcv.err
 	}
-	regexAnchored := fmt.Sprintf("^(?:%s)$", re)
-	r, err := regexp.Compile(regexAnchored)
+	r, err := regexp.Compile(re)
 	rcv = &regexpCacheValue{
 		r:   r,
 		err: err,
diff --git a/app/vmselect/promql/transform.go b/app/vmselect/promql/transform.go
index 8608e7394..21aec90e9 100644
--- a/app/vmselect/promql/transform.go
+++ b/app/vmselect/promql/transform.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"math"
 	"math/rand"
+	"regexp"
 	"sort"
 	"strconv"
 	"strings"
@@ -61,6 +62,7 @@ var transformFuncs = map[string]transformFunc{
 	"label_keep":         transformLabelKeep,
 	"label_copy":         transformLabelCopy,
 	"label_move":         transformLabelMove,
+	"label_transform":    transformLabelTransform,
 	"union":              transformUnion,
 	"":                   transformUnion, // empty func is a synonim to union
 	"keep_last_value":    transformKeepLastValue,
@@ -816,6 +818,31 @@ func transformLabelJoin(tfa *transformFuncArg) ([]*timeseries, error) {
 	return rvs, nil
 }
 
+func transformLabelTransform(tfa *transformFuncArg) ([]*timeseries, error) {
+	args := tfa.args
+	if err := expectTransformArgsNum(args, 4); err != nil {
+		return nil, err
+	}
+	label, err := getString(args[1], 1)
+	if err != nil {
+		return nil, err
+	}
+	regex, err := getString(args[2], 2)
+	if err != nil {
+		return nil, err
+	}
+	replacement, err := getString(args[3], 3)
+	if err != nil {
+		return nil, err
+	}
+
+	r, err := compileRegexp(regex)
+	if err != nil {
+		return nil, fmt.Errorf(`cannot compile regex %q: %s`, regex, err)
+	}
+	return labelReplace(args[0], label, r, label, replacement)
+}
+
 func transformLabelReplace(tfa *transformFuncArg) ([]*timeseries, error) {
 	args := tfa.args
 	if err := expectTransformArgsNum(args, 5); err != nil {
@@ -842,11 +869,12 @@ func transformLabelReplace(tfa *transformFuncArg) ([]*timeseries, error) {
 	if err != nil {
 		return nil, fmt.Errorf(`cannot compile regex %q: %s`, regex, err)
 	}
+	return labelReplace(args[0], srcLabel, r, dstLabel, replacement)
+}
 
+func labelReplace(tss []*timeseries, srcLabel string, r *regexp.Regexp, dstLabel, replacement string) ([]*timeseries, error) {
 	replacementBytes := []byte(replacement)
-
-	rvs := args[0]
-	for _, ts := range rvs {
+	for _, ts := range tss {
 		mn := &ts.MetricName
 		dstValue := getDstValue(mn, dstLabel)
 		srcValue := mn.GetTagValue(srcLabel)
@@ -856,7 +884,7 @@ func transformLabelReplace(tfa *transformFuncArg) ([]*timeseries, error) {
 			mn.RemoveTag(dstLabel)
 		}
 	}
-	return rvs, nil
+	return tss, nil
 }
 
 func transformLn(v float64) float64 {