From c324376999f7c4b3ced9bf66e3792373d18cd73b Mon Sep 17 00:00:00 2001 From: Matias Chomicki Date: Mon, 9 Dec 2024 16:54:25 +0000 Subject: [PATCH] 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 --- .../app/plugins/datasource/loki/datasource.ts | 6 +-- .../datasource/loki/queryUtils.test.ts | 39 ++++++++++++++- .../app/plugins/datasource/loki/queryUtils.ts | 50 ++++++++++++++++++- 3 files changed, 88 insertions(+), 7 deletions(-) diff --git a/public/app/plugins/datasource/loki/datasource.ts b/public/app/plugins/datasource/loki/datasource.ts index dc52aca0a11..4a436dc4b87 100644 --- a/public/app/plugins/datasource/loki/datasource.ts +++ b/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); diff --git a/public/app/plugins/datasource/loki/queryUtils.test.ts b/public/app/plugins/datasource/loki/queryUtils.test.ts index 1dc2c64fc0e..e901b77c55b 100644 --- a/public/app/plugins/datasource/loki/queryUtils.test.ts +++ b/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); + }); +}); diff --git a/public/app/plugins/datasource/loki/queryUtils.ts b/public/app/plugins/datasource/loki/queryUtils.ts index 1a5f01bf25c..64ad1a8a4ae 100644 --- a/public/app/plugins/datasource/loki/queryUtils.ts +++ b/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;