SceneQueryRunner: decouple time range comparisons

packages/scenes/src/components/SceneTimeRangeCompare.tsx

@@ -1,17 +1,17 @@
-import { DateTime, dateTime, GrafanaTheme2, rangeUtil, TimeRange } from '@grafana/data';
+import { DataQueryRequest, DateTime, dateTime, FieldType, GrafanaTheme2, rangeUtil, TimeRange } from '@grafana/data';
+import { config } from '@grafana/runtime';
 import { ButtonGroup, ButtonSelect, Checkbox, ToolbarButton, useStyles2 } from '@grafana/ui';
 import React from 'react';
 import { sceneGraph } from '../core/sceneGraph';
 import { SceneObjectBase } from '../core/SceneObjectBase';
 import { SceneComponentProps, SceneObjectState, SceneObjectUrlValues } from '../core/types';
+import { DataQueryExtended } from '../querying/SceneQueryRunner';
+import { SupplementaryRequest, ProcessorFunc, SupplementaryRequestProvider } from '../querying/SupplementaryRequestProvider';
 import { SceneObjectUrlSyncConfig } from '../services/SceneObjectUrlSyncConfig';
+import { getCompareSeriesRefId } from '../utils/getCompareSeriesRefId';
 import { parseUrlParam } from '../utils/parseUrlParam';
 import { css } from '@emotion/css';
 
-export interface TimeRangeCompareProvider {
-  getCompareTimeRange(timeRange: TimeRange): TimeRange | undefined;
-}
-
 interface SceneTimeRangeCompareState extends SceneObjectState {
   compareWith?: string;
   compareOptions: Array<{ label: string; value: string }>;
@@ -38,8 +38,8 @@ export const DEFAULT_COMPARE_OPTIONS = [
 
 export class SceneTimeRangeCompare
   extends SceneObjectBase<SceneTimeRangeCompareState>
-  implements TimeRangeCompareProvider
-{
+  implements SupplementaryRequestProvider<SceneTimeRangeCompareState> {
+
   static Component = SceneTimeRangeCompareRenderer;
   protected _urlSync = new SceneObjectUrlSyncConfig(this, { keys: ['compareWith'] });
 
@@ -94,6 +94,33 @@ export class SceneTimeRangeCompare
     this.setState({ compareWith: undefined });
   };
 
+  // Get a time shifted request to compare with the primary request.
+  public getSupplementaryRequests(request: DataQueryRequest): SupplementaryRequest[] {
+    const extraRequests: SupplementaryRequest[] = [];
+    const compareRange = this.getCompareTimeRange(request.range);
+    if (!compareRange) {
+      return extraRequests;
+    }
+
+    const targets = request.targets.filter((query: DataQueryExtended) => query.timeRangeCompare !== false);
+    if (targets.length) {
+      extraRequests.push({
+        req: {
+          ...request,
+          targets,
+          range: compareRange,
+        },
+        processor: timeShiftAlignmentProcessor,
+      });
+    }
+    return extraRequests;
+  }
+
+  // The query runner should rerun the comparison query if the compareWith value has changed.
+  public shouldRerun(prev: SceneTimeRangeCompareState, next: SceneTimeRangeCompareState): boolean {
+    return prev.compareWith !== next.compareWith;
+  }
+
   public getCompareTimeRange(timeRange: TimeRange): TimeRange | undefined {
     let compareFrom: DateTime;
     let compareTo: DateTime;
@@ -149,6 +176,43 @@ export class SceneTimeRangeCompare
   }
 }
 
+// Processor function for use with time shifted comparison series.
+// This aligns the secondary series with the primary and adds custom
+// metadata and config to the secondary series' fields so that it is
+// rendered appropriately.
+const timeShiftAlignmentProcessor: ProcessorFunc = (primary, secondary) => {
+  const diff = secondary.timeRange.from.diff(primary.timeRange.from);
+  secondary.series.forEach((series) => {
+    series.refId = getCompareSeriesRefId(series.refId || '');
+    series.meta = {
+      ...series.meta,
+      // @ts-ignore Remove when https://github.com/grafana/grafana/pull/71129 is released
+      timeCompare: {
+        diffMs: diff,
+        isTimeShiftQuery: true,
+      },
+    };
+    series.fields.forEach((field) => {
+      // Align compare series time stamps with reference series
+      if (field.type === FieldType.time) {
+        field.values = field.values.map((v) => {
+          return diff < 0 ? v - diff : v + diff;
+        });
+      }
+
+      field.config = {
+        ...field.config,
+        color: {
+          mode: 'fixed',
+          fixedColor: config.theme.palette.gray60,
+        },
+      };
+      return field;
+    });
+  });
+  return secondary;
+}
+
 function SceneTimeRangeCompareRenderer({ model }: SceneComponentProps<SceneTimeRangeCompare>) {
   const styles = useStyles2(getStyles);
   const { compareWith, compareOptions } = model.useState();

packages/scenes/src/index.ts

@@ -28,6 +28,7 @@ export { SceneTimeRange } from './core/SceneTimeRange';
 export { SceneTimeZoneOverride } from './core/SceneTimeZoneOverride';
 
 export { SceneQueryRunner, type QueryRunnerState } from './querying/SceneQueryRunner';
+export { type SupplementaryRequest, type SupplementaryRequestProvider, type ProcessorFunc } from './querying/SupplementaryRequestProvider';
 export { SceneDataLayerSet, SceneDataLayerSetBase } from './querying/SceneDataLayerSet';
 export { SceneDataLayerBase } from './querying/layers/SceneDataLayerBase';
 export { SceneDataLayerControls } from './querying/layers/SceneDataLayerControls';

packages/scenes/src/querying/SceneQueryRunner.test.ts

@@ -32,8 +32,10 @@ import { emptyPanelData } from '../core/SceneDataNode';
 import { GroupByVariable } from '../variables/groupby/GroupByVariable';
 import { SceneQueryController, SceneQueryStateControllerState } from '../behaviors/SceneQueryController';
 import { activateFullSceneTree } from '../utils/test/activateFullSceneTree';
-import { SceneDeactivationHandler } from '../core/types';
+import { SceneDeactivationHandler, SceneObjectState } from '../core/types';
 import { LocalValueVariable } from '../variables/variants/LocalValueVariable';
+import { SceneObjectBase } from '../core/SceneObjectBase';
+import { SupplementaryRequest, SupplementaryRequestProvider } from './SupplementaryRequestProvider';
 
 const getDataSourceMock = jest.fn().mockReturnValue({
   uid: 'test-uid',
@@ -97,6 +99,7 @@ const runRequestMock = jest.fn().mockImplementation((ds: DataSourceApi, request:
     state: LoadingState.Loading,
     series: [],
     annotations: [],
+    request,
     timeRange: request.range,
   };
 
@@ -1081,6 +1084,88 @@ describe('SceneQueryRunner', () => {
     });
   });
 
+  describe('supplementary requests', () => {
+    test('should run and rerun supplementary requests', async () => {
+      const timeRange = new SceneTimeRange({
+        from: '2023-08-24T05:00:00.000Z',
+        to: '2023-08-24T07:00:00.000Z',
+      });
+
+      const queryRunner = new SceneQueryRunner({
+        queries: [{ refId: 'A' }],
+      });
+      const provider = new TestSupplementaryRequestProvider({ foo: 1 }, true);
+      const scene = new EmbeddedScene({
+        $timeRange: timeRange,
+        $data: queryRunner,
+        controls: [provider],
+        body: new SceneCanvasText({ text: 'hello' }),
+      });
+
+      // activate the scene, which will also activate the provider
+      // and the provider will run the supplementary request
+      scene.activate();
+      await new Promise((r) => setTimeout(r, 1));
+
+      expect(runRequestMock.mock.calls.length).toEqual(2);
+      let runRequestCall = runRequestMock.mock.calls[0];
+      let supplementaryRunRequestCall = runRequestMock.mock.calls[1];
+      expect(runRequestCall[1].targets[0].refId).toEqual('A');
+      expect(supplementaryRunRequestCall[1].targets[0].refId).toEqual('Supplementary');
+      expect(supplementaryRunRequestCall[1].targets[0].foo).toEqual(1);
+
+      // change the state of the provider, which will trigger the activation
+      // handler to run the supplementary request again.
+      provider.setState({ foo: 2 });
+      await new Promise((r) => setTimeout(r, 1));
+
+      expect(runRequestMock.mock.calls.length).toEqual(4);
+      runRequestCall = runRequestMock.mock.calls[2];
+      supplementaryRunRequestCall = runRequestMock.mock.calls[3];
+      expect(runRequestCall[1].targets[0].refId).toEqual('A');
+      expect(supplementaryRunRequestCall[1].targets[0].refId).toEqual('Supplementary');
+      expect(supplementaryRunRequestCall[1].targets[0].foo).toEqual(2);
+    });
+
+    test('should not rerun supplementary requests when providers say not to', async () => {
+      const timeRange = new SceneTimeRange({
+        from: '2023-08-24T05:00:00.000Z',
+        to: '2023-08-24T07:00:00.000Z',
+      });
+
+      const queryRunner = new SceneQueryRunner({
+        queries: [{ refId: 'A' }],
+      });
+      const provider = new TestSupplementaryRequestProvider({ foo: 1 }, false);
+      const scene = new EmbeddedScene({
+        $timeRange: timeRange,
+        $data: queryRunner,
+        controls: [provider],
+        body: new SceneCanvasText({ text: 'hello' }),
+      });
+
+      // activate the scene, which will also activate the provider
+      // and the provider will run the supplementary request
+      scene.activate();
+      await new Promise((r) => setTimeout(r, 1));
+
+      expect(runRequestMock.mock.calls.length).toEqual(2);
+      let runRequestCall = runRequestMock.mock.calls[0];
+      let supplementaryRunRequestCall = runRequestMock.mock.calls[1];
+      expect(runRequestCall[1].targets[0].refId).toEqual('A');
+      expect(supplementaryRunRequestCall[1].targets[0].refId).toEqual('Supplementary');
+      expect(supplementaryRunRequestCall[1].targets[0].foo).toEqual(1);
+
+      // change the state of the provider, which will trigger the activation
+      // handler to run the supplementary request again. The provider will
+      // return false from shouldRun, so we should not see any more queries.
+      provider.setState({ foo: 2 });
+      await new Promise((r) => setTimeout(r, 1));
+
+      expect(runRequestMock.mock.calls.length).toEqual(2);
+    });
+  });
+
   describe('time frame comparison', () => {
     test('should run query with time range comparison', async () => {
       const timeRange = new SceneTimeRange({
@@ -2190,3 +2275,30 @@ class CustomDataSource extends RuntimeDataSource {
     return of({ data: [{ refId: 'A', fields: [{ name: 'time', type: FieldType.time, values: [123] }] }] });
   }
 }
+
+interface TestSupplementaryRequestProviderState extends SceneObjectState {
+  foo: number;
+}
+
+class TestSupplementaryRequestProvider extends SceneObjectBase<TestSupplementaryRequestProviderState> implements SupplementaryRequestProvider<{}> {
+  private _shouldRerun: boolean;
+
+  public constructor(state: { foo: number; }, shouldRerun: boolean) {
+    super(state);
+    this._shouldRerun = shouldRerun;
+  }
+
+  public getSupplementaryRequests(): SupplementaryRequest[] {
+    return [{
+      req: {
+        targets: [
+          // @ts-expect-error
+          { refId: 'Supplementary', foo: this.state.foo },
+        ],
+      }
+    }];
+  }
+  public shouldRerun(prev: {}, next: {}): boolean {
+    return this._shouldRerun;
+  }
+}