From 526dc4bbe7d60b7f7a5ed753754e2a5fd0009215 Mon Sep 17 00:00:00 2001
From: Cyril Tovena <cyril.tovena@gmail.com>
Date: Tue, 25 May 2021 16:00:42 +0200
Subject: [PATCH] Fixes parser labels hint for grouping. (#3763)

* Fixes parser labels hint for grouping.

This PR include those now.

Fixes #3707

Signed-off-by: Cyril Tovena <cyril.tovena@gmail.com>

* Fixes array allocations.

Signed-off-by: Cyril Tovena <cyril.tovena@gmail.com>
---
 pkg/logql/log/parser_hints.go      | 51 +++++++++++++++---------------
 pkg/logql/log/parser_hints_test.go |  7 ++++
 2 files changed, 32 insertions(+), 26 deletions(-)

diff --git a/pkg/logql/log/parser_hints.go b/pkg/logql/log/parser_hints.go
index 55e85994106f..5a116c76bb06 100644
--- a/pkg/logql/log/parser_hints.go
+++ b/pkg/logql/log/parser_hints.go
@@ -63,32 +63,14 @@ func (p *parserHint) NoLabels() bool {
 
 // newParserHint creates a new parser hint using the list of labels that are seen and required in a query.
 func newParserHint(requiredLabelNames, groups []string, without, noLabels bool, metricLabelName string) *parserHint {
-	// If a parsed label collides with a stream label we add the `_extracted` suffix to it, however hints
-	// are used by the parsers before we know they will collide with a stream label and hence before the
-	// _extracted suffix is added. Therefore we must strip the _extracted suffix from any required labels
-	// that were parsed from somewhere in the query, say in a filter or an aggregation clause.
-	// Because it's possible for a valid json or logfmt key to already end with _extracted, we'll just
-	// leave the existing entry ending with _extracted but also add a version with the suffix removed.
-	extractedLabels := []string{}
-	for _, l := range requiredLabelNames {
-		if strings.HasSuffix(l, "_extracted") {
-			extractedLabels = append(extractedLabels, strings.TrimSuffix(l, "_extracted"))
-		}
-	}
-	if len(extractedLabels) > 0 {
-		requiredLabelNames = append(requiredLabelNames, extractedLabels...)
-	}
-
-	if len(groups) > 0 {
-		requiredLabelNames = append(requiredLabelNames, groups...)
-	}
-	if metricLabelName != "" {
-		requiredLabelNames = append(requiredLabelNames, metricLabelName)
-	}
-	requiredLabelNames = uniqueString(requiredLabelNames)
+	hints := make([]string, 0, 2*(len(requiredLabelNames)+len(groups)+1))
+	hints = appendLabelHints(hints, requiredLabelNames...)
+	hints = appendLabelHints(hints, groups...)
+	hints = appendLabelHints(hints, metricLabelName)
+	hints = uniqueString(hints)
 	if noLabels {
-		if len(requiredLabelNames) > 0 {
-			return &parserHint{requiredLabels: requiredLabelNames}
+		if len(hints) > 0 {
+			return &parserHint{requiredLabels: hints}
 		}
 		return &parserHint{noLabels: true}
 	}
@@ -98,5 +80,22 @@ func newParserHint(requiredLabelNames, groups []string, without, noLabels bool,
 	if without || len(groups) == 0 {
 		return noParserHints
 	}
-	return &parserHint{requiredLabels: requiredLabelNames}
+	return &parserHint{requiredLabels: hints}
+}
+
+// appendLabelHints Appends the label to the list of hints with and without the duplicate suffix.
+// If a parsed label collides with a stream label we add the `_extracted` suffix to it, however hints
+// are used by the parsers before we know they will collide with a stream label and hence before the
+// _extracted suffix is added. Therefore we must strip the _extracted suffix from any required labels
+// that were parsed from somewhere in the query, say in a filter or an aggregation clause.
+// Because it's possible for a valid json or logfmt key to already end with _extracted, we'll just
+// leave the existing entry ending with _extracted but also add a version with the suffix removed.
+func appendLabelHints(dst []string, src ...string) []string {
+	for _, l := range src {
+		dst = append(dst, l)
+		if strings.HasSuffix(l, duplicateSuffix) {
+			dst = append(dst, strings.TrimSuffix(l, duplicateSuffix))
+		}
+	}
+	return dst
 }
diff --git a/pkg/logql/log/parser_hints_test.go b/pkg/logql/log/parser_hints_test.go
index fc639461ca82..6a1ada86ce0a 100644
--- a/pkg/logql/log/parser_hints_test.go
+++ b/pkg/logql/log/parser_hints_test.go
@@ -198,6 +198,13 @@ func Test_ParserHints(t *testing.T) {
 			1.0,
 			`{cluster_extracted="us-east-west"}`,
 		},
+		{
+			`sum by (cluster_extracted)(count_over_time({app="nginx"} | unpack[1m]))`,
+			packedLine,
+			true,
+			1.0,
+			`{cluster_extracted="us-east-west"}`,
+		},
 		{
 			`sum(rate({app="nginx"} | unpack | nonexistant_field="foo" [1m]))`,
 			packedLine,