diff --git a/pkg/tsdb/cloudwatch/get_metric_query_batches.go b/pkg/tsdb/cloudwatch/get_metric_query_batches.go index 924966bc119..2b4f0b47667 100644 --- a/pkg/tsdb/cloudwatch/get_metric_query_batches.go +++ b/pkg/tsdb/cloudwatch/get_metric_query_batches.go @@ -13,72 +13,79 @@ var nonWordRegex = regexp.MustCompile(`\W+`) // getMetricQueryBatches separates queries into batches if necessary. Metric Insight queries cannot run together, and math expressions must be run // with all the queries they reference. func getMetricQueryBatches(queries []*models.CloudWatchQuery, logger log.Logger) [][]*models.CloudWatchQuery { - metricInsightIndices := []int{} - mathIndices := []int{} - for i, query := range queries { - switch query.GetGetMetricDataAPIMode() { - case models.GMDApiModeSQLExpression: - metricInsightIndices = append(metricInsightIndices, i) - case models.GMDApiModeMathExpression: - mathIndices = append(mathIndices, i) - default: - } - } - // We only need multiple batches if there are multiple metrics insight queries - if len(metricInsightIndices) <= 1 { + if !hasMultipleMetricInsights(queries) { return [][]*models.CloudWatchQuery{queries} } logger.Debug("Separating queries into batches") - // Map ids to their queries - idToIndex := map[string]int{} - for i, query := range queries { - if query.Id != "" { - idToIndex[query.Id] = i + + // set up list of math expression queries since below we loop over them to get what query IDs they reference + var mathQueries []*models.CloudWatchQuery + for _, query := range queries { + if query.GetGetMetricDataAPIMode() == models.GMDApiModeMathExpression { + mathQueries = append(mathQueries, query) } } - // Find and track which queries are referenced by math queries - queryReferences := make([][]int, len(queries)) - isReferenced := make([]bool, len(queries)) - for _, idx := range mathIndices { - tokens := nonWordRegex.Split(queries[idx].Expression, -1) - references := []int{} - for _, token := range tokens { - ref, found := idToIndex[token] + // put queries into a set in order to facilitate lookup below + idToQuery := make(map[string]*models.CloudWatchQuery, len(queries)) + for _, q := range queries { + idToQuery[q.Id] = q + } + + // gets query IDs which are referenced in math expressions + mathQueryIdToReferences := make(map[string][]*models.CloudWatchQuery) + // we will use this set of referenced queries to determine the root queries below + referencedQueries := make(map[string]bool) + for _, mathQuery := range mathQueries { + substrings := nonWordRegex.Split(mathQuery.Expression, -1) + for _, id := range substrings { + query, found := idToQuery[id] if found { - references = append(references, ref) - isReferenced[ref] = true + mathQueryIdToReferences[mathQuery.Id] = append(mathQueryIdToReferences[mathQuery.Id], query) + referencedQueries[query.Id] = true } } - queryReferences[idx] = references } - // Create a new batch for every query not used in another query batches := [][]*models.CloudWatchQuery{} - for i, used := range isReferenced { - if !used { - batches = append(batches, getReferencedQueries(queries, queryReferences, i)) + for _, query := range queries { + // if a query is not referenced, then it is a "root" query + if _, ok := referencedQueries[query.Id]; !ok { + batches = append(batches, getConnectedQueries(query, mathQueryIdToReferences)) } } return batches } -// getReferencedQueries gets all the queries referenced by startQuery and its referenced queries -func getReferencedQueries(queries []*models.CloudWatchQuery, queryReferences [][]int, startQuery int) []*models.CloudWatchQuery { - usedQueries := make([]bool, len(queries)) - batch := []*models.CloudWatchQuery{} +// getConnectedQueries does a breadth-first search to find all the query ids connected to the root id by references. The root id is also returned in the response. +func getConnectedQueries(root *models.CloudWatchQuery, queryReferences map[string][]*models.CloudWatchQuery) []*models.CloudWatchQuery { + visited := map[string]bool{root.Id: true} + queriesToReturn := []*models.CloudWatchQuery{} - queriesToAdd := []int{startQuery} - usedQueries[startQuery] = true - for i := 0; i < len(queriesToAdd); i++ { - batch = append(batch, queries[queriesToAdd[i]]) - for _, queryIdx := range queryReferences[queriesToAdd[i]] { - if !usedQueries[queryIdx] { - usedQueries[queryIdx] = true - queriesToAdd = append(queriesToAdd, queryIdx) + queriesToVisit := []*models.CloudWatchQuery{root} + for i := 0; i < len(queriesToVisit); i++ { + currentQuery := queriesToVisit[i] + queriesToReturn = append(queriesToReturn, currentQuery) + for _, queryRef := range queryReferences[currentQuery.Id] { + if !visited[queryRef.Id] { + visited[queryRef.Id] = true + queriesToVisit = append(queriesToVisit, queryRef) } } } - return batch + return queriesToReturn +} + +func hasMultipleMetricInsights(queries []*models.CloudWatchQuery) bool { + count := 0 + for _, query := range queries { + if query.GetGetMetricDataAPIMode() == models.GMDApiModeSQLExpression { + count++ + } + if count > 1 { + return true + } + } + return false } diff --git a/pkg/tsdb/cloudwatch/get_metric_query_batches_test.go b/pkg/tsdb/cloudwatch/get_metric_query_batches_test.go index 59023ba2843..f00fb982856 100644 --- a/pkg/tsdb/cloudwatch/get_metric_query_batches_test.go +++ b/pkg/tsdb/cloudwatch/get_metric_query_batches_test.go @@ -22,6 +22,7 @@ func TestGetMetricQueryBatches(t *testing.T) { MetricQueryType: models.MetricQueryTypeQuery, Id: "i3", } + metricStat := models.CloudWatchQuery{ MetricQueryType: models.MetricQueryTypeSearch, MetricEditorMode: models.MetricEditorModeBuilder, @@ -33,6 +34,24 @@ func TestGetMetricQueryBatches(t *testing.T) { Expression: "PERIOD(i1)", Id: "m1", } + m99_ref_m98 := models.CloudWatchQuery{ + MetricQueryType: models.MetricQueryTypeSearch, + MetricEditorMode: models.MetricEditorModeRaw, + Expression: "PERIOD(m98)", + Id: "m99", + } + m98_ref_m88 := models.CloudWatchQuery{ + MetricQueryType: models.MetricQueryTypeSearch, + MetricEditorMode: models.MetricEditorModeRaw, + Expression: "PERIOD(m88)", + Id: "m98", + } + m88_ref_m98 := models.CloudWatchQuery{ + MetricQueryType: models.MetricQueryTypeSearch, + MetricEditorMode: models.MetricEditorModeRaw, + Expression: "PERIOD(m98)", + Id: "m88", + } m2_ref_i1 := models.CloudWatchQuery{ MetricQueryType: models.MetricQueryTypeSearch, MetricEditorMode: models.MetricEditorModeRaw, @@ -58,6 +77,22 @@ func TestGetMetricQueryBatches(t *testing.T) { Id: "m5", } + t.Run("m99 ref m98 which ref m88 which ref m98, with 2 insights", func(t *testing.T) { + batch := []*models.CloudWatchQuery{ + &insight1, + &insight2, + &m99_ref_m98, + &m98_ref_m88, + &m88_ref_m98, + } + + result := getMetricQueryBatches(batch, logger) + assert.Len(t, result, 3) + assert.ElementsMatch(t, []*models.CloudWatchQuery{&insight1}, result[0]) + assert.ElementsMatch(t, []*models.CloudWatchQuery{&insight2}, result[1]) + assert.ElementsMatch(t, []*models.CloudWatchQuery{&m99_ref_m98, &m98_ref_m88, &m88_ref_m98}, result[2]) + }) + t.Run("zero insight queries should not separate into batches", func(t *testing.T) { batch := []*models.CloudWatchQuery{ &metricStat,