Loki: Improve getLogQueryFromMetricsQuery (#75492)

* Query utils: fix bug in getLogQueryFromMetricsQuery function

* Query utils: implement getLogQueryFromMetricsQueryAtPosition

* Loki autocomplete situation: use getLogQueryFromMetricsQueryAtPosition

* Fix typo

* getLogQueryFromMetricsQueryAtPosition: use find instead of filter

public/app/plugins/datasource/loki/components/monaco-query-field/monaco-completion-provider/situation.ts

@@ -30,7 +30,7 @@ import {
   LabelExtractionExpressionList,
 } from '@grafana/lezer-logql';
 
-import { getLogQueryFromMetricsQuery, getNodesFromQuery } from '../../../queryUtils';
+import { getLogQueryFromMetricsQueryAtPosition, getNodesFromQuery } from '../../../queryUtils';
 
 type Direction = 'parent' | 'firstChild' | 'lastChild' | 'nextSibling';
 type NodeType = number;
@@ -303,7 +303,7 @@ function getLabels(selectorNode: SyntaxNode, text: string): Label[] {
 function resolveAfterUnwrap(node: SyntaxNode, text: string, pos: number): Situation | null {
   return {
     type: 'AFTER_UNWRAP',
-    logQuery: getLogQueryFromMetricsQuery(text).trim(),
+    logQuery: getLogQueryFromMetricsQueryAtPosition(text, pos).trim(),
   };
 }
 
@@ -353,7 +353,7 @@ function resolveLabelsForGrouping(node: SyntaxNode, text: string, pos: number):
 
   return {
     type: 'IN_GROUPING',
-    logQuery: getLogQueryFromMetricsQuery(text).trim(),
+    logQuery: getLogQueryFromMetricsQueryAtPosition(text, pos).trim(),
   };
 }
 
@@ -465,7 +465,7 @@ function resolveLogfmtParser(_: SyntaxNode, text: string, cursorPosition: number
     type: 'IN_LOGFMT',
     otherLabels,
     flags,
-    logQuery: getLogQueryFromMetricsQuery(text).trim(),
+    logQuery: getLogQueryFromMetricsQueryAtPosition(text, position).trim(),
   };
 }
 
@@ -539,7 +539,7 @@ function resolveLogOrLogRange(node: SyntaxNode, text: string, pos: number, after
     type: 'AFTER_SELECTOR',
     afterPipe,
     hasSpace: text.charAt(pos - 1) === ' ',
-    logQuery: getLogQueryFromMetricsQuery(text).trim(),
+    logQuery: getLogQueryFromMetricsQueryAtPosition(text, pos).trim(),
   };
 }
 
@@ -582,7 +582,7 @@ function resolveSelector(node: SyntaxNode, text: string, pos: number): Situation
 }
 
 function resolveAfterKeepAndDrop(node: SyntaxNode, text: string, pos: number): Situation | null {
-  let logQuery = getLogQueryFromMetricsQuery(text).trim();
+  let logQuery = getLogQueryFromMetricsQueryAtPosition(text, pos).trim();
   let keepAndDropParent: SyntaxNode | null = null;
   let parent = node.parent;
   while (parent !== null) {

public/app/plugins/datasource/loki/queryUtils.test.ts

@@ -18,6 +18,7 @@ import {
   getNormalizedLokiQuery,
   getNodePositionsFromQuery,
   formatLogqlQuery,
+  getLogQueryFromMetricsQueryAtPosition,
 } from './queryUtils';
 import { LokiQuery, LokiQueryType } from './types';
 
@@ -427,6 +428,39 @@ describe('getLogQueryFromMetricsQuery', () => {
       )
     ).toBe('{label="$var"} | logfmt | __error__=``');
   });
+  it('does not return a query when there is no log query', () => {
+    expect(getLogQueryFromMetricsQuery('1+1')).toBe('');
+    expect(getLogQueryFromMetricsQuery('count_over_time([1s])')).toBe('');
+  });
+});
+
+describe('getLogQueryFromMetricsQueryAtPosition', () => {
+  it('works like getLogQueryFromMetricsQuery for simple queries', () => {
+    expect(
+      getLogQueryFromMetricsQueryAtPosition('count_over_time({job="grafana"} | logfmt | label="value" [1m])', 57)
+    ).toBe('{job="grafana"} | logfmt | label="value"');
+    expect(getLogQueryFromMetricsQueryAtPosition('count_over_time({job="grafana"} [1m])', 37)).toBe('{job="grafana"}');
+    expect(
+      getLogQueryFromMetricsQueryAtPosition(
+        'sum(quantile_over_time(0.5, {label="$var"} | logfmt | __error__=`` | unwrap latency | __error__=`` [$__interval]))',
+        45
+      )
+    ).toBe('{label="$var"} | logfmt | __error__=``');
+  });
+  it.each([
+    [
+      'count_over_time({place="moon"} | json test="test" [1m]) + avg_over_time({place="luna"} | logfmt test="test" [1m])',
+      '{place="moon"} | json test="test"',
+      49,
+    ],
+    [
+      'count_over_time({place="moon"} | json test="test" [1m]) + avg_over_time({place="luna"} | logfmt test="test" [1m])',
+      '{place="luna"} | logfmt test="test"',
+      107,
+    ],
+  ])('gets the right query for complex queries', (metric: string, log: string, position: number) => {
+    expect(getLogQueryFromMetricsQueryAtPosition(metric, position)).toBe(log);
+  });
 });
 
 describe('getNodePositionsFromQuery', () => {

public/app/plugins/datasource/loki/queryUtils.ts

@@ -231,7 +231,7 @@ export function getLogQueryFromMetricsQuery(query: string): string {
   // Log query in metrics query composes of Selector & PipelineExpr
   const selectorNode = getNodeFromQuery(query, Selector);
   if (!selectorNode) {
-    return query;
+    return '';
   }
   const selector = query.substring(selectorNode.from, selectorNode.to);
 
@@ -241,6 +241,20 @@ export function getLogQueryFromMetricsQuery(query: string): string {
   return `${selector} ${pipelineExpr}`.trim();
 }
 
+export function getLogQueryFromMetricsQueryAtPosition(query: string, position: number): string {
+  if (isLogsQuery(query)) {
+    return query;
+  }
+
+  const metricQuery = getNodesFromQuery(query, [MetricExpr])
+    .reverse() // So we don't get the root metric node
+    .find((node) => node.from <= position && node.to >= position);
+  if (!metricQuery) {
+    return '';
+  }
+  return getLogQueryFromMetricsQuery(query.substring(metricQuery.from, metricQuery.to));
+}
+
 export function isQueryWithLabelFilter(query: string): boolean {
   return isQueryWithNode(query, LabelFilter);
 }