Plugins: Fix support for adhoc filters with raw queries in InfluxDB (#101966)

Plugins: Fix support for adhoc filters with raw queries in InfluxDB

Fixes #101635.

pkg/tsdb/influxdb/models/model_parser.go

@@ -36,6 +36,11 @@ func QueryParse(query backend.DataQuery, logger log.Logger) (*Query, error) {
 	measurement := model.Get("measurement").MustString("")
 	resultFormat := model.Get("resultFormat").MustString("")
 
+	adhocFilters, err := parseAdhocFilters(model.Get("adhocFilters").MustArray())
+	if err != nil {
+		return nil, errors.Join(ErrInvalidQuery, err)
+	}
+
 	tags, err := parseTags(model)
 	if err != nil {
 		return nil, errors.Join(ErrInvalidQuery, err)
@@ -84,6 +89,7 @@ func QueryParse(query backend.DataQuery, logger log.Logger) (*Query, error) {
 		OrderByTime:  orderByTime,
 		ResultFormat: resultFormat,
 		Statement:    statement,
+		AdhocFilters: adhocFilters,
 	}, nil
 }
 
@@ -111,6 +117,38 @@ func parseSelects(model *simplejson.Json) ([]*Select, error) {
 	return result, nil
 }
 
+func parseAdhocFilters(adhocFilters []any) ([]*Tag, error) {
+	result := make([]*Tag, 0, len(adhocFilters))
+	for _, t := range adhocFilters {
+		tagJson := simplejson.NewFromAny(t)
+		tag := &Tag{}
+		var err error
+
+		tag.Key, err = tagJson.Get("key").String()
+		if err != nil {
+			return nil, err
+		}
+
+		tag.Value, err = tagJson.Get("value").String()
+		if err != nil {
+			return nil, err
+		}
+
+		operator, err := tagJson.Get("operator").String()
+		if err == nil {
+			tag.Operator = operator
+		}
+
+		condition, err := tagJson.Get("condition").String()
+		if err == nil {
+			tag.Condition = condition
+		}
+
+		result = append(result, tag)
+	}
+	return result, nil
+}
+
 func parseTags(model *simplejson.Json) ([]*Tag, error) {
 	tags := model.Get("tags").MustArray()
 	result := make([]*Tag, 0, len(tags))

pkg/tsdb/influxdb/models/models.go

@@ -10,6 +10,7 @@ type Query struct {
 	Measurement  string
 	Policy       string
 	Tags         []*Tag
+	AdhocFilters []*Tag
 	GroupBy      []*QueryPart
 	Selects      []*Select
 	RawQuery     string

pkg/tsdb/influxdb/models/query.go

@@ -6,6 +6,7 @@ import (
 	"strconv"
 	"strings"
 	"time"
+	"unicode/utf8"
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
 	"github.com/grafana/grafana-plugin-sdk-go/backend/gtime"
@@ -15,12 +16,62 @@ var (
 	regexpOperatorPattern           = regexp.MustCompile(`^\/.*\/$`)
 	regexpMeasurementPattern        = regexp.MustCompile(`^\/.*\/$`)
 	regexMatcherWithStartEndPattern = regexp.MustCompile(`^/\^(.*)\$/$`)
+	regexpRawQueryWhere             = regexp.MustCompile(`(?i)WHERE`)
 )
 
+func (query *Query) getRawQueryWithAdhocFilters() string {
+	if len(query.AdhocFilters) == 0 {
+		return query.RawQuery
+	}
+
+	// If there's no where to be found (get it?), just append.
+	whereIndexes := regexpRawQueryWhere.FindAllStringIndex(query.RawQuery, -1)
+	if len(whereIndexes) == 0 {
+		query.RawQuery += " WHERE "
+		query.RawQuery += strings.Join(query.renderAdhocFilters(), " ")
+		return query.RawQuery
+	}
+
+	// Walk through raw query and determine if the 'WHERE' strings
+	// we found above are quoted or not. We'll insert adhoc filters
+	// after the first unquoted 'WHERE' string we find. Influxql supports
+	// subqueries, so valid queries can have multiple where clauses,
+	// but this code does not attempt to support that.
+	byteIndex := 0
+	insideQuotes := false
+	var previousRuneValue rune
+	whereIndex, whereIndexes := whereIndexes[0], whereIndexes[1:]
+	for _, runeValue := range query.RawQuery {
+		if previousRuneValue != '\\' && (runeValue == '"' || runeValue == '\'') {
+			insideQuotes = !insideQuotes
+		}
+		previousRuneValue = runeValue
+
+		if byteIndex == whereIndex[0] {
+			if !insideQuotes {
+				alteredQuery := query.RawQuery[:whereIndex[1]]
+				alteredQuery += " "
+				alteredQuery += strings.Join(query.renderAdhocFilters(), " ")
+				alteredQuery += " AND"
+				alteredQuery += query.RawQuery[whereIndex[1]:]
+				return alteredQuery
+			}
+			if len(whereIndexes) == 0 {
+				query.RawQuery += " WHERE "
+				query.RawQuery += strings.Join(query.renderAdhocFilters(), " ")
+				return query.RawQuery
+			}
+			whereIndex, whereIndexes = whereIndexes[0], whereIndexes[1:]
+		}
+		byteIndex += utf8.RuneLen(runeValue)
+	}
+	return query.RawQuery
+}
+
 func (query *Query) Build(queryContext *backend.QueryDataRequest) (string, error) {
 	var res string
 	if query.UseRawQuery && query.RawQuery != "" {
-		res = query.RawQuery
+		res = query.getRawQueryWithAdhocFilters()
 	} else {
 		res = query.renderSelectors(queryContext)
 		res += query.renderMeasurement()
@@ -44,9 +95,13 @@ func (query *Query) Build(queryContext *backend.QueryDataRequest) (string, error
 	return res, nil
 }
 
-func (query *Query) renderTags() []string {
-	res := make([]string, 0, len(query.Tags))
-	for i, tag := range query.Tags {
+func (query *Query) renderAdhocFilters() []string {
+	return renderTags(query.AdhocFilters)
+}
+
+func renderTags(tags []*Tag) []string {
+	res := make([]string, 0, len(tags))
+	for i, tag := range tags {
 		str := ""
 
 		if i > 0 {
@@ -131,10 +186,13 @@ func (query *Query) renderTags() []string {
 
 		res = append(res, fmt.Sprintf(`%s%s %s %s`, str, escapedKey, tag.Operator, textValue))
 	}
-
 	return res
 }
 
+func (query *Query) renderTags() []string {
+	return renderTags(query.Tags)
+}
+
 func (query *Query) renderTimeFilter(queryContext *backend.QueryDataRequest) string {
 	from, to := epochMStoInfluxTime(&queryContext.Queries[0].TimeRange)
 	return fmt.Sprintf("time >= %s and time <= %s", from, to)

pkg/tsdb/influxdb/models/query_test.go

@@ -1,11 +1,13 @@
 package models
 
 import (
+	"fmt"
 	"strings"
 	"testing"
 	"time"
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
+	"github.com/influxdata/influxql"
 	"github.com/stretchr/testify/require"
 )
 
@@ -329,3 +331,75 @@ func TestRemoveRegexWrappers(t *testing.T) {
 		require.Equal(t, expected, result)
 	})
 }
+
+func TestRawQueryWithAdhocFilters(t *testing.T) {
+	noAdhocFilters := []*Tag{}
+	adhocFilter := &Tag{
+		Key:   "lol",
+		Value: "sob",
+	}
+
+	cases := []struct {
+		Name          string
+		RawQuery      string
+		ExpectedQuery string
+		AdhocFilters  []*Tag
+	}{
+		{
+			Name:          "no filters",
+			RawQuery:      `SELECT "gauge" FROM "a"."b" WHERE time >= 123456780 and time <= 123456789`,
+			ExpectedQuery: `SELECT "gauge" FROM "a"."b" WHERE time >= 123456780 and time <= 123456789`,
+			AdhocFilters:  noAdhocFilters,
+		},
+		{
+			Name:          "multi-byte chars, single filter",
+			RawQuery:      `SELECT "gauge" as "世界" FROM "a"."b"`,
+			ExpectedQuery: `SELECT "gauge" as "世界" FROM "a"."b" WHERE "lol" = 'sob'`,
+			AdhocFilters:  []*Tag{adhocFilter},
+		},
+		{
+			Name:          "multi-byte chars, single filter, existing where condition",
+			RawQuery:      `SELECT "gauge" as "世界" FROM "a"."b" WHERE time >= 123456780 AND time <= 123456789`,
+			ExpectedQuery: `SELECT "gauge" as "世界" FROM "a"."b" WHERE "lol" = 'sob' AND time >= 123456780 AND time <= 123456789`,
+			AdhocFilters:  []*Tag{adhocFilter},
+		},
+		{
+			Name:          "quoted where, single filter",
+			RawQuery:      `SELECT "gauge" as "where is my thing" FROM "a"."b"`,
+			ExpectedQuery: `SELECT "gauge" as "where is my thing" FROM "a"."b" WHERE "lol" = 'sob'`,
+			AdhocFilters:  []*Tag{adhocFilter},
+		},
+		{
+			Name:          "quoted where, single filter, existing where condition",
+			RawQuery:      `SELECT "gauge" as "where is my thing" FROM "a"."b" WHERE time >= 123456780 AND time <= 123456789`,
+			ExpectedQuery: `SELECT "gauge" as "where is my thing" FROM "a"."b" WHERE "lol" = 'sob' AND time >= 123456780 AND time <= 123456789`,
+			AdhocFilters:  []*Tag{adhocFilter},
+		},
+		{
+			Name:          "multiple filters",
+			RawQuery:      `SELECT "gauge" FROM "a"."b"`,
+			ExpectedQuery: `SELECT "gauge" FROM "a"."b" WHERE "lol" = 'sob' AND "lol" = 'sob'`,
+			AdhocFilters:  []*Tag{adhocFilter, adhocFilter},
+		},
+		{
+			Name:          "multiple filters, existing where condition",
+			RawQuery:      `SELECT "gauge" FROM "a"."b" WHERE time >= 123456780 AND time <= 123456789`,
+			ExpectedQuery: `SELECT "gauge" FROM "a"."b" WHERE "lol" = 'sob' AND "lol" = 'sob' AND time >= 123456780 AND time <= 123456789`,
+			AdhocFilters:  []*Tag{adhocFilter, adhocFilter},
+		},
+	}
+
+	for _, c := range cases {
+		t.Run(fmt.Sprintf("raw query + adhoc filters: %s", c.Name), func(t *testing.T) {
+			query := &Query{
+				RawQuery:     c.RawQuery,
+				UseRawQuery:  true,
+				AdhocFilters: c.AdhocFilters,
+			}
+			actualQuery := query.getRawQueryWithAdhocFilters()
+			require.Equal(t, c.ExpectedQuery, actualQuery)
+			_, err := influxql.ParseStatement(actualQuery)
+			require.NoError(t, err)
+		})
+	}
+}