From c9bc937919d882f55876320aac4a2e3d2c3e04a1 Mon Sep 17 00:00:00 2001
From: Ivana Huckova <30407135+ivanahuckova@users.noreply.github.com>
Date: Wed, 31 Jan 2024 13:39:07 +0100
Subject: [PATCH] Loki: Limit running of samples based on query and time range
 (#81585)

* Loki: Limit running of samples based on query and time range

* Update
---
 .../components/LokiQueryBuilder.test.tsx      | 100 +++++++++++++++++-
 .../components/LokiQueryBuilder.tsx           |  15 ++-
 2 files changed, 112 insertions(+), 3 deletions(-)

diff --git a/public/app/plugins/datasource/loki/querybuilder/components/LokiQueryBuilder.test.tsx b/public/app/plugins/datasource/loki/querybuilder/components/LokiQueryBuilder.test.tsx
index 024a7cddd5d..3cec60ded95 100644
--- a/public/app/plugins/datasource/loki/querybuilder/components/LokiQueryBuilder.test.tsx
+++ b/public/app/plugins/datasource/loki/querybuilder/components/LokiQueryBuilder.test.tsx
@@ -3,11 +3,12 @@ import userEvent from '@testing-library/user-event';
 import React from 'react';
 
 import { dateTime } from '@grafana/data';
+import { config } from '@grafana/runtime';
 
 import { createLokiDatasource } from '../../__mocks__/datasource';
 import { LokiOperationId, LokiVisualQuery } from '../types';
 
-import { LokiQueryBuilder } from './LokiQueryBuilder';
+import { LokiQueryBuilder, TIME_SPAN_TO_TRIGGER_SAMPLES } from './LokiQueryBuilder';
 import { EXPLAIN_LABEL_FILTER_CONTENT } from './LokiQueryBuilderExplained';
 
 const MISSING_LABEL_FILTER_ERROR_MESSAGE = 'Select at least 1 label filter (label and value)';
@@ -40,6 +41,14 @@ const createDefaultProps = () => {
 };
 
 describe('LokiQueryBuilder', () => {
+  const originalLokiQueryHints = config.featureToggles.lokiQueryHints;
+  beforeEach(() => {
+    config.featureToggles.lokiQueryHints = true;
+  });
+
+  afterEach(() => {
+    config.featureToggles.lokiQueryHints = originalLokiQueryHints;
+  });
   it('tries to load labels when no labels are selected', async () => {
     const props = createDefaultProps();
     props.datasource.getDataSamples = jest.fn().mockResolvedValue([]);
@@ -134,6 +143,95 @@ describe('LokiQueryBuilder', () => {
       expect(screen.queryByText(EXPLAIN_LABEL_FILTER_CONTENT)).not.toBeInTheDocument();
     });
   });
+
+  it('re-runs sample query when query changes', async () => {
+    const query = {
+      labels: [{ label: 'foo', op: '=', value: 'bar' }],
+      operations: [{ id: LokiOperationId.LineContains, params: ['error'] }],
+    };
+    const props = createDefaultProps();
+    props.datasource.getDataSamples = jest.fn().mockResolvedValue([]);
+
+    const { rerender } = render(<LokiQueryBuilder {...props} query={query} />);
+    rerender(
+      <LokiQueryBuilder
+        {...props}
+        query={{ ...query, labels: [...query.labels, { label: 'xyz', op: '=', value: 'abc' }] }}
+      />
+    );
+
+    await waitFor(() => {
+      expect(props.datasource.getDataSamples).toHaveBeenCalledTimes(2);
+    });
+  });
+
+  it('does not re-run sample query when query does not change', async () => {
+    const query = {
+      labels: [{ label: 'foo', op: '=', value: 'bar' }],
+      operations: [{ id: LokiOperationId.LineContains, params: ['error'] }],
+    };
+    const props = createDefaultProps();
+    props.datasource.getDataSamples = jest.fn().mockResolvedValue([]);
+
+    const { rerender } = render(<LokiQueryBuilder {...props} query={query} />);
+    rerender(<LokiQueryBuilder {...props} query={query} />);
+
+    await waitFor(() => {
+      expect(props.datasource.getDataSamples).toHaveBeenCalledTimes(1);
+    });
+  });
+
+  it('re-run sample query when time range changes over 5 minutes', async () => {
+    const query = {
+      labels: [{ label: 'foo', op: '=', value: 'bar' }],
+      operations: [{ id: LokiOperationId.LineContains, params: ['error'] }],
+    };
+    const props = createDefaultProps();
+    const updatedFrom = dateTime(props.timeRange.from.valueOf() + TIME_SPAN_TO_TRIGGER_SAMPLES + 1000);
+    const updatedTo = dateTime(props.timeRange.to.valueOf() + TIME_SPAN_TO_TRIGGER_SAMPLES + 1000);
+    const updatedTimeRange = {
+      from: updatedFrom,
+      to: updatedTo,
+      raw: {
+        from: updatedFrom,
+        to: updatedTo,
+      },
+    };
+    props.datasource.getDataSamples = jest.fn().mockResolvedValue([]);
+
+    const { rerender } = render(<LokiQueryBuilder {...props} query={query} />);
+    rerender(<LokiQueryBuilder {...props} query={query} timeRange={updatedTimeRange} />);
+
+    await waitFor(() => {
+      expect(props.datasource.getDataSamples).toHaveBeenCalledTimes(2);
+    });
+  });
+
+  it('does not re-run sample query when time range changes less than 5 minutes', async () => {
+    const query = {
+      labels: [{ label: 'foo', op: '=', value: 'bar' }],
+      operations: [{ id: LokiOperationId.LineContains, params: ['error'] }],
+    };
+    const props = createDefaultProps();
+    const updatedFrom = dateTime(props.timeRange.from.valueOf() + TIME_SPAN_TO_TRIGGER_SAMPLES - 1000);
+    const updatedTo = dateTime(props.timeRange.to.valueOf() + TIME_SPAN_TO_TRIGGER_SAMPLES - 1000);
+    const updatedTimeRange = {
+      from: updatedFrom,
+      to: updatedTo,
+      raw: {
+        from: updatedFrom,
+        to: updatedTo,
+      },
+    };
+    props.datasource.getDataSamples = jest.fn().mockResolvedValue([]);
+
+    const { rerender } = render(<LokiQueryBuilder {...props} query={query} />);
+    rerender(<LokiQueryBuilder {...props} query={query} timeRange={updatedTimeRange} />);
+
+    await waitFor(() => {
+      expect(props.datasource.getDataSamples).toHaveBeenCalledTimes(1);
+    });
+  });
 });
 
 const getSelectParent = (input: HTMLElement) =>
diff --git a/public/app/plugins/datasource/loki/querybuilder/components/LokiQueryBuilder.tsx b/public/app/plugins/datasource/loki/querybuilder/components/LokiQueryBuilder.tsx
index 6c10532c5d0..0261c8f7bdb 100644
--- a/public/app/plugins/datasource/loki/querybuilder/components/LokiQueryBuilder.tsx
+++ b/public/app/plugins/datasource/loki/querybuilder/components/LokiQueryBuilder.tsx
@@ -1,4 +1,6 @@
+import { isEqual } from 'lodash';
 import React, { useEffect, useMemo, useState } from 'react';
+import { usePrevious } from 'react-use';
 
 import { DataSourceApi, getDefaultTimeRange, LoadingState, PanelData, SelectableValue, TimeRange } from '@grafana/data';
 import {
@@ -29,6 +31,7 @@ import { LokiOperationId, LokiVisualQuery } from '../types';
 import { EXPLAIN_LABEL_FILTER_CONTENT } from './LokiQueryBuilderExplained';
 import { NestedQueryList } from './NestedQueryList';
 
+export const TIME_SPAN_TO_TRIGGER_SAMPLES = 5 * 60 * 1000;
 export interface Props {
   query: LokiVisualQuery;
   datasource: LokiDatasource;
@@ -41,6 +44,8 @@ export const LokiQueryBuilder = React.memo<Props>(
   ({ datasource, query, onChange, onRunQuery, showExplain, timeRange }) => {
     const [sampleData, setSampleData] = useState<PanelData>();
     const [highlightedOp, setHighlightedOp] = useState<QueryBuilderOperation | undefined>(undefined);
+    const prevQuery = usePrevious(query);
+    const prevTimeRange = usePrevious(timeRange);
 
     const onChangeLabels = (labels: QueryBuilderLabelFilter[]) => {
       onChange({ ...query, labels });
@@ -109,10 +114,16 @@ export const LokiQueryBuilder = React.memo<Props>(
         setSampleData(sampleData);
       };
 
-      if (config.featureToggles.lokiQueryHints) {
+      const updateBasedOnChangedTimeRange =
+        prevTimeRange &&
+        timeRange &&
+        (Math.abs(timeRange.to.valueOf() - prevTimeRange.to.valueOf()) > TIME_SPAN_TO_TRIGGER_SAMPLES ||
+          Math.abs(timeRange.from.valueOf() - prevTimeRange.from.valueOf()) > TIME_SPAN_TO_TRIGGER_SAMPLES);
+      const updateBasedOnChangedQuery = !isEqual(prevQuery, query);
+      if (config.featureToggles.lokiQueryHints && (updateBasedOnChangedTimeRange || updateBasedOnChangedQuery)) {
         onGetSampleData().catch(console.error);
       }
-    }, [datasource, query, timeRange]);
+    }, [datasource, query, timeRange, prevQuery, prevTimeRange]);
 
     const lang = { grammar: logqlGrammar, name: 'logql' };
     return (