DashboardDataSource: Improve handling of source provider activation, and add tests (#81034)

* DashboardDataSource: Improve handling of source provider activation, and add tests

* Add setContainerWidth call

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

@@ -0,0 +1,93 @@
+import {
+  arrayToDataFrame,
+  DataQueryResponse,
+  DataSourceInstanceSettings,
+  getDefaultTimeRange,
+  LoadingState,
+  standardTransformersRegistry,
+} from '@grafana/data';
+import { SceneDataNode, SceneDataTransformer, SceneFlexItem, SceneFlexLayout, VizPanel } from '@grafana/scenes';
+import { getVizPanelKeyForPanelId } from 'app/features/dashboard-scene/utils/utils';
+import { getStandardTransformers } from 'app/features/transformers/standardTransformers';
+
+import { DashboardDatasource } from './datasource';
+import { DashboardQuery } from './types';
+
+standardTransformersRegistry.setInit(getStandardTransformers);
+
+describe('DashboardDatasource', () => {
+  it("should look up the other panel and subscribe to it's data", async () => {
+    const { observable } = setup({ refId: 'A', panelId: 1 });
+
+    let rsp: DataQueryResponse | undefined;
+
+    observable.subscribe({ next: (data) => (rsp = data) });
+
+    expect(rsp?.data[0].fields[0].values).toEqual([1, 2, 3]);
+  });
+
+  it('Can subscribe to panel data + transforms', async () => {
+    const { observable } = setup({ refId: 'A', panelId: 1, withTransforms: true });
+
+    let rsp: DataQueryResponse | undefined;
+
+    observable.subscribe({ next: (data) => (rsp = data) });
+
+    expect(rsp?.data[0].fields[1].values).toEqual([3]);
+  });
+
+  it('Should activate source provder on observable subscribe and and deactivate when completed (if only activator)', async () => {
+    const { observable, sourceData } = setup({ refId: 'A', panelId: 1, withTransforms: true });
+
+    const test = observable.subscribe({ next: () => {} });
+
+    expect(sourceData.isActive).toBe(true);
+
+    test.unsubscribe();
+
+    expect(sourceData.isActive).toBe(false);
+  });
+});
+
+function setup(query: DashboardQuery) {
+  const sourceData = new SceneDataTransformer({
+    $data: new SceneDataNode({
+      data: {
+        series: [arrayToDataFrame([1, 2, 3])],
+        state: LoadingState.Done,
+        timeRange: getDefaultTimeRange(),
+        structureRev: 11,
+      },
+    }),
+    transformations: [{ id: 'reduce', options: {} }],
+  });
+
+  const scene = new SceneFlexLayout({
+    children: [
+      new SceneFlexItem({
+        body: new VizPanel({
+          key: getVizPanelKeyForPanelId(1),
+          $data: sourceData,
+        }),
+      }),
+    ],
+  });
+
+  const ds = new DashboardDatasource({} as DataSourceInstanceSettings);
+
+  const observable = ds.query({
+    timezone: 'utc',
+    targets: [query],
+    requestId: '',
+    interval: '',
+    intervalMs: 0,
+    range: getDefaultTimeRange(),
+    scopedVars: {
+      __sceneObject: { value: scene },
+    },
+    app: '',
+    startTime: 0,
+  });
+
+  return { observable, sourceData };
+}

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

@@ -1,4 +1,4 @@
-import { Observable, map, of } from 'rxjs';
+import { Observable, defer, finalize, map, of } from 'rxjs';
 
 import {
   DataSourceApi,
@@ -8,11 +8,7 @@ import {
   TestDataSourceResponse,
 } from '@grafana/data';
 import { SceneDataProvider, SceneDataTransformer, SceneObject } from '@grafana/scenes';
-import {
-  findVizPanelByKey,
-  getQueryRunnerFor,
-  getVizPanelKeyForPanelId,
-} from 'app/features/dashboard-scene/utils/utils';
+import { findVizPanelByKey, getVizPanelKeyForPanelId } from 'app/features/dashboard-scene/utils/utils';
 
 import { DashboardQuery } from './types';
 
@@ -52,34 +48,35 @@ export class DashboardDatasource extends DataSourceApi<DashboardQuery> {
       return of({ data: [], error: { message: 'Could not find source panel' } });
     }
 
-    let sourceDataProvider: SceneDataProvider | undefined = getQueryRunnerFor(sourcePanel);
+    let sourceDataProvider: SceneDataProvider | undefined = sourcePanel.state.$data;
+
+    if (!query.withTransforms && sourceDataProvider instanceof SceneDataTransformer) {
+      sourceDataProvider = sourceDataProvider.state.$data;
+    }
 
     if (!sourceDataProvider || !sourceDataProvider.getResultsStream) {
       return of({ data: [] });
     }
 
-    if (query.withTransforms && sourceDataProvider.parent) {
-      const transformer = sourceDataProvider.parent;
-      if (transformer && transformer instanceof SceneDataTransformer) {
-        sourceDataProvider = transformer;
+    return defer(() => {
+      if (!sourceDataProvider!.isActive && sourceDataProvider?.setContainerWidth) {
+        sourceDataProvider?.setContainerWidth(500);
       }
-    }
 
-    if (!sourceDataProvider?.isActive) {
-      sourceDataProvider?.activate();
-      sourceDataProvider.setContainerWidth!(500);
-    }
+      const cleanUp = sourceDataProvider!.activate();
 
-    return sourceDataProvider.getResultsStream!().pipe(
-      map((result) => {
-        return {
-          data: result.data.series,
-          state: result.data.state,
-          errors: result.data.errors,
-          error: result.data.error,
-        };
-      })
-    );
+      return sourceDataProvider!.getResultsStream!().pipe(
+        map((result) => {
+          return {
+            data: result.data.series,
+            state: result.data.state,
+            errors: result.data.errors,
+            error: result.data.error,
+          };
+        }),
+        finalize(cleanUp)
+      );
+    });
   }
 
   private findSourcePanel(scene: SceneObject, panelId: number) {