Dashboards: Allow DashboardDS subqueries in MixedDS (#97116)

* Allow dashboardDS to run in mixedDS

* Make mixedDS panel wait for dashboardDS panel to load first

* cleanup

* cleanup

* refresh dashboardDS queries within mixedDS when source panel changes

* more tests

* fix

* fixes scenario where source returns an error

* do not allow dashboardDS references to mixedDS targets that contain other dashboardDS panels

* test

* lint

* Show invalid panels as invalid and with a message

* refactor

* avoid bunching shared dashboard queries

* skip instead of debouncing to avoid stale data

* debouce dashboard ds result stream when coming from mixed ds

* restore unnecessarily touched files

* fix import

* increase debounce interval value to account for slower machines

---------

Co-authored-by: Sergej-Vlasov <sergej.s.vlasov@gmail.com>

public/app/features/dashboard-scene/scene/DashboardDatasourceBehaviour.test.tsx

@@ -13,6 +13,7 @@ import { getPanelPlugin } from '@grafana/data/test/__mocks__/pluginMocks';
 import { setPluginImportUtils } from '@grafana/runtime';
 import { SceneDataTransformer, SceneFlexLayout, SceneQueryRunner, VizPanel } from '@grafana/scenes';
 import { SHARED_DASHBOARD_QUERY, DASHBOARD_DATASOURCE_PLUGIN_ID } from 'app/plugins/datasource/dashboard/constants';
+import { MIXED_DATASOURCE_NAME } from 'app/plugins/datasource/mixed/MixedDataSource';
 
 import { activateFullSceneTree } from '../utils/test-utils';
 
@@ -46,6 +47,18 @@ const dashboardDs: DataSourceApi = {
   },
 } as DataSourceApi<DataQuery, DataSourceJsonData, {}>;
 
+const mixedDs: DataSourceApi = {
+  meta: {
+    id: 'mixed',
+  },
+  name: MIXED_DATASOURCE_NAME,
+  type: MIXED_DATASOURCE_NAME,
+  uid: MIXED_DATASOURCE_NAME,
+  getRef: () => {
+    return { type: MIXED_DATASOURCE_NAME, uid: MIXED_DATASOURCE_NAME };
+  },
+} as DataSourceApi<DataQuery, DataSourceJsonData, {}>;
+
 setPluginImportUtils({
   importPanelPlugin: (id: string) => Promise.resolve(getPanelPlugin({})),
   getPanelPluginFromCache: (id: string) => undefined,
@@ -85,6 +98,10 @@ jest.mock('@grafana/runtime', () => ({
           return dashboardDs;
         }
 
+        if (ref.uid === MIXED_DATASOURCE_NAME) {
+          return mixedDs;
+        }
+
         return null;
       },
       getInstanceSettings: jest.fn().mockResolvedValue({ uid: 'ds1' }),
@@ -449,6 +466,7 @@ describe('DashboardDatasourceBehaviour', () => {
     });
 
     it('should wait for library panel to load before running queries', async () => {
+      jest.spyOn(console, 'error').mockImplementation();
       const libPanelBehavior = new LibraryPanelBehavior({
         isLoaded: false,
         title: 'Panel title',
@@ -510,6 +528,74 @@ describe('DashboardDatasourceBehaviour', () => {
       expect(spyRunQueries).toHaveBeenCalledTimes(1);
     });
   });
+
+  describe('DashboardDS within MixedDS', () => {
+    it('Should re-run query of MixedDS panel that contains a dashboardDS when source query re-runs', async () => {
+      jest.spyOn(console, 'error').mockImplementation();
+      const sourcePanel = new VizPanel({
+        title: 'Panel A',
+        pluginId: 'table',
+        key: 'panel-1',
+        $data: new SceneDataTransformer({
+          transformations: [],
+          $data: new SceneQueryRunner({
+            datasource: { uid: 'grafana' },
+            queries: [{ refId: 'A', queryType: 'randomWalk' }],
+          }),
+        }),
+      });
+
+      const dashboardDSPanel = new VizPanel({
+        title: 'Panel B',
+        pluginId: 'table',
+        key: 'panel-2',
+        $data: new SceneDataTransformer({
+          transformations: [],
+          $data: new SceneQueryRunner({
+            datasource: { uid: MIXED_DATASOURCE_NAME },
+            queries: [
+              {
+                datasource: { uid: SHARED_DASHBOARD_QUERY },
+                refId: 'B',
+                panelId: 1,
+              },
+            ],
+            $behaviors: [new DashboardDatasourceBehaviour({})],
+          }),
+        }),
+      });
+
+      const scene = new DashboardScene({
+        title: 'hello',
+        uid: 'dash-1',
+        meta: {
+          canEdit: true,
+        },
+        body: DefaultGridLayoutManager.fromVizPanels([sourcePanel, dashboardDSPanel]),
+      });
+
+      const sceneDeactivate = activateFullSceneTree(scene);
+
+      await new Promise((r) => setTimeout(r, 1));
+
+      // spy on runQueries that will be called by the behaviour
+      const spy = jest
+        .spyOn(dashboardDSPanel.state.$data!.state.$data as SceneQueryRunner, 'runQueries')
+        .mockImplementation();
+
+      // deactivate scene to mimic going into panel edit
+      sceneDeactivate();
+      // run source panel queries and update request ID
+      (sourcePanel.state.$data!.state.$data as SceneQueryRunner).runQueries();
+
+      await new Promise((r) => setTimeout(r, 1));
+
+      // activate scene to mimic coming back from panel edit
+      activateFullSceneTree(scene);
+
+      expect(spy).toHaveBeenCalled();
+    });
+  });
 });
 
 async function buildTestScene() {

public/app/features/dashboard-scene/scene/DashboardDatasourceBehaviour.tsx

@@ -2,6 +2,7 @@ import { Unsubscribable } from 'rxjs';
 
 import { SceneObjectBase, SceneObjectState, SceneQueryRunner, VizPanel } from '@grafana/scenes';
 import { SHARED_DASHBOARD_QUERY } from 'app/plugins/datasource/dashboard/constants';
+import { MIXED_DATASOURCE_NAME } from 'app/plugins/datasource/mixed/MixedDataSource';
 
 import {
   findVizPanelByKey,
@@ -25,24 +26,24 @@ export class DashboardDatasourceBehaviour extends SceneObjectBase<DashboardDatas
   }
 
   private _activationHandler() {
-    const dashboardDsQueryRunner = this.parent;
+    const queryRunner = this.parent;
     let libraryPanelSub: Unsubscribable;
     let dashboard: DashboardScene;
-    if (!(dashboardDsQueryRunner instanceof SceneQueryRunner)) {
+    if (!(queryRunner instanceof SceneQueryRunner)) {
       throw new Error('DashboardDatasourceBehaviour must be attached to a SceneQueryRunner');
     }
 
-    if (dashboardDsQueryRunner.state.datasource?.uid !== SHARED_DASHBOARD_QUERY) {
+    if (!this.containsDashboardDSQueries(queryRunner)) {
       return;
     }
 
     try {
-      dashboard = getDashboardSceneFor(dashboardDsQueryRunner);
+      dashboard = getDashboardSceneFor(queryRunner);
     } catch {
       return;
     }
 
-    const dashboardQuery = dashboardDsQueryRunner.state.queries.find((query) => query.panelId !== undefined);
+    const dashboardQuery = queryRunner.state.queries.find((query) => query.panelId !== undefined);
 
     if (!dashboardQuery) {
       return;
@@ -61,7 +62,7 @@ export class DashboardDatasourceBehaviour extends SceneObjectBase<DashboardDatas
     const libraryPanelBehaviour = getLibraryPanelBehavior(sourcePanel);
     if (libraryPanelBehaviour && !libraryPanelBehaviour.state.isLoaded) {
       libraryPanelSub = libraryPanelBehaviour.subscribeToState((newLibPanel) => {
-        this.handleLibPanelStateUpdates(newLibPanel, dashboardDsQueryRunner, sourcePanel);
+        this.handleLibPanelStateUpdates(newLibPanel, queryRunner, sourcePanel);
       });
       return;
     }
@@ -73,7 +74,7 @@ export class DashboardDatasourceBehaviour extends SceneObjectBase<DashboardDatas
     }
 
     if (this.prevRequestId && this.prevRequestId !== sourcePanelQueryRunner.state.data?.request?.requestId) {
-      dashboardDsQueryRunner.runQueries();
+      queryRunner.runQueries();
     }
 
     return () => {
@@ -84,6 +85,21 @@ export class DashboardDatasourceBehaviour extends SceneObjectBase<DashboardDatas
     };
   }
 
+  private containsDashboardDSQueries(queryRunner: SceneQueryRunner): boolean {
+    if (queryRunner.state.datasource?.uid === SHARED_DASHBOARD_QUERY) {
+      return true;
+    }
+
+    if (
+      queryRunner.state.datasource?.uid === MIXED_DATASOURCE_NAME &&
+      queryRunner.state.queries.some((query) => query.datasource?.uid === SHARED_DASHBOARD_QUERY)
+    ) {
+      return true;
+    }
+
+    return false;
+  }
+
   private handleLibPanelStateUpdates(
     newLibPanel: LibraryPanelBehaviorState,
     dashboardDsQueryRunner: SceneQueryRunner,

public/app/features/query/components/QueryEditorRowHeader.tsx

@@ -139,7 +139,13 @@ const renderDataSource = <TQuery extends DataQuery>(
 
   return (
     <div className={styles.itemWrapper}>
-      <DataSourcePicker variables={true} alerting={alerting} current={dataSource.name} onChange={onChangeDataSource} />
+      <DataSourcePicker
+        dashboard={true}
+        variables={true}
+        alerting={alerting}
+        current={dataSource.name}
+        onChange={onChangeDataSource}
+      />
     </div>
   );
 };

public/app/plugins/datasource/dashboard/DashboardQueryEditor.test.tsx

@@ -11,8 +11,9 @@ import {
   createDashboardModelFixture,
   createPanelSaveModel,
 } from '../../../features/dashboard/state/__fixtures__/dashboardFixtures';
+import { MIXED_DATASOURCE_NAME } from '../mixed/MixedDataSource';
 
-import { DashboardQueryEditor } from './DashboardQueryEditor';
+import { DashboardQueryEditor, INVALID_PANEL_DESCRIPTION } from './DashboardQueryEditor';
 import { SHARED_DASHBOARD_QUERY } from './constants';
 import { DashboardDatasource } from './datasource';
 
@@ -61,6 +62,21 @@ describe('DashboardQueryEditor', () => {
           type: 'timeseries',
           title: 'Another panel',
         }),
+        createPanelSaveModel({
+          datasource: {
+            uid: MIXED_DATASOURCE_NAME,
+          },
+          targets: [
+            {
+              datasource: {
+                uid: SHARED_DASHBOARD_QUERY,
+              },
+            },
+          ],
+          id: 3,
+          type: 'timeseries',
+          title: 'A mixed DS with dashboard DS query panel',
+        }),
         createPanelSaveModel({
           datasource: {
             uid: SHARED_DASHBOARD_QUERY,
@@ -95,7 +111,37 @@ describe('DashboardQueryEditor', () => {
     const anotherPanel = await screen.findByText('Another panel');
     expect(anotherPanel).toBeInTheDocument();
 
-    expect(screen.queryByText('A dashboard query panel')).not.toBeInTheDocument();
+    expect(screen.queryByText('A dashboard query panel')?.nextElementSibling).toHaveTextContent(
+      INVALID_PANEL_DESCRIPTION
+    );
+  });
+
+  it('does not show a panel with either SHARED_DASHBOARD_QUERY datasource or MixedDS with SHARED_DASHBOARD_QUERY as an option in the dropdown', async () => {
+    render(
+      <DashboardQueryEditor
+        datasource={{} as DashboardDatasource}
+        query={mockQueries[0]}
+        data={mockPanelData}
+        onChange={mockOnChange}
+        onRunQuery={mockOnRunQueries}
+      />
+    );
+    const select = screen.getByText('Choose panel');
+
+    await userEvent.click(select);
+
+    const myFirstPanel = await screen.findByText('My first panel');
+    expect(myFirstPanel).toBeInTheDocument();
+
+    const anotherPanel = await screen.findByText('Another panel');
+    expect(anotherPanel).toBeInTheDocument();
+
+    expect(screen.queryByText('A dashboard query panel')?.nextElementSibling).toHaveTextContent(
+      INVALID_PANEL_DESCRIPTION
+    );
+    expect(screen.queryByText('A mixed DS with dashboard DS query panel')?.nextElementSibling).toHaveTextContent(
+      INVALID_PANEL_DESCRIPTION
+    );
   });
 
   it('does not show the current panelInEdit as an option in the dropdown', async () => {
@@ -118,6 +164,8 @@ describe('DashboardQueryEditor', () => {
     const anotherPanel = await screen.findByText('Another panel');
     expect(anotherPanel).toBeInTheDocument();
 
-    expect(screen.queryByText('A dashboard query panel')).not.toBeInTheDocument();
+    expect(screen.queryByText('A dashboard query panel')?.nextElementSibling).toHaveTextContent(
+      INVALID_PANEL_DESCRIPTION
+    );
   });
 });

public/app/plugins/datasource/dashboard/DashboardQueryEditor.tsx

@@ -14,6 +14,8 @@ import { DashboardScene } from 'app/features/dashboard-scene/scene/DashboardScen
 import { getDatasourceSrv } from 'app/features/plugins/datasource_srv';
 import { filterPanelDataToQuery } from 'app/features/query/components/QueryEditorRow';
 
+import { MIXED_DATASOURCE_NAME } from '../mixed/MixedDataSource';
+
 import { SHARED_DASHBOARD_QUERY } from './constants';
 import { DashboardDatasource } from './datasource';
 import { DashboardQuery, ResultInfo } from './types';
@@ -39,6 +41,8 @@ const topics = [
   { label: 'Annotations', value: true, description: 'Include annotations as regular data' },
 ];
 
+export const INVALID_PANEL_DESCRIPTION = 'Contains a shared dashboard query';
+
 export function DashboardQueryEditor({ data, query, onChange, onRunQuery }: Props) {
   const { value: defaultDatasource } = useAsync(() => getDatasourceSrv().get());
 
@@ -104,6 +108,13 @@ export function DashboardQueryEditor({ data, query, onChange, onRunQuery }: Prop
     [query, onUpdateQuery]
   );
 
+  const isMixedDSWithDashboardQueries = (panel: PanelModel) => {
+    return (
+      panel.datasource?.uid === MIXED_DATASOURCE_NAME &&
+      panel.targets.some((t) => t.datasource?.uid === SHARED_DASHBOARD_QUERY)
+    );
+  };
+
   const getPanelDescription = useCallback(
     (panel: PanelModel): string => {
       const datasource = panel.datasource ?? defaultDatasource;
@@ -120,18 +131,24 @@ export function DashboardQueryEditor({ data, query, onChange, onRunQuery }: Prop
     () =>
       dashboard?.panels
         .filter(
-          (panel) =>
-            config.panels[panel.type] &&
-            panel.targets &&
-            !isPanelInEdit(panel.id, dashboard.panelInEdit?.id) &&
-            panel.datasource?.uid !== SHARED_DASHBOARD_QUERY
+          (panel) => config.panels[panel.type] && panel.targets && !isPanelInEdit(panel.id, dashboard.panelInEdit?.id)
         )
-        .map((panel) => ({
-          value: panel.id,
-          label: panel.title ?? 'Panel ' + panel.id,
-          description: getPanelDescription(panel),
-          imgUrl: config.panels[panel.type].info.logos.small,
-        })) ?? [],
+        .map((panel) => {
+          let description = getPanelDescription(panel);
+          let isDisabled = false;
+          if (panel.datasource?.uid === SHARED_DASHBOARD_QUERY || isMixedDSWithDashboardQueries(panel)) {
+            description = INVALID_PANEL_DESCRIPTION;
+            isDisabled = true;
+          }
+
+          return {
+            value: panel.id,
+            label: panel.title ?? 'Panel ' + panel.id,
+            imgUrl: config.panels[panel.type].info.logos.small,
+            description,
+            isDisabled,
+          };
+        }) ?? [],
     [dashboard, getPanelDescription]
   );
 

public/app/plugins/datasource/dashboard/datasource.test.ts

@@ -1,3 +1,5 @@
+import { first } from 'rxjs';
+
 import {
   arrayToDataFrame,
   DataQueryResponse,
@@ -19,9 +21,19 @@ import {
 import { getVizPanelKeyForPanelId } from 'app/features/dashboard-scene/utils/utils';
 import { getStandardTransformers } from 'app/features/transformers/standardTransformers';
 
+import { MIXED_REQUEST_PREFIX } from '../mixed/MixedDataSource';
+
 import { DashboardDatasource } from './datasource';
 import { DashboardQuery } from './types';
 
+jest.mock('rxjs', () => {
+  const original = jest.requireActual('rxjs');
+  return {
+    ...original,
+    first: jest.fn(original.first),
+  };
+});
+
 standardTransformersRegistry.setInit(getStandardTransformers);
 setPluginImportUtils({
   importPanelPlugin: (id: string) => Promise.resolve(getPanelPlugin({})),
@@ -29,6 +41,10 @@ setPluginImportUtils({
 });
 
 describe('DashboardDatasource', () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+  });
+
   it("should look up the other panel and subscribe to it's data", async () => {
     const { observable } = setup({ refId: 'A', panelId: 1 });
 
@@ -70,9 +86,25 @@ describe('DashboardDatasource', () => {
 
     expect(sourceData.isActive).toBe(false);
   });
+
+  it('Should emit only the first value and complete if used within MixedDS', async () => {
+    const { observable } = setup({ refId: 'A', panelId: 1 }, `${MIXED_REQUEST_PREFIX}1`);
+
+    observable.subscribe({ next: () => {} });
+
+    expect(first).toHaveBeenCalled();
+  });
+
+  it('Should not get the first emission if requestId does not contain the MixedDS prefix', async () => {
+    const { observable } = setup({ refId: 'A', panelId: 1 });
+
+    observable.subscribe({ next: () => {} });
+
+    expect(first).not.toHaveBeenCalled();
+  });
 });
 
-function setup(query: DashboardQuery) {
+function setup(query: DashboardQuery, requestId?: string) {
   const sourceData = new SceneDataTransformer({
     $data: new SceneDataNode({
       data: {
@@ -101,7 +133,7 @@ function setup(query: DashboardQuery) {
   const observable = ds.query({
     timezone: 'utc',
     targets: [query],
-    requestId: '',
+    requestId: requestId ?? '',
     interval: '',
     intervalMs: 0,
     range: getDefaultTimeRange(),

public/app/plugins/datasource/dashboard/datasource.ts

@@ -1,4 +1,4 @@
-import { Observable, defer, finalize, map, of } from 'rxjs';
+import { Observable, debounce, defer, finalize, first, interval, map, of } from 'rxjs';
 
 import {
   DataSourceApi,
@@ -10,6 +10,7 @@ import {
   DataTopic,
   PanelData,
   DataFrame,
+  LoadingState,
 } from '@grafana/data';
 import { SceneDataProvider, SceneDataTransformer, SceneObject } from '@grafana/scenes';
 import {
@@ -18,6 +19,8 @@ import {
   getVizPanelKeyForPanelId,
 } from 'app/features/dashboard-scene/utils/utils';
 
+import { MIXED_REQUEST_PREFIX } from '../mixed/MixedDataSource';
+
 import { DashboardQuery } from './types';
 
 /**
@@ -36,10 +39,6 @@ export class DashboardDatasource extends DataSourceApi<DashboardQuery> {
     const sceneScopedVar: ScopedVar | undefined = options.scopedVars?.__sceneObject;
     let scene: SceneObject | undefined = sceneScopedVar ? (sceneScopedVar.value.valueOf() as SceneObject) : undefined;
 
-    if (options.requestId.indexOf('mixed') > -1) {
-      throw new Error('Dashboard data source cannot be used with Mixed data source.');
-    }
-
     if (!scene) {
       throw new Error('Can only be called from a scene');
     }
@@ -88,6 +87,7 @@ export class DashboardDatasource extends DataSourceApi<DashboardQuery> {
             key: 'source-ds-provider',
           };
         }),
+        this.emitFirstLoadedDataIfMixedDS(options.requestId),
         finalize(() => cleanUp?.())
       );
     });
@@ -112,6 +112,45 @@ export class DashboardDatasource extends DataSourceApi<DashboardQuery> {
     return findVizPanelByKey(scene, getVizPanelKeyForPanelId(panelId));
   }
 
+  private emitFirstLoadedDataIfMixedDS(
+    requestId: string
+  ): (source: Observable<DataQueryResponse>) => Observable<DataQueryResponse> {
+    return (source: Observable<DataQueryResponse>) => {
+      if (requestId.includes(MIXED_REQUEST_PREFIX)) {
+        let count = 0;
+
+        return source.pipe(
+          /*
+           * We can have the following piped values scenarios:
+           * Loading -> Done         - initial load
+           * Done -> Loading -> Done - refresh
+           * Done                    - adding another query in editor
+           *
+           * When we see Done as a first element this is because of ReplaySubject in SceneQueryRunner
+           *
+           * we use first(...) below to emit correct result which is last value with Done/Error states
+           *
+           * to avoid emitting first Done/Error (due to ReplaySubject) we selectively debounce only first value with such states
+           */
+          debounce((val) => {
+            if ([LoadingState.Done, LoadingState.Error].includes(val.state!) && count === 0) {
+              count++;
+              // in the refresh scenario we need to debounce first Done/Error until Loading arrives
+              //   400ms here is a magic number that was sufficient enough with the 20x cpu throttle
+              //   this still might affect slower machines but the issue affects only panel view/edit modes
+              return interval(400);
+            }
+            count++;
+            return interval(0);
+          }),
+          first((val) => val.state === LoadingState.Done || val.state === LoadingState.Error)
+        );
+      }
+
+      return source;
+    };
+  }
+
   testDatasource(): Promise<TestDataSourceResponse> {
     return Promise.resolve({ message: '', status: '' });
   }

public/app/plugins/datasource/mixed/MixedDataSource.ts

@@ -15,9 +15,13 @@ import {
 import { getDataSourceSrv, getTemplateSrv, toDataQueryError } from '@grafana/runtime';
 import { CustomFormatterVariable } from '@grafana/scenes';
 
-export const MIXED_DATASOURCE_NAME = '-- Mixed --';
+import { SHARED_DASHBOARD_QUERY } from '../dashboard/constants';
 
-export const mixedRequestId = (queryIdx: number, requestId?: string) => `mixed-${queryIdx}-${requestId || ''}`;
+export const MIXED_DATASOURCE_NAME = '-- Mixed --';
+export const MIXED_REQUEST_PREFIX = 'mixed-';
+
+export const mixedRequestId = (queryIdx: number, requestId?: string) =>
+  `${MIXED_REQUEST_PREFIX}${queryIdx}-${requestId || ''}`;
 
 export interface BatchedQueries {
   datasource: Promise<DataSourceApi>;
@@ -45,7 +49,15 @@ export class MixedDatasource extends DataSourceApi<DataQuery> {
     const batches: BatchedQueries[] = [];
 
     for (const key in sets) {
-      batches.push(...this.getBatchesForQueries(sets[key], request));
+      // dashboard ds expects to have only 1 query with const query = options.targets[0]; therefore
+      //   we should not batch them together
+      if (key === SHARED_DASHBOARD_QUERY) {
+        sets[key].forEach((a) => {
+          batches.push(...this.getBatchesForQueries([a], request));
+        });
+      } else {
+        batches.push(...this.getBatchesForQueries(sets[key], request));
+      }
     }
 
     // Missing UIDs?