Dashboard - Schema V2 Stateless (ds ref independent) queries (#101661)

- Rename `initializeMapping` into `initializeElementMapping`
- Create new  `initializeDSReferencesMapping` to track queries without explicit datasources
- Fix panel datasource detection to use default datasource when none is specified
- Improve datasource resolution by searching for best match by query kind
- Add unit test

---------

Co-authored-by: alexandra vargas <alexa1866@gmail.com>
Co-authored-by: Alexa V <239999+axelavargas@users.noreply.github.com>
Co-authored-by: Ivan Ortega Alba <ivanortegaalba@gmail.com>

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

@@ -686,7 +686,8 @@ export class DashboardScene extends SceneObjectBase<DashboardSceneState> impleme
     saveModel?: Dashboard | DashboardV2Spec,
     meta?: DashboardMeta | DashboardWithAccessInfo<DashboardV2Spec>['metadata']
   ): void {
-    this.serializer.initializeMapping(saveModel);
+    this.serializer.initializeElementMapping(saveModel);
+    this.serializer.initializeDSReferencesMapping(saveModel);
     const sortedModel = sortedDeepCloneWithoutNulls(saveModel);
     this.serializer.initialSaveModel = sortedModel;
     this.serializer.metadata = meta;

public/app/features/dashboard-scene/serialization/DashboardSceneSerializer.test.ts

@@ -392,7 +392,7 @@ describe('DashboardSceneSerializer', () => {
           ],
         };
 
-        serializer.initializeMapping(saveModel);
+        serializer.initializeElementMapping(saveModel);
         const mapping = serializer.getElementPanelMapping();
 
         expect(mapping.size).toBe(2);
@@ -401,10 +401,10 @@ describe('DashboardSceneSerializer', () => {
       });
 
       it('should handle empty or undefined panels in initializeMapping', () => {
-        serializer.initializeMapping(undefined);
+        serializer.initializeElementMapping(undefined);
         expect(serializer.getElementPanelMapping().size).toBe(0);
 
-        serializer.initializeMapping({
+        serializer.initializeElementMapping({
           title: 'hello',
           uid: 'my-uid',
           schemaVersion: 30,
@@ -424,7 +424,7 @@ describe('DashboardSceneSerializer', () => {
           ],
         };
 
-        serializer.initializeMapping(saveModel);
+        serializer.initializeElementMapping(saveModel);
 
         expect(serializer.getPanelIdForElement('panel-1')).toBe(1);
         expect(serializer.getPanelIdForElement('panel-2')).toBe(2);
@@ -441,7 +441,7 @@ describe('DashboardSceneSerializer', () => {
             { id: 2, title: 'Panel 2', type: 'text' },
           ],
         };
-        serializer.initializeMapping(saveModel);
+        serializer.initializeElementMapping(saveModel);
 
         expect(serializer.getElementIdForPanel(1)).toBe('panel-1');
         expect(serializer.getElementIdForPanel(2)).toBe('panel-2');
@@ -923,7 +923,7 @@ describe('DashboardSceneSerializer', () => {
       });
 
       it('should initialize panel mapping correctly', () => {
-        serializer.initializeMapping(saveModel);
+        serializer.initializeElementMapping(saveModel);
         const mapping = serializer.getElementPanelMapping();
 
         expect(mapping.size).toBe(2);
@@ -932,15 +932,15 @@ describe('DashboardSceneSerializer', () => {
       });
 
       it('should handle empty or undefined elements in initializeMapping', () => {
-        serializer.initializeMapping({} as DashboardV2Spec);
+        serializer.initializeElementMapping({} as DashboardV2Spec);
         expect(serializer.getElementPanelMapping().size).toBe(0);
 
-        serializer.initializeMapping({ elements: {} } as DashboardV2Spec);
+        serializer.initializeElementMapping({ elements: {} } as DashboardV2Spec);
         expect(serializer.getElementPanelMapping().size).toBe(0);
       });
 
       it('should get panel id for element correctly', () => {
-        serializer.initializeMapping(saveModel);
+        serializer.initializeElementMapping(saveModel);
 
         expect(serializer.getPanelIdForElement('element-panel-a')).toBe(1);
         expect(serializer.getPanelIdForElement('element-panel-b')).toBe(2);
@@ -948,7 +948,7 @@ describe('DashboardSceneSerializer', () => {
       });
 
       it('should get element id for panel correctly', () => {
-        serializer.initializeMapping(saveModel);
+        serializer.initializeElementMapping(saveModel);
 
         expect(serializer.getElementIdForPanel(1)).toBe('element-panel-a');
         expect(serializer.getElementIdForPanel(2)).toBe('element-panel-b');
@@ -958,6 +958,147 @@ describe('DashboardSceneSerializer', () => {
     });
   });
 
+  describe('Datasource References Mapping', () => {
+    describe('V2DashboardSerializer', () => {
+      let serializer: V2DashboardSerializer;
+
+      beforeEach(() => {
+        serializer = new V2DashboardSerializer();
+      });
+
+      it('should initialize datasource references mapping correctly for panels with undefined datasources', () => {
+        const saveModel: DashboardV2Spec = {
+          ...defaultDashboardV2Spec(),
+          title: 'Test Dashboard',
+          elements: {
+            'panel-1': {
+              kind: 'Panel',
+              spec: {
+                id: 1,
+                title: 'Panel 1',
+                description: '',
+                links: [],
+                vizConfig: {
+                  kind: 'timeseries',
+                  spec: {
+                    pluginVersion: '1.0.0',
+                    options: {},
+                    fieldConfig: { defaults: {}, overrides: [] },
+                  },
+                },
+                data: {
+                  kind: 'QueryGroup',
+                  spec: {
+                    queries: [
+                      {
+                        kind: 'PanelQuery',
+                        spec: {
+                          refId: 'A',
+                          hidden: false,
+                          // No datasource defined
+                          query: { kind: 'sql', spec: {} },
+                        },
+                      },
+                      {
+                        kind: 'PanelQuery',
+                        spec: {
+                          refId: 'B',
+                          hidden: false,
+                          datasource: { uid: 'datasource-1', type: 'prometheus' },
+                          query: { kind: 'prometheus', spec: {} },
+                        },
+                      },
+                    ],
+                    queryOptions: {},
+                    transformations: [],
+                  },
+                },
+              },
+            },
+            'panel-2': {
+              kind: 'Panel',
+              spec: {
+                id: 2,
+                title: 'Panel 2',
+                description: '',
+                links: [],
+                vizConfig: {
+                  kind: 'timeseries',
+                  spec: {
+                    pluginVersion: '1.0.0',
+                    options: {},
+                    fieldConfig: { defaults: {}, overrides: [] },
+                  },
+                },
+                data: {
+                  kind: 'QueryGroup',
+                  spec: {
+                    queries: [
+                      {
+                        kind: 'PanelQuery',
+                        spec: {
+                          refId: 'C',
+                          hidden: false,
+                          // No datasource defined
+                          query: { kind: 'sql', spec: {} },
+                        },
+                      },
+                    ],
+                    queryOptions: {},
+                    transformations: [],
+                  },
+                },
+              },
+            },
+          },
+        };
+
+        serializer.initializeElementMapping(saveModel);
+        serializer.initializeDSReferencesMapping(saveModel);
+
+        const dsReferencesMap = serializer.getDSReferencesMapping();
+
+        // Panel 1 should have refId A in the map (no datasource)
+        expect(dsReferencesMap.panels.has('panel-1')).toBe(true);
+        expect(dsReferencesMap.panels.get('panel-1')?.has('A')).toBe(true);
+        expect(dsReferencesMap.panels.get('panel-1')?.has('B')).toBe(false); // Has datasource defined
+
+        // Panel 2 should have refId C in the map
+        expect(dsReferencesMap.panels.has('panel-2')).toBe(true);
+        expect(dsReferencesMap.panels.get('panel-2')?.has('C')).toBe(true);
+      });
+
+      it('should handle empty or undefined elements in initializeDSReferencesMapping', () => {
+        serializer.initializeDSReferencesMapping(undefined);
+        expect(serializer.getDSReferencesMapping().panels.size).toBe(0);
+
+        serializer.initializeDSReferencesMapping({} as DashboardV2Spec);
+        expect(serializer.getDSReferencesMapping().panels.size).toBe(0);
+
+        serializer.initializeDSReferencesMapping({ elements: {} } as DashboardV2Spec);
+        expect(serializer.getDSReferencesMapping().panels.size).toBe(0);
+      });
+    });
+
+    describe('V1DashboardSerializer', () => {
+      let serializer: V1DashboardSerializer;
+
+      beforeEach(() => {
+        serializer = new V1DashboardSerializer();
+      });
+
+      it('should return empty mapping object for V1 serializer', () => {
+        serializer.initializeDSReferencesMapping(undefined);
+        expect(serializer.getDSReferencesMapping()).toEqual({
+          panels: expect.any(Map),
+          variables: expect.any(Set),
+          annotations: expect.any(Set),
+        });
+        expect(serializer.getDSReferencesMapping().panels.size).toBe(0);
+      });
+    });
+  });
+
   describe('onSaveComplete', () => {
     it('should set the initialSaveModel correctly', () => {
       const serializer = new V2DashboardSerializer();

public/app/features/dashboard-scene/serialization/DashboardSceneSerializer.ts

@@ -30,7 +30,8 @@ export interface DashboardSceneSerializerLike<T, M, I = T> {
    */
   initialSaveModel?: I;
   metadata?: M;
-  initializeMapping(saveModel: T | undefined): void;
+  initializeElementMapping(saveModel: T | undefined): void;
+  initializeDSReferencesMapping(saveModel: T | undefined): void;
   getSaveModel: (s: DashboardScene) => T;
   getSaveAsModel: (s: DashboardScene, options: SaveDashboardAsOptions) => T;
   getDashboardChangesFromScene: (
@@ -47,6 +48,7 @@ export interface DashboardSceneSerializerLike<T, M, I = T> {
   getPanelIdForElement: (elementId: string) => number | undefined;
   getElementIdForPanel: (panelId: number) => string | undefined;
   getElementPanelMapping: () => Map<string, number>;
+  getDSReferencesMapping: () => DSReferencesMapping;
 }
 
 interface DashboardTrackingInfo {
@@ -58,12 +60,23 @@ interface DashboardTrackingInfo {
   settings_livenow?: boolean;
 }
 
+interface DSReferencesMapping {
+  panels: Map<string, Set<string>>;
+  variables: Set<string>;
+  annotations: Set<string>;
+}
+
 export class V1DashboardSerializer implements DashboardSceneSerializerLike<Dashboard, DashboardMeta> {
   initialSaveModel?: Dashboard;
   metadata?: DashboardMeta;
   protected elementPanelMap = new Map<string, number>();
+  protected defaultDsReferencesMap = {
+    panels: new Map<string, Set<string>>(), // refIds as keys
+    variables: new Set<string>(), // variable names as keys
+    annotations: new Set<string>(), // annotation names as keys
+  };
 
-  initializeMapping(saveModel: Dashboard | undefined) {
+  initializeElementMapping(saveModel: Dashboard | undefined) {
     this.elementPanelMap.clear();
 
     if (!saveModel || !saveModel.panels) {
@@ -81,6 +94,15 @@ export class V1DashboardSerializer implements DashboardSceneSerializerLike<Dashb
     return this.elementPanelMap;
   }
 
+  initializeDSReferencesMapping(saveModel: Dashboard | undefined) {
+    // To be implemented in a different PR
+    return {};
+  }
+
+  getDSReferencesMapping() {
+    return this.defaultDsReferencesMap;
+  }
+
   getPanelIdForElement(elementId: string) {
     return this.elementPanelMap.get(elementId);
   }
@@ -185,12 +207,18 @@ export class V2DashboardSerializer
   initialSaveModel?: DashboardV2Spec | Dashboard;
   metadata?: DashboardWithAccessInfo<DashboardV2Spec>['metadata'];
   protected elementPanelMap = new Map<string, number>();
+  // map of elementId that will contain all the queries, variables and annotations that dont have a ds defined
+  protected defaultDsReferencesMap = {
+    panels: new Map<string, Set<string>>(), // refIds as keys
+    variables: new Set<string>(), // variable names as keys
+    annotations: new Set<string>(), // annotation names as keys
+  };
 
   getElementPanelMapping() {
     return this.elementPanelMap;
   }
 
-  initializeMapping(saveModel: DashboardV2Spec | undefined) {
+  initializeElementMapping(saveModel: DashboardV2Spec | undefined) {
     this.elementPanelMap.clear();
 
     if (!saveModel || !saveModel.elements) {
@@ -206,6 +234,42 @@ export class V2DashboardSerializer
     });
   }
 
+  initializeDSReferencesMapping(saveModel: DashboardV2Spec | undefined) {
+    // initialize the object
+    this.defaultDsReferencesMap = {
+      panels: new Map<string, Set<string>>(),
+      variables: new Set<string>(),
+      annotations: new Set<string>(),
+    };
+
+    // get all the element keys
+    const elementKeys = Object.keys(saveModel?.elements || {});
+    elementKeys.forEach((key) => {
+      const elementPanel = saveModel?.elements[key];
+      if (elementPanel?.kind === 'Panel') {
+        // check if the elementPanel.spec.datasource is defined
+        const panelQueries = elementPanel.spec.data.spec.queries;
+
+        for (const query of panelQueries) {
+          if (!query.spec.datasource) {
+            const elementId = this.getElementIdForPanel(elementPanel.spec.id);
+            if (!this.defaultDsReferencesMap.panels.has(elementId)) {
+              this.defaultDsReferencesMap.panels.set(elementId, new Set());
+            }
+
+            const panelDsqueries = this.defaultDsReferencesMap.panels.get(elementId)!;
+
+            panelDsqueries.add(query.spec.refId);
+          }
+        }
+      }
+    });
+  }
+
+  getDSReferencesMapping() {
+    return this.defaultDsReferencesMap;
+  }
+
   getPanelIdForElement(elementId: string) {
     return this.elementPanelMap.get(elementId);
   }

public/app/features/dashboard-scene/serialization/layoutSerializers/utils.ts

@@ -182,13 +182,21 @@ function getPanelDataSource(panel: PanelKind): DataSourceRef | undefined {
 
   panel.spec.data.spec.queries.forEach((query) => {
     if (!datasource) {
-      datasource = query.spec.datasource;
+      if (!query.spec.datasource?.uid) {
+        const defaultDatasource = config.bootData.settings.defaultDatasource;
+        const dsList = config.bootData.settings.datasources;
+        // this is look up by type
+        const bestGuess = Object.values(dsList).find((ds) => ds.meta.id === query.spec.query.kind);
+        datasource = bestGuess ? { uid: bestGuess.uid, type: bestGuess.meta.id } : dsList[defaultDatasource];
+      } else {
+        datasource = query.spec.datasource;
+      }
     } else if (datasource.uid !== query.spec.datasource?.uid || datasource.type !== query.spec.datasource?.type) {
       isMixedDatasource = true;
     }
   });
 
-  return isMixedDatasource ? { type: 'mixed', uid: MIXED_DATASOURCE_NAME } : undefined;
+  return isMixedDatasource ? { type: 'mixed', uid: MIXED_DATASOURCE_NAME } : datasource;
 }
 
 function panelQueryKindToSceneQuery(query: PanelQueryKind): SceneDataQuery {

public/app/features/dashboard-scene/serialization/transformSaveModelSchemaV2ToScene.ts

@@ -11,10 +11,6 @@ import {
   IntervalVariable,
   QueryVariable,
   SceneDataLayerControls,
-  SceneDataProvider,
-  SceneDataQuery,
-  SceneDataTransformer,
-  SceneQueryRunner,
   SceneRefreshPicker,
   SceneTimePicker,
   SceneTimeRange,
@@ -23,7 +19,6 @@ import {
   TextBoxVariable,
   VariableValueSelectors,
 } from '@grafana/scenes';
-import { DataSourceRef } from '@grafana/schema/dist/esm/index.gen';
 import {
   AdhocVariableKind,
   ConstantVariableKind,
@@ -42,7 +37,6 @@ import {
   IntervalVariableKind,
   LibraryPanelKind,
   PanelKind,
-  PanelQueryKind,
   QueryVariableKind,
   TextVariableKind,
 } from '@grafana/schema/src/schema/dashboard/v2alpha0';
@@ -55,14 +49,12 @@ import {
   DeprecatedInternalId,
 } from 'app/features/apiserver/types';
 import { DashboardWithAccessInfo } from 'app/features/dashboard/api/types';
-import { MIXED_DATASOURCE_NAME } from 'app/plugins/datasource/mixed/MixedDataSource';
 import { DashboardMeta } from 'app/types';
 
 import { addPanelsOnLoadBehavior } from '../addToDashboard/addPanelsOnLoadBehavior';
 import { DashboardAnnotationsDataLayer } from '../scene/DashboardAnnotationsDataLayer';
 import { DashboardControls } from '../scene/DashboardControls';
 import { DashboardDataLayerSet } from '../scene/DashboardDataLayerSet';
-import { DashboardDatasourceBehaviour } from '../scene/DashboardDatasourceBehaviour';
 import { registerDashboardMacro } from '../scene/DashboardMacro';
 import { DashboardReloadBehavior } from '../scene/DashboardReloadBehavior';
 import { DashboardScene } from '../scene/DashboardScene';
@@ -218,71 +210,6 @@ export function transformSaveModelSchemaV2ToScene(dto: DashboardWithAccessInfo<D
   return dashboardScene;
 }
 
-function getPanelDataSource(panel: PanelKind): DataSourceRef | undefined {
-  if (!panel.spec.data?.spec.queries?.length) {
-    return undefined;
-  }
-
-  let datasource: DataSourceRef | undefined = undefined;
-  let isMixedDatasource = false;
-
-  panel.spec.data.spec.queries.forEach((query) => {
-    if (!datasource) {
-      datasource = query.spec.datasource;
-    } else if (datasource.uid !== query.spec.datasource?.uid || datasource.type !== query.spec.datasource?.type) {
-      isMixedDatasource = true;
-    }
-  });
-
-  return isMixedDatasource ? { type: 'mixed', uid: MIXED_DATASOURCE_NAME } : datasource;
-}
-
-function panelQueryKindToSceneQuery(query: PanelQueryKind): SceneDataQuery {
-  return {
-    refId: query.spec.refId,
-    datasource: query.spec.datasource,
-    hide: query.spec.hidden,
-    ...query.spec.query.spec,
-  };
-}
-
-export function createPanelDataProvider(panelKind: PanelKind): SceneDataProvider | undefined {
-  const panel = panelKind.spec;
-  const targets = panel.data?.spec.queries ?? [];
-  // Skip setting query runner for panels without queries
-  if (!targets?.length) {
-    return undefined;
-  }
-
-  // Skip setting query runner for panel plugins with skipDataQuery
-  if (config.panels[panel.vizConfig.kind]?.skipDataQuery) {
-    return undefined;
-  }
-
-  let dataProvider: SceneDataProvider | undefined = undefined;
-  const datasource = getPanelDataSource(panelKind);
-
-  dataProvider = new SceneQueryRunner({
-    datasource,
-    queries: targets.map(panelQueryKindToSceneQuery),
-    maxDataPoints: panel.data.spec.queryOptions.maxDataPoints ?? undefined,
-    maxDataPointsFromWidth: true,
-    cacheTimeout: panel.data.spec.queryOptions.cacheTimeout,
-    queryCachingTTL: panel.data.spec.queryOptions.queryCachingTTL,
-    minInterval: panel.data.spec.queryOptions.interval ?? undefined,
-    dataLayerFilter: {
-      panelId: panel.id,
-    },
-    $behaviors: [new DashboardDatasourceBehaviour({})],
-  });
-
-  // Wrap inner data provider in a data transformer
-  return new SceneDataTransformer({
-    $data: dataProvider,
-    transformations: panel.data.spec.transformations.map((transformation) => transformation.spec),
-  });
-}
-
 function getVariables(dashboard: DashboardV2Spec, isSnapshot: boolean): SceneVariableSet | undefined {
   let variables: SceneVariableSet | undefined;
 

public/app/features/dashboard-scene/serialization/transformSceneToSaveModelSchemaV2.test.ts

@@ -17,6 +17,8 @@ import {
   SceneVariableSet,
   TextBoxVariable,
   VizPanel,
+  SceneDataQuery,
+  SceneQueryRunner,
 } from '@grafana/scenes';
 import {
   DashboardCursorSync as DashboardCursorSyncV1,
@@ -48,7 +50,38 @@ import { TabItem } from '../scene/layout-tabs/TabItem';
 import { TabsLayoutManager } from '../scene/layout-tabs/TabsLayoutManager';
 import { DashboardLayoutManager } from '../scene/types/DashboardLayoutManager';
 
-import { transformSceneToSaveModelSchemaV2 } from './transformSceneToSaveModelSchemaV2';
+import { getPersistedDSForQuery, transformSceneToSaveModelSchemaV2 } from './transformSceneToSaveModelSchemaV2';
+
+// Mock dependencies
+jest.mock('../utils/dashboardSceneGraph', () => {
+  const original = jest.requireActual('../utils/dashboardSceneGraph');
+  return {
+    ...original,
+    dashboardSceneGraph: {
+      ...original.dashboardSceneGraph,
+      getElementIdentifierForVizPanel: jest.fn().mockImplementation((panel) => {
+        // Return the panel key if it exists, otherwise use panel-1 as default
+        return panel?.state?.key || 'panel-1';
+      }),
+    },
+  };
+});
+
+jest.mock('../utils/utils', () => {
+  const original = jest.requireActual('../utils/utils');
+  return {
+    ...original,
+    getDashboardSceneFor: jest.fn().mockImplementation(() => ({
+      serializer: {
+        getDSReferencesMapping: jest.fn().mockReturnValue({
+          panels: new Map([['panel-1', new Set(['A'])]]),
+          variables: new Set(),
+          annotations: new Set(),
+        }),
+      },
+    })),
+  };
+});
 
 function setupDashboardScene(state: Partial<DashboardSceneState>): DashboardScene {
   return new DashboardScene(state);
@@ -100,7 +133,7 @@ describe('transformSceneToSaveModelSchemaV2', () => {
     // The intention is to have a complete dashboard scene
     // with all the possible properties set
     dashboardScene = setupDashboardScene({
-      $data: new DashboardDataLayerSet({ annotationLayers }),
+      $data: new DashboardDataLayerSet({ annotationLayers: createAnnotationLayers() }),
       id: 1,
       title: 'Test Dashboard',
       description: 'Test Description',
@@ -377,6 +410,84 @@ describe('transformSceneToSaveModelSchemaV2', () => {
     // check annotation layer 3 with no datasource has the default datasource defined as type
     expect(result.annotations?.[2].spec.datasource?.type).toBe('loki');
   });
+
+  describe('getPersistedDSForQuery', () => {
+    it('should respect datasource reference mapping when determining query datasource', () => {
+      // Setup test data
+      const queryWithoutDS: SceneDataQuery = {
+        refId: 'A',
+        // No datasource defined originally
+      };
+      const queryWithDS: SceneDataQuery = {
+        refId: 'B',
+        datasource: { uid: 'prometheus', type: 'prometheus' },
+      };
+
+      // Mock query runner with runtime-resolved datasource
+      const queryRunner = new SceneQueryRunner({
+        queries: [queryWithoutDS, queryWithDS],
+        datasource: { uid: 'default-ds', type: 'default' },
+      });
+
+      // Get a reference to the DS references mapping
+      const dsReferencesMap = new Set(['A']);
+
+      // Test the query without DS originally - should return undefined
+      const resultA = getPersistedDSForQuery(queryWithoutDS, queryRunner, dsReferencesMap);
+      expect(resultA).toBeUndefined();
+
+      // Test the query with DS originally - should return the original datasource
+      const resultB = getPersistedDSForQuery(queryWithDS, queryRunner, dsReferencesMap);
+      expect(resultB).toEqual({ uid: 'prometheus', type: 'prometheus' });
+
+      // Test a query with no DS originally but not in the mapping - should get the runner's datasource
+      const queryNotInMapping: SceneDataQuery = {
+        refId: 'C',
+        // No datasource, but not in mapping
+      };
+      const resultC = getPersistedDSForQuery(queryNotInMapping, queryRunner, dsReferencesMap);
+      expect(resultC).toEqual({ uid: 'default-ds', type: 'default' });
+    });
+  });
+
+  describe('getDatasourceForQueries', () => {
+    it('should respect datasource reference mapping when determining which queries should have datasources saved', () => {
+      // Setup test data
+      const queryWithoutDS: SceneDataQuery = {
+        refId: 'A',
+        // No datasource defined originally
+      };
+      const queryWithDS: SceneDataQuery = {
+        refId: 'B',
+        datasource: { uid: 'prometheus', type: 'prometheus' },
+      };
+
+      // Mock query runner with runtime-resolved datasource
+      const queryRunner = new SceneQueryRunner({
+        queries: [queryWithoutDS, queryWithDS],
+        datasource: { uid: 'default-ds', type: 'default' },
+      });
+
+      // Get a reference to the DS references mapping
+      const dsReferencesMap = new Set(['A']);
+
+      // Test the query without DS originally - should return undefined
+      const resultA = getPersistedDSForQuery(queryWithoutDS, queryRunner, dsReferencesMap);
+      expect(resultA).toBeUndefined();
+
+      // Test the query with DS originally - should return the original datasource
+      const resultB = getPersistedDSForQuery(queryWithDS, queryRunner, dsReferencesMap);
+      expect(resultB).toEqual({ uid: 'prometheus', type: 'prometheus' });
+
+      // Test a query with no DS originally but not in the mapping - should get the runner's datasource
+      const queryNotInMapping: SceneDataQuery = {
+        refId: 'C',
+        // No datasource, but not in mapping
+      };
+      const resultC = getPersistedDSForQuery(queryNotInMapping, queryRunner, dsReferencesMap);
+      expect(resultC).toEqual({ uid: 'default-ds', type: 'default' });
+    });
+  });
 });
 
 function getMinimalSceneState(body: DashboardLayoutManager): Partial<DashboardSceneState> {
@@ -571,49 +682,50 @@ describe('dynamic layouts', () => {
   });
 });
 
-const annotationLayer1 = new DashboardAnnotationsDataLayer({
-  key: 'layer1',
-  query: {
-    datasource: {
-      type: 'grafana',
-      uid: '-- Grafana --',
-    },
-    name: 'query1',
-    enable: true,
-    iconColor: 'red',
-  },
-  name: 'layer1',
-  isEnabled: true,
-  isHidden: false,
-});
-
-const annotationLayer2 = new DashboardAnnotationsDataLayer({
-  key: 'layer2',
-  query: {
-    datasource: {
-      type: 'prometheus',
-      uid: 'abcdef',
-    },
-    name: 'query2',
-    enable: true,
-    iconColor: 'blue',
-  },
-  name: 'layer2',
-  isEnabled: true,
-  isHidden: true,
-});
-
-// this could happen if a dahboard was created from code and the datasource was not defined
-const annotationLayer3NoDsDefined = new DashboardAnnotationsDataLayer({
-  key: 'layer3',
-  query: {
-    name: 'query3',
-    enable: true,
-    iconColor: 'green',
-  },
-  name: 'layer3',
-  isEnabled: true,
-  isHidden: true,
-});
-
-const annotationLayers = [annotationLayer1, annotationLayer2, annotationLayer3NoDsDefined];
+// Instead of reusing annotation layer objects, create a factory function to generate new ones each time
+function createAnnotationLayers() {
+  return [
+    new DashboardAnnotationsDataLayer({
+      key: 'layer1',
+      query: {
+        datasource: {
+          type: 'grafana',
+          uid: '-- Grafana --',
+        },
+        name: 'query1',
+        enable: true,
+        iconColor: 'red',
+      },
+      name: 'layer1',
+      isEnabled: true,
+      isHidden: false,
+    }),
+    new DashboardAnnotationsDataLayer({
+      key: 'layer2',
+      query: {
+        datasource: {
+          type: 'prometheus',
+          uid: 'abcdef',
+        },
+        name: 'query2',
+        enable: true,
+        iconColor: 'blue',
+      },
+      name: 'layer2',
+      isEnabled: true,
+      isHidden: true,
+    }),
+    // this could happen if a dahboard was created from code and the datasource was not defined
+    new DashboardAnnotationsDataLayer({
+      key: 'layer3',
+      query: {
+        name: 'query3',
+        enable: true,
+        iconColor: 'green',
+      },
+      name: 'layer3',
+      isEnabled: true,
+      isHidden: true,
+    }),
+  ];
+}

public/app/features/dashboard-scene/serialization/transformSceneToSaveModelSchemaV2.ts

@@ -7,6 +7,7 @@ import {
   dataLayers,
   SceneDataQuery,
   SceneDataTransformer,
+  SceneQueryRunner,
   SceneVariableSet,
   VizPanel,
 } from '@grafana/scenes';
@@ -45,7 +46,13 @@ import { DashboardDataLayerSet } from '../scene/DashboardDataLayerSet';
 import { DashboardScene, DashboardSceneState } from '../scene/DashboardScene';
 import { PanelTimeRange } from '../scene/PanelTimeRange';
 import { dashboardSceneGraph } from '../utils/dashboardSceneGraph';
-import { getLibraryPanelBehavior, getPanelIdForVizPanel, getQueryRunnerFor, isLibraryPanel } from '../utils/utils';
+import {
+  getDashboardSceneFor,
+  getLibraryPanelBehavior,
+  getPanelIdForVizPanel,
+  getQueryRunnerFor,
+  isLibraryPanel,
+} from '../utils/utils';
 
 import { getLayout } from './layoutSerializers/utils';
 import { sceneVariablesSetToSchemaV2Variables } from './sceneVariablesSetToVariables';
@@ -239,16 +246,16 @@ function getVizPanelQueries(vizPanel: VizPanel): PanelQueryKind[] {
   const queries: PanelQueryKind[] = [];
   const queryRunner = getQueryRunnerFor(vizPanel);
   const vizPanelQueries = queryRunner?.state.queries;
-  const datasource = queryRunner?.state.datasource ?? getDefaultDataSourceRef();
-
+  const autoAssignedPanelDSRef = getAutoAssignedPanelDSRef(vizPanel);
   if (vizPanelQueries) {
     vizPanelQueries.forEach((query) => {
+      const queryDatasource = getPersistedDSForQuery(query, queryRunner, autoAssignedPanelDSRef);
       const dataQuery: DataQueryKind = {
         kind: getDataQueryKind(query),
         spec: omit(query, 'datasource', 'refId', 'hide'),
       };
       const querySpec: PanelQuerySpec = {
-        datasource: query.datasource ?? datasource,
+        datasource: queryDatasource,
         query: dataQuery,
         refId: query.refId,
         hidden: Boolean(query.hide),
@@ -584,3 +591,41 @@ function validateRowsLayout(layout: unknown) {
     throw new Error('Layout spec items is not an array');
   }
 }
+
+/**
+ * Get a collection of panel queries refIds
+ * the refIds are the ones which did not have a datasource set
+ * @returns a set of panel queries refIds
+ */
+function getAutoAssignedPanelDSRef(vizPanel: VizPanel) {
+  const elementKey = dashboardSceneGraph.getElementIdentifierForVizPanel(vizPanel);
+  const scene = getDashboardSceneFor(vizPanel);
+  const elementMapReferences = scene.serializer.getDSReferencesMapping();
+
+  const panelQueries = elementMapReferences.panels.get(elementKey);
+  return panelQueries;
+}
+
+/**
+ * Get the persisted datasource for a query
+ * When a query is created it could not have a datasource set
+ * we want to respect that and not overwrite it with the auto assigned datasources
+ * resolved in runtime
+ * @param query
+ * @param queryRunner
+ * @param autoAssignedPanelDsRef
+ * @returns
+ */
+export function getPersistedDSForQuery(
+  query: SceneDataQuery,
+  queryRunner: SceneQueryRunner,
+  autoAssignedPanelDsRef: Set<string> | undefined
+) {
+  // if the query has a refId and it is in the panelDsReferences then it did NOT have a datasource
+  const hasMatchingRefId = autoAssignedPanelDsRef?.has(query.refId);
+  if (hasMatchingRefId) {
+    return undefined;
+  }
+
+  return query.datasource || queryRunner?.state?.datasource;
+}