From e17fd5e8ad42d78752909ab90d9c6a73613efc30 Mon Sep 17 00:00:00 2001
From: Victor Marin <36818606+mdvictor@users.noreply.github.com>
Date: Tue, 11 Feb 2025 13:39:00 +0200
Subject: [PATCH] DashboardDS: Re-run dashboard queries within MixedDS on
 transformation reprocessing (#100370)

* fix scenario where results subscription is lost due to transformations

* fix
---
 .../DashboardDatasourceBehaviour.test.tsx     | 63 +++++++++++++++++++
 .../scene/DashboardDatasourceBehaviour.tsx    | 29 ++++++---
 2 files changed, 85 insertions(+), 7 deletions(-)

diff --git a/public/app/features/dashboard-scene/scene/DashboardDatasourceBehaviour.test.tsx b/public/app/features/dashboard-scene/scene/DashboardDatasourceBehaviour.test.tsx
index 1b183577b1f..59a27eacaad 100644
--- a/public/app/features/dashboard-scene/scene/DashboardDatasourceBehaviour.test.tsx
+++ b/public/app/features/dashboard-scene/scene/DashboardDatasourceBehaviour.test.tsx
@@ -6,6 +6,7 @@ import {
   DataSourceApi,
   DataSourceJsonData,
   DataSourceRef,
+  getDefaultTimeRange,
   LoadingState,
   PanelData,
 } from '@grafana/data';
@@ -593,6 +594,68 @@ describe('DashboardDatasourceBehaviour', () => {
       expect(spy).toHaveBeenCalled();
     });
   });
+
+  it('Should re-run query after transformations reprocess', async () => {
+    const sourcePanel = new VizPanel({
+      title: 'Panel A',
+      pluginId: 'table',
+      key: 'panel-1',
+      $data: new SceneDataTransformer({
+        transformations: [{ id: 'transformA', options: {} }],
+        $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]),
+    });
+
+    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();
+
+    // transformations are reprocessed (e.g. variable change) and data is updated so
+    // we re-run the queries in the dashboardDS panel because we lose the subscription
+    // in mixed DS scenario
+    (sourcePanel.state.$data as SceneDataTransformer).setState({
+      data: { state: LoadingState.Done, series: [], timeRange: getDefaultTimeRange() },
+    });
+
+    expect(spy).toHaveBeenCalled();
+  });
 });
 
 async function buildTestScene() {
diff --git a/public/app/features/dashboard-scene/scene/DashboardDatasourceBehaviour.tsx b/public/app/features/dashboard-scene/scene/DashboardDatasourceBehaviour.tsx
index fcbab5cc59b..940b4077d8d 100644
--- a/public/app/features/dashboard-scene/scene/DashboardDatasourceBehaviour.tsx
+++ b/public/app/features/dashboard-scene/scene/DashboardDatasourceBehaviour.tsx
@@ -1,6 +1,6 @@
 import { Unsubscribable } from 'rxjs';
 
-import { SceneObjectBase, SceneObjectState, SceneQueryRunner, VizPanel } from '@grafana/scenes';
+import { SceneDataTransformer, 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';
 
@@ -28,6 +28,7 @@ export class DashboardDatasourceBehaviour extends SceneObjectBase<DashboardDatas
   private _activationHandler() {
     const queryRunner = this.parent;
     let libraryPanelSub: Unsubscribable;
+    let transformerSub: Unsubscribable;
     let dashboard: DashboardScene;
     if (!(queryRunner instanceof SceneQueryRunner)) {
       throw new Error('DashboardDatasourceBehaviour must be attached to a SceneQueryRunner');
@@ -73,6 +74,20 @@ export class DashboardDatasourceBehaviour extends SceneObjectBase<DashboardDatas
       throw new Error('Could not find SceneQueryRunner for panel');
     }
 
+    const dataTransformer = sourcePanelQueryRunner.parent;
+
+    if (dataTransformer instanceof SceneDataTransformer && dataTransformer.state.transformations.length) {
+      // in mixed DS scenario we complete the observable and merge data, so on a variable change
+      // the data transformer will emit but there will be no subscription and thus not visual update
+      // on the panel. Similar thing happens when going to edit mode and back, where we unsubscribe and
+      // since we never re-run the query, only reprocess the transformations, the panel will not update.
+      transformerSub = dataTransformer.subscribeToState((newState, oldState) => {
+        if (newState.data !== oldState.data) {
+          queryRunner.runQueries();
+        }
+      });
+    }
+
     if (this.prevRequestId && this.prevRequestId !== sourcePanelQueryRunner.state.data?.request?.requestId) {
       queryRunner.runQueries();
     }
@@ -82,6 +97,10 @@ export class DashboardDatasourceBehaviour extends SceneObjectBase<DashboardDatas
       if (libraryPanelSub) {
         libraryPanelSub.unsubscribe();
       }
+
+      if (transformerSub) {
+        transformerSub.unsubscribe();
+      }
     };
   }
 
@@ -90,14 +109,10 @@ export class DashboardDatasourceBehaviour extends SceneObjectBase<DashboardDatas
       return true;
     }
 
-    if (
+    return (
       queryRunner.state.datasource?.uid === MIXED_DATASOURCE_NAME &&
       queryRunner.state.queries.some((query) => query.datasource?.uid === SHARED_DASHBOARD_QUERY)
-    ) {
-      return true;
-    }
-
-    return false;
+    );
   }
 
   private handleLibPanelStateUpdates(