Loki shard splitting: enable for supported metric queries (#97675)

* Loki: remove app restriction for sharding

* Loki: allow some metric queries to be sharded

* Enable metric queries wrapped by sum

* Prettier

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

@@ -361,11 +361,7 @@ export class LokiDatasource
       return this.runLiveQueryThroughBackend(fixedRequest);
     }
 
-    if (
-      config.featureToggles.lokiShardSplitting &&
-      requestSupportsSharding(fixedRequest.targets) &&
-      fixedRequest.app === CoreApp.Explore
-    ) {
+    if (config.featureToggles.lokiShardSplitting && requestSupportsSharding(fixedRequest.targets)) {
       return runShardSplitQuery(this, fixedRequest);
     } else if (config.featureToggles.lokiQuerySplitting && requestSupportsSplitting(fixedRequest.targets)) {
       return runSplitQuery(this, fixedRequest);

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

@@ -18,8 +18,9 @@ import {
   getNodePositionsFromQuery,
   getLogQueryFromMetricsQueryAtPosition,
   interpolateShardingSelector,
+  requestSupportsSharding,
 } from './queryUtils';
-import { LokiQuery, LokiQueryType } from './types';
+import { LokiQuery, LokiQueryDirection, LokiQueryType } from './types';
 
 describe('getHighlighterExpressionsFromQuery', () => {
   it('returns no expressions for empty query', () => {
@@ -590,3 +591,39 @@ describe('interpolateShardingSelector', () => {
     });
   });
 });
+
+describe('requestSupportsSharding', () => {
+  it('supports log queries with Scan direction', () => {
+    expect(requestSupportsSharding([{ refId: 'A', expr: '{place="luna"}', direction: LokiQueryDirection.Scan }])).toBe(
+      true
+    );
+  });
+
+  it('declines log queries without Scan direction', () => {
+    expect(
+      requestSupportsSharding([{ refId: 'A', expr: '{place="luna"}', direction: LokiQueryDirection.Backward }])
+    ).toBe(false);
+  });
+
+  it.each([
+    'count_over_time({place="luna"}[1m])',
+    'sum_over_time({place="luna"}[1m])',
+    'sum by (level) (count_over_time({place="luna"}[1m]))',
+    'sum by (level) (rate({place="luna"}[1m]))',
+    'sum(sum by (level) (avg_over_time({place="luna"}[1m])))',
+    'sum(rate({place="luna"}[1m]))',
+  ])('allows supported metric queries', (expr: string) => {
+    expect(requestSupportsSharding([{ refId: 'A', expr }])).toBe(true);
+  });
+
+  it.each([
+    'avg_over_time({place="luna"}[1m])',
+    'avg(sum_over_time({place="luna"}[1m]))',
+    'avg(rate({place="luna"}[1m]))',
+    'count_over_time({place="luna"}[1m]) / count_over_time({place="luna"}[1m])',
+    'avg(sum by (level) (avg_over_time({place="luna"}[1m])))',
+    'sum(rate({place="luna"}[1m])) / sum(rate({place="luna"}[1m]))',
+  ])('declines supported metric queries', (expr: string) => {
+    expect(requestSupportsSharding([{ refId: 'A', expr }])).toBe(false);
+  });
+});

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

@@ -22,6 +22,9 @@ import {
   Json,
   OrFilter,
   FilterOp,
+  RangeOp,
+  VectorOp,
+  BinOpExpr,
 } from '@grafana/lezer-logql';
 import { DataQuery } from '@grafana/schema';
 
@@ -319,11 +322,56 @@ export function requestSupportsSharding(allQueries: LokiQuery[]) {
     .filter((query) => query.queryType !== LokiQueryType.Instant)
     .filter((query) => !query.refId.includes('do-not-shard'))
     .filter((query) => query.expr)
-    .filter((query) => query.direction === LokiQueryDirection.Scan && isLogsQuery(query.expr));
+    .filter(
+      (query) =>
+        (query.direction === LokiQueryDirection.Scan && isLogsQuery(query.expr)) || metricSupportsSharding(query.expr)
+    );
 
   return queries.length > 0;
 }
 
+function metricSupportsSharding(query: string) {
+  if (isLogsQuery(query)) {
+    return false;
+  }
+  query = query.trim().toLowerCase();
+
+  const disallowed = getNodesFromQuery(query, [BinOpExpr]);
+  if (disallowed.length > 0) {
+    return false;
+  }
+
+  /**
+   * If there are VectorAggregationExpr, we want to make sure that the leftmost VectorOp is sum, meaning that
+   * it's wrapped in a sum. E.g.
+   * Disallowed: avg(sum by (level) (avg_over_time({place="luna"}[1m])))
+   * Allowed: sum(sum by (level) (avg_over_time({place="luna"}[1m])))
+   */
+  const vectorOps = getNodesFromQuery(query, [VectorOp]);
+  const supportedVectorOps = vectorOps.filter((node) => getNodeString(query, node) === 'sum');
+  const unsupportedVectorOps = vectorOps.filter((node) => getNodeString(query, node) !== 'sum');
+  const supportedWrappingVectorOpps = supportedVectorOps.filter((supportedOp) =>
+    unsupportedVectorOps.every((unsupportedOp) => supportedOp.from < unsupportedOp.from)
+  );
+  if (unsupportedVectorOps.length > 0) {
+    return supportedWrappingVectorOpps.length > 0;
+  }
+
+  const rangeOps = getNodesFromQuery(query, [RangeOp]);
+  const supportedRangeOps = ['count_over_time', 'sum_over_time', 'bytes_over_time'];
+  for (const node of rangeOps) {
+    if (!supportedRangeOps.includes(getNodeString(query, node))) {
+      return supportedWrappingVectorOpps.length > 0;
+    }
+  }
+
+  return true;
+}
+
+function getNodeString(query: string, node: SyntaxNode) {
+  return query.substring(node.from, node.to);
+}
+
 export const isLokiQuery = (query: DataQuery): query is LokiQuery => {
   if (!query) {
     return false;