Saved Queries: Fix Change DS during replace query did not update DataSourcePicker (#110299)

* Pass onUpdateDatasources from PanelDataQueriesTab to QueryEditorRows to keep ds up to date * Add unit tests * add type ds to make the test clearer
2025-09-01 14:53:07 +02:00
parent 66a99d0ae8
commit af893344f2
3 changed files with 194 additions and 8 deletions
@@ -175,6 +175,9 @@ const MixedDs = {
    id: 'grafana',
    mixed: true,
  },
+  getRef: () => {
+    return { type: 'datasource', uid: '-- Mixed --' };
+  },
 };

 const MixedDsSettingsMock = {
@@ -760,6 +763,67 @@ describe('PanelDataQueriesTab', () => {
        });
      });
    });
+
+    describe('updateDatasourceIfNeeded', () => {
+      it('should update datasource when different datasource reference is provided', async () => {
+        const { queriesTab } = await setupScene('panel-1');
+
+        // Initially should have testdata datasource
+        expect(queriesTab.state.datasource?.uid).toBe('gdev-testdata');
+
+        // Call updateDatasourceIfNeeded with prometheus datasource
+        await queriesTab.updateDatasourceIfNeeded({ uid: 'gdev-prometheus' });
+
+        // Should update to prometheus datasource
+        expect(queriesTab.state.datasource?.uid).toBe('gdev-prometheus');
+        expect(queriesTab.state.dsSettings?.uid).toBe('gdev-prometheus');
+      });
+
+      it('should not update datasource when same datasource reference is provided', async () => {
+        const { queriesTab } = await setupScene('panel-1');
+
+        // Initially should have testdata datasource
+        expect(queriesTab.state.datasource?.uid).toBe('gdev-testdata');
+
+        const originalDatasource = queriesTab.state.datasource;
+        const originalDsSettings = queriesTab.state.dsSettings;
+
+        // Call updateDatasourceIfNeeded with same datasource
+        await queriesTab.updateDatasourceIfNeeded({ uid: 'gdev-testdata' });
+
+        // Should not change the datasource
+        expect(queriesTab.state.datasource).toBe(originalDatasource);
+        expect(queriesTab.state.dsSettings).toBe(originalDsSettings);
+      });
+
+      it('should update datasource to mixed when mixed datasource reference is provided', async () => {
+        const { queriesTab } = await setupScene('panel-1');
+
+        // Initially should have testdata datasource
+        expect(queriesTab.state.datasource?.uid).toBe('gdev-testdata');
+
+        // Call updateDatasourceIfNeeded with mixed datasource
+        await queriesTab.updateDatasourceIfNeeded({ uid: '-- Mixed --' });
+
+        // Should update to mixed datasource
+        expect(queriesTab.state.datasource?.uid).toBe('-- Mixed --');
+        expect(queriesTab.state.dsSettings?.uid).toBe('-- Mixed --');
+      });
+
+      it('should handle case when datasource instance settings are not found', async () => {
+        const { queriesTab } = await setupScene('panel-1');
+
+        // Initially should have testdata datasource
+        expect(queriesTab.state.datasource?.uid).toBe('gdev-testdata');
+
+        // Call updateDatasourceIfNeeded with non-existent datasource
+        await queriesTab.updateDatasourceIfNeeded({ uid: 'non-existent-ds' });
+
+        // Should fall back to default datasource (since mock returns default when not found)
+        expect(queriesTab.state.datasource?.uid).toBe('gdev-testdata');
+        expect(queriesTab.state.dsSettings?.uid).toBe('gdev-testdata');
+      });
+    });
  });
 });

@@ -12,7 +12,7 @@ import {
  SceneObjectState,
  SceneDataQuery,
 } from '@grafana/scenes';
-import { DataQuery } from '@grafana/schema';
+import { DataQuery, DataSourceRef } from '@grafana/schema';
 import { Button, Stack, Tab } from '@grafana/ui';
 import { addQuery } from 'app/core/utils/query';
 import { getLastUsedDatasourceFromStorage } from 'app/features/dashboard/utils/dashboard';
@@ -315,6 +315,21 @@ export class PanelDataQueriesTab extends SceneObjectBase<PanelDataQueriesTabStat
  public get queryRunner(): SceneQueryRunner {
    return getQueryRunnerFor(this.state.panelRef.resolve())!;
  }
+
+  /**
+   * Updates panel datasource when saved queries from different datasources are used.
+   * Handles transition between single and mixed datasource modes.
+   */
+  public updateDatasourceIfNeeded = async (newDatasourceRef: DataSourceRef): Promise<void> => {
+    const { datasource } = this.state;
+    const shouldChangeDatasource = datasource?.uid !== newDatasourceRef.uid;
+    if (shouldChangeDatasource) {
+      const newDatasource = getDatasourceSrv().getInstanceSettings(newDatasourceRef);
+      if (newDatasource) {
+        await this.onChangeDataSource(newDatasource);
+      }
+    }
+  };
 }

 export function PanelDataQueriesTabRendered({ model }: SceneComponentProps<PanelDataQueriesTab>) {
@@ -339,6 +354,7 @@ export function PanelDataQueriesTabRendered({ model }: SceneComponentProps<Panel
    );
    const newQueries = addQuery(enrichedQueries, query);
    model.onQueriesChange(newQueries);
+
    if (query.datasource?.uid) {
      const uniqueDatasources = new Set(
        newQueries.map((q) => q.datasource?.uid).filter((uid) => uid !== ExpressionDatasourceUID)
@@ -347,13 +363,7 @@ export function PanelDataQueriesTabRendered({ model }: SceneComponentProps<Panel
      const newDatasourceRef = {
        uid: isMixed ? MIXED_DATASOURCE_NAME : query.datasource.uid,
      };
-      const shouldChangeDatasource = datasource.uid !== newDatasourceRef.uid;
-      if (shouldChangeDatasource) {
-        const newDatasource = getDatasourceSrv().getInstanceSettings(newDatasourceRef);
-        if (newDatasource) {
-          await model.onChangeDataSource(newDatasource);
-        }
-      }
+      await model.updateDatasourceIfNeeded(newDatasourceRef);
    }
  };

@@ -376,6 +386,7 @@ export function PanelDataQueriesTabRendered({ model }: SceneComponentProps<Panel
        onAddQuery={model.onAddQuery}
        onQueriesChange={model.onQueriesChange}
        onRunQueries={model.onRunQueries}
+        onUpdateDatasources={queryLibraryEnabled ? model.updateDatasourceIfNeeded : undefined}
        app={CoreApp.PanelEditor}
      />

@@ -87,6 +87,117 @@ describe('QueryEditorRows', () => {
    );
  });

+  it('Should call onUpdateDatasources when replacing query with different datasource creates mixed scenario', () => {
+    const onQueriesChangeMock = jest.fn();
+    const onUpdateDatasourcesMock = jest.fn();
+
+    const testProps = {
+      ...props,
+      onQueriesChange: onQueriesChangeMock,
+      onUpdateDatasources: onUpdateDatasourcesMock,
+      dsSettings: { ...props.dsSettings, uid: 'current-datasource' },
+      queries: [
+        { datasource: { uid: 'current-datasource', type: 'alertmanager' }, refId: 'A' },
+        { datasource: { uid: 'current-datasource', type: 'alertmanager' }, refId: 'B' },
+      ],
+    };
+
+    const component = new QueryEditorRows(testProps);
+    const replacementQuery = {
+      refId: 'A',
+      datasource: { uid: 'different-datasource', type: 'prometheus' },
+      expr: 'new query content',
+    };
+
+    component.onReplaceQuery(replacementQuery, 0);
+
+    expect(onUpdateDatasourcesMock).toHaveBeenCalledWith({
+      uid: MIXED_DATASOURCE_NAME,
+    });
+  });
+
+  it('Should call onUpdateDatasources when replacing query results in single different datasource', () => {
+    const onQueriesChangeMock = jest.fn();
+    const onUpdateDatasourcesMock = jest.fn();
+
+    const testProps = {
+      ...props,
+      onQueriesChange: onQueriesChangeMock,
+      onUpdateDatasources: onUpdateDatasourcesMock,
+      dsSettings: { ...props.dsSettings, uid: 'current-datasource' },
+      queries: [{ datasource: { uid: 'current-datasource', type: 'alertmanager' }, refId: 'A' }],
+    };
+
+    const component = new QueryEditorRows(testProps);
+    const replacementQuery = {
+      refId: 'A',
+      datasource: { uid: 'different-datasource', type: 'prometheus' },
+      expr: 'new query content',
+    };
+
+    component.onReplaceQuery(replacementQuery, 0);
+
+    expect(onUpdateDatasourcesMock).toHaveBeenCalledWith({
+      uid: 'different-datasource',
+    });
+  });
+
+  it('Should not call onUpdateDatasources when replacing query with same datasource', () => {
+    const onQueriesChangeMock = jest.fn();
+    const onUpdateDatasourcesMock = jest.fn();
+
+    const testProps = {
+      ...props,
+      onQueriesChange: onQueriesChangeMock,
+      onUpdateDatasources: onUpdateDatasourcesMock,
+      dsSettings: { ...props.dsSettings, uid: 'same-datasource' },
+      queries: [
+        { datasource: { uid: 'same-datasource', type: 'prometheus' }, refId: 'A' },
+        { datasource: { uid: 'same-datasource', type: 'prometheus' }, refId: 'B' },
+      ],
+    };
+
+    const component = new QueryEditorRows(testProps);
+    const replacementQuery = {
+      refId: 'A',
+      datasource: { uid: 'same-datasource', type: 'prometheus' },
+      expr: 'new query content',
+    };
+
+    component.onReplaceQuery(replacementQuery, 0);
+
+    expect(onUpdateDatasourcesMock).not.toHaveBeenCalled();
+  });
+
+  it('Should call onUpdateDatasources with mixed datasource when replacing creates mixed scenario', () => {
+    const onQueriesChangeMock = jest.fn();
+    const onUpdateDatasourcesMock = jest.fn();
+
+    const testProps = {
+      ...props,
+      onQueriesChange: onQueriesChangeMock,
+      onUpdateDatasources: onUpdateDatasourcesMock,
+      dsSettings: { ...props.dsSettings, uid: 'current-datasource' },
+      queries: [
+        { datasource: { uid: 'datasource-1', type: 'loki' }, refId: 'A' },
+        { datasource: { uid: 'datasource-2', type: 'test-data' }, refId: 'B' },
+      ],
+    };
+
+    const component = new QueryEditorRows(testProps);
+    const replacementQuery = {
+      refId: 'A',
+      datasource: { uid: 'datasource-3', type: 'prometheus' },
+      expr: 'new query content',
+    };
+
+    component.onReplaceQuery(replacementQuery, 0);
+
+    expect(onUpdateDatasourcesMock).toHaveBeenCalledWith({
+      uid: MIXED_DATASOURCE_NAME,
+    });
+  });
+
  it('Should render queries', async () => {
    const {
      renderResult: { rerender },