SceneQueryRunner: decouple time range comparisons

packages/scenes/src/components/SceneTimeRangeCompare.tsx

@@ -1,16 +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 { ExtraQueryDescriptor, ExtraQueryDataProcessor, ExtraQueryProvider } from '../querying/ExtraQueryProvider';
 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;
-}
+import { of } from 'rxjs';
 
 interface SceneTimeRangeCompareState extends SceneObjectState {
   compareWith?: string;
@@ -38,8 +39,8 @@ export const DEFAULT_COMPARE_OPTIONS = [
 
 export class SceneTimeRangeCompare
   extends SceneObjectBase<SceneTimeRangeCompareState>
-  implements TimeRangeCompareProvider
-{
+  implements ExtraQueryProvider<SceneTimeRangeCompareState> {
+
   static Component = SceneTimeRangeCompareRenderer;
   protected _urlSync = new SceneObjectUrlSyncConfig(this, { keys: ['compareWith'] });
 
@@ -94,6 +95,33 @@ export class SceneTimeRangeCompare
     this.setState({ compareWith: undefined });
   };
 
+  // Get a time shifted request to compare with the primary request.
+  public getExtraQueries(request: DataQueryRequest): ExtraQueryDescriptor[] {
+    const extraQueries: ExtraQueryDescriptor[] = [];
+    const compareRange = this.getCompareTimeRange(request.range);
+    if (!compareRange) {
+      return extraQueries;
+    }
+
+    const targets = request.targets.filter((query: DataQueryExtended) => query.timeRangeCompare !== false);
+    if (targets.length) {
+      extraQueries.push({
+        req: {
+          ...request,
+          targets,
+          range: compareRange,
+        },
+        processor: timeShiftAlignmentProcessor,
+      });
+    }
+    return extraQueries;
+  }
+
+  // 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 +177,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: ExtraQueryDataProcessor = (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 of(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 ExtraQueryDescriptor, type ExtraQueryProvider, type ExtraQueryDataProcessor } from './querying/ExtraQueryProvider';
 export { SceneDataLayerSet, SceneDataLayerSetBase } from './querying/SceneDataLayerSet';
 export { SceneDataLayerBase } from './querying/layers/SceneDataLayerBase';
 export { SceneDataLayerControls } from './querying/layers/SceneDataLayerControls';

packages/scenes/src/querying/ExtraQueryProvider.ts

@@ -0,0 +1,47 @@
+import { DataQueryRequest, PanelData } from "@grafana/data";
+import { Observable } from "rxjs";
+
+import { SceneObjectBase } from "../core/SceneObjectBase";
+import { SceneObjectState } from "../core/types";
+
+// A processor function called by the query runner with responses
+// to any extra requests.
+//
+// A processor function should accept two arguments: the data returned by the
+// _primary_ query, and the data returned by the `ExtraQueryProvider`'s
+// _secondary_ query. It should return a new `PanelData` representing the processed output.
+// It should _not_ modify the primary PanelData.
+//
+// Examples of valid processing include alignment of data between primary and secondary
+// (see the `timeShiftAlignmentProcessor` returned by `SceneTimeRangeCompare`), or doing
+// some more advanced processing such as fitting a time series model on the secondary data.
+//
+// See the docs for `extraQueryProcessingOperator` for more information.
+export type ExtraQueryDataProcessor = (primary: PanelData, secondary: PanelData) => Observable<PanelData>;
+
+// An extra request that should be run by a query runner, and an optional
+// processor that should be called with the response data.
+export interface ExtraQueryDescriptor {
+  // The extra request to add.
+  req: DataQueryRequest;
+  // An optional function used to process the data before passing it
+  // to any transformations or visualizations.
+  processor?: ExtraQueryDataProcessor;
+}
+
+// Indicates that this type wants to add extra requests, along with
+// optional processing functions, to a query runner.
+export interface ExtraQueryProvider<T extends SceneObjectState> extends SceneObjectBase<T> {
+  // Get any extra requests and their required processors.
+  getExtraQueries(request: DataQueryRequest): ExtraQueryDescriptor[];
+  // Determine whether a query should be rerun.
+  //
+  // When the provider's state changes this function will be passed both the previous and the
+  // next state. The implementation can use this to determine whether the change should trigger
+  // a rerun of the query or not.
+  shouldRerun(prev: T, next: T): boolean;
+}
+
+export function isExtraQueryProvider(obj: any): obj is ExtraQueryProvider<any> {
+  return typeof obj === 'object' && 'getExtraQueries' in obj;
+}

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 { ExtraQueryDescriptor, ExtraQueryProvider } from './ExtraQueryProvider';
 
 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('extra requests', () => {
+    test('should run and rerun extra 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 TestExtraQueryProvider({ 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 extra request
+      scene.activate();
+      await new Promise((r) => setTimeout(r, 1));
+
+      expect(runRequestMock.mock.calls.length).toEqual(2);
+      let runRequestCall = runRequestMock.mock.calls[0];
+      let extraRunRequestCall = runRequestMock.mock.calls[1];
+      expect(runRequestCall[1].targets[0].refId).toEqual('A');
+      expect(extraRunRequestCall[1].targets[0].refId).toEqual('Extra');
+      expect(extraRunRequestCall[1].targets[0].foo).toEqual(1);
+
+      // change the state of the provider, which will trigger the activation
+      // handler to run the extra 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];
+      extraRunRequestCall = runRequestMock.mock.calls[3];
+      expect(runRequestCall[1].targets[0].refId).toEqual('A');
+      expect(extraRunRequestCall[1].targets[0].refId).toEqual('Extra');
+      expect(extraRunRequestCall[1].targets[0].foo).toEqual(2);
+    });
+
+    test('should not rerun extra 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 TestExtraQueryProvider({ 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 extra request
+      scene.activate();
+      await new Promise((r) => setTimeout(r, 1));
+
+      expect(runRequestMock.mock.calls.length).toEqual(2);
+      let runRequestCall = runRequestMock.mock.calls[0];
+      let extraRunRequestCall = runRequestMock.mock.calls[1];
+      expect(runRequestCall[1].targets[0].refId).toEqual('A');
+      expect(extraRunRequestCall[1].targets[0].refId).toEqual('Extra');
+      expect(extraRunRequestCall[1].targets[0].foo).toEqual(1);
+
+      // change the state of the provider, which will trigger the activation
+      // handler to run the extra 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,31 @@ class CustomDataSource extends RuntimeDataSource {
     return of({ data: [{ refId: 'A', fields: [{ name: 'time', type: FieldType.time, values: [123] }] }] });
   }
 }
+
+interface TestExtraQueryProviderState extends SceneObjectState {
+  foo: number;
+}
+
+class TestExtraQueryProvider extends SceneObjectBase<TestExtraQueryProviderState> implements ExtraQueryProvider<{}> {
+  private _shouldRerun: boolean;
+
+  public constructor(state: { foo: number; }, shouldRerun: boolean) {
+    super(state);
+    this._shouldRerun = shouldRerun;
+  }
+
+  public getExtraQueries(): ExtraQueryDescriptor[] {
+    return [{
+      req: {
+        targets: [
+          // @ts-expect-error
+          { refId: 'Extra', foo: this.state.foo },
+        ],
+      },
+      processor: (primary, secondary) => of({ ...primary, ...secondary }),
+    }];
+  }
+  public shouldRerun(prev: {}, next: {}): boolean {
+    return this._shouldRerun;
+  }
+}