Dashboard - Schema V2 Stateless variables (ds ref independent) (#102571)

* Create basic logic for stateless variables

* Refactor function to get DS from queries and variables

* Simplify code, remove unnecessary wrappers

* Fix typo

* improve naming

* Add unit test for variables

* Fix tests

* Add extra logic to keep backwards compatibility, variables can be used without dashboard scene

* Adjust logic for variables ds

* Refactor code, pass dsReferenceMap

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

@@ -60,7 +60,7 @@ interface DashboardTrackingInfo {
   settings_livenow?: boolean;
 }
 
-interface DSReferencesMapping {
+export interface DSReferencesMapping {
   panels: Map<string, Set<string>>;
   variables: Set<string>;
   annotations: Set<string>;
@@ -264,6 +264,16 @@ export class V2DashboardSerializer
         }
       }
     });
+
+    // initialize autossigned variable ds references map
+    if (saveModel?.variables) {
+      for (const variable of saveModel.variables) {
+        // for query variables that dont have a ds defined add them to the list
+        if (variable.kind === 'QueryVariable' && !variable.spec.datasource) {
+          this.defaultDsReferencesMap.variables.add(variable.spec.name);
+        }
+      }
+    }
   }
 
   getDSReferencesMapping() {

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

@@ -17,6 +17,7 @@ import {
   LibraryPanelKind,
   PanelKind,
   PanelQueryKind,
+  QueryVariableKind,
 } from '@grafana/schema/dist/esm/schema/dashboard/v2alpha0';
 import { MIXED_DATASOURCE_NAME } from 'app/plugins/datasource/mixed/MixedDataSource';
 
@@ -199,6 +200,23 @@ function getPanelDataSource(panel: PanelKind): DataSourceRef | undefined {
   return isMixedDatasource ? { type: 'mixed', uid: MIXED_DATASOURCE_NAME } : datasource;
 }
 
+export function getRuntimeVariableDataSource(variable: QueryVariableKind): DataSourceRef | undefined {
+  let datasource: DataSourceRef | undefined = undefined;
+
+  if (!datasource) {
+    if (!variable.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 === variable.spec.query.kind);
+      datasource = bestGuess ? { uid: bestGuess.uid, type: bestGuess.meta.id } : dsList[defaultDatasource];
+    } else {
+      datasource = variable.spec.datasource;
+    }
+  }
+  return datasource;
+}
+
 function panelQueryKindToSceneQuery(query: PanelQueryKind): SceneDataQuery {
   return {
     refId: query.spec.refId,

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

@@ -22,7 +22,8 @@ import {
 
 import { getIntervalsQueryFromNewIntervalModel } from '../utils/utils';
 
-import { getDataQueryKind, getDataQuerySpec } from './transformSceneToSaveModelSchemaV2';
+import { DSReferencesMapping } from './DashboardSceneSerializer';
+import { getDataQueryKind, getDataQuerySpec, getElementDatasource } from './transformSceneToSaveModelSchemaV2';
 import {
   transformVariableRefreshToEnum,
   transformVariableHideToEnum,
@@ -49,6 +50,7 @@ export function sceneVariablesSetToVariables(set: SceneVariables, keepQueryOptio
       hide: variable.state.hide || OldVariableHide.dontHide,
       type: variable.state.type,
     };
+
     if (sceneUtils.isQueryVariable(variable)) {
       let options: VariableOption[] = [];
       // Not sure if we actually have to still support this option given
@@ -67,7 +69,7 @@ export function sceneVariablesSetToVariables(set: SceneVariables, keepQueryOptio
         options,
         query: variable.state.query,
         definition: variable.state.definition,
-        datasource: variable.state.datasource,
+        datasource: getElementDatasource(set, variable, 'variable'),
         sort: variable.state.sort,
         refresh: variable.state.refresh,
         regex: variable.state.regex,
@@ -224,7 +226,8 @@ function variableValueOptionsToVariableOptions(varState: MultiValueVariable['sta
 
 export function sceneVariablesSetToSchemaV2Variables(
   set: SceneVariables,
-  keepQueryOptions?: boolean
+  keepQueryOptions?: boolean,
+  dsReferencesMapping?: DSReferencesMapping
 ): Array<
   | QueryVariableKind
   | TextVariableKind
@@ -293,7 +296,7 @@ export function sceneVariablesSetToSchemaV2Variables(
           options,
           query: dataQuery,
           definition: variable.state.definition,
-          datasource: variable.state.datasource || {},
+          datasource: getElementDatasource(set, variable, 'variable', undefined, dsReferencesMapping),
           sort: transformSortVariableToEnum(variable.state.sort),
           refresh: transformVariableRefreshToEnum(variable.state.refresh),
           regex: variable.state.regex,

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

@@ -64,6 +64,7 @@ import { getIntervalsFromQueryString } from '../utils/utils';
 
 import { SnapshotVariable } from './custom-variables/SnapshotVariable';
 import { layoutSerializerRegistry } from './layoutSerializers/layoutSerializerRegistry';
+import { getRuntimeVariableDataSource } from './layoutSerializers/utils';
 import { registerPanelInteractionsReporter } from './transformSaveModelToScene';
 import {
   transformCursorSyncV2ToV1,
@@ -292,7 +293,7 @@ function createSceneVariableFromVariableModel(variable: TypedVariableModelV2): S
       value: variable.spec.current?.value ?? '',
       text: variable.spec.current?.text ?? '',
       query: getDataQueryForVariable(variable),
-      datasource: variable.spec.datasource,
+      datasource: getRuntimeVariableDataSource(variable),
       sort: transformSortVariableToEnumV1(variable.spec.sort),
       refresh: transformVariableRefreshToEnumV1(variable.spec.refresh),
       regex: variable.spec.regex,

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

@@ -19,6 +19,7 @@ import {
   VizPanel,
   SceneDataQuery,
   SceneQueryRunner,
+  sceneUtils,
 } from '@grafana/scenes';
 import {
   DashboardCursorSync as DashboardCursorSyncV1,
@@ -50,7 +51,11 @@ import { TabItem } from '../scene/layout-tabs/TabItem';
 import { TabsLayoutManager } from '../scene/layout-tabs/TabsLayoutManager';
 import { DashboardLayoutManager } from '../scene/types/DashboardLayoutManager';
 
-import { getPersistedDSForQuery, transformSceneToSaveModelSchemaV2 } from './transformSceneToSaveModelSchemaV2';
+import {
+  getPersistedDSFor,
+  getElementDatasource,
+  transformSceneToSaveModelSchemaV2,
+} from './transformSceneToSaveModelSchemaV2';
 
 // Mock dependencies
 jest.mock('../utils/dashboardSceneGraph', () => {
@@ -411,7 +416,7 @@ describe('transformSceneToSaveModelSchemaV2', () => {
     expect(result.annotations?.[2].spec.datasource?.type).toBe('loki');
   });
 
-  describe('getPersistedDSForQuery', () => {
+  describe('getPersistedDSFor query', () => {
     it('should respect datasource reference mapping when determining query datasource', () => {
       // Setup test data
       const queryWithoutDS: SceneDataQuery = {
@@ -433,11 +438,11 @@ describe('transformSceneToSaveModelSchemaV2', () => {
       const dsReferencesMap = new Set(['A']);
 
       // Test the query without DS originally - should return undefined
-      const resultA = getPersistedDSForQuery(queryWithoutDS, queryRunner, dsReferencesMap);
+      const resultA = getPersistedDSFor(queryWithoutDS, dsReferencesMap, 'query', queryRunner);
       expect(resultA).toBeUndefined();
 
       // Test the query with DS originally - should return the original datasource
-      const resultB = getPersistedDSForQuery(queryWithDS, queryRunner, dsReferencesMap);
+      const resultB = getPersistedDSFor(queryWithDS, dsReferencesMap, 'query', queryRunner);
       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
@@ -445,51 +450,200 @@ describe('transformSceneToSaveModelSchemaV2', () => {
         refId: 'C',
         // No datasource, but not in mapping
       };
-      const resultC = getPersistedDSForQuery(queryNotInMapping, queryRunner, dsReferencesMap);
+      const resultC = getPersistedDSFor(queryNotInMapping, dsReferencesMap, 'query', queryRunner);
       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',
+  describe('getPersistedDSFor variable', () => {
+    it('should respect datasource reference mapping when determining variable datasource', () => {
+      // Setup test data - variable without datasource
+      const variableWithoutDS = new QueryVariable({
+        name: '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' },
+      // Variable with datasource
+      const variableWithDS = new QueryVariable({
+        name: 'B',
+        datasource: { uid: 'prometheus', type: 'prometheus' },
       });
 
       // 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);
+      // Test the variable without DS originally - should return undefined
+      const resultA = getPersistedDSFor(variableWithoutDS, dsReferencesMap, 'variable');
       expect(resultA).toBeUndefined();
 
-      // Test the query with DS originally - should return the original datasource
-      const resultB = getPersistedDSForQuery(queryWithDS, queryRunner, dsReferencesMap);
+      // Test the variable with DS originally - should return the original datasource
+      const resultB = getPersistedDSFor(variableWithDS, dsReferencesMap, 'variable');
       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',
+      // Test a variable with no DS originally but not in the mapping - should get empty object
+      const variableNotInMapping = new QueryVariable({
+        name: 'C',
         // No datasource, but not in mapping
-      };
-      const resultC = getPersistedDSForQuery(queryNotInMapping, queryRunner, dsReferencesMap);
-      expect(resultC).toEqual({ uid: 'default-ds', type: 'default' });
+      });
+      const resultC = getPersistedDSFor(variableNotInMapping, dsReferencesMap, 'variable');
+      expect(resultC).toEqual({});
     });
   });
 });
 
+describe('getElementDatasource', () => {
+  it('should handle panel query datasources correctly', () => {
+    // Create test elements
+    const vizPanel = new VizPanel({
+      key: 'panel-1',
+      pluginId: 'timeseries',
+    });
+
+    const queryWithDS: SceneDataQuery = {
+      refId: 'B',
+      datasource: { uid: 'prometheus', type: 'prometheus' },
+    };
+
+    const queryWithoutDS: SceneDataQuery = {
+      refId: 'A',
+    };
+
+    // Mock query runner
+    const queryRunner = new SceneQueryRunner({
+      queries: [queryWithoutDS, queryWithDS],
+      datasource: { uid: 'default-ds', type: 'default' },
+    });
+
+    // Mock dsReferencesMapping
+    const dsReferencesMapping = {
+      panels: new Map(new Set([['panel-1', new Set<string>(['A'])]])),
+      variables: new Set<string>(),
+      annotations: new Set<string>(),
+    };
+
+    // Call the function with the panel and query with DS
+    const resultWithDS = getElementDatasource(vizPanel, queryWithDS, 'panel', queryRunner, dsReferencesMapping);
+    expect(resultWithDS).toEqual({ uid: 'prometheus', type: 'prometheus' });
+
+    // Call the function with the panel and query without DS
+    const resultWithoutDS = getElementDatasource(vizPanel, queryWithoutDS, 'panel', queryRunner, dsReferencesMapping);
+    expect(resultWithoutDS).toBeUndefined();
+  });
+
+  it('should handle variable datasources correctly', () => {
+    // Create a variable set
+    const variableSet = new SceneVariableSet({
+      variables: [
+        new QueryVariable({
+          name: 'A',
+          // No datasource
+        }),
+        new QueryVariable({
+          name: 'B',
+          datasource: { uid: 'prometheus', type: 'prometheus' },
+        }),
+      ],
+    });
+
+    // Variable with DS
+    const variableWithDS = variableSet.getByName('B');
+
+    // Variable without DS
+    const variableWithoutDS = variableSet.getByName('A');
+
+    // Mock dsReferencesMapping
+    const dsReferencesMapping = {
+      panels: new Map(new Set([['panel-1', new Set<string>(['A'])]])),
+      variables: new Set<string>(['A']),
+      annotations: new Set<string>(),
+    };
+
+    // Call the function with variables
+    if (variableWithDS && sceneUtils.isQueryVariable(variableWithDS)) {
+      const resultWithDS = getElementDatasource(
+        variableSet,
+        variableWithDS,
+        'variable',
+        undefined,
+        dsReferencesMapping
+      );
+      expect(resultWithDS).toEqual({ uid: 'prometheus', type: 'prometheus' });
+    }
+
+    if (variableWithoutDS && sceneUtils.isQueryVariable(variableWithoutDS)) {
+      // Test with auto-assigned variable (in the mapping)
+      const resultWithoutDS = getElementDatasource(variableSet, variableWithoutDS, 'variable');
+      expect(resultWithoutDS).toEqual(undefined);
+    }
+  });
+
+  it('should return undefined for non-query variables', () => {
+    // Create a variable set with non-query variable
+    const variableSet = new SceneVariableSet({
+      variables: [
+        new ConstantVariable({
+          name: 'constant',
+          value: 'value',
+        }),
+      ],
+    });
+
+    // Non-query variable
+    const constantVar = variableSet.getByName('constant');
+
+    // Call the function
+    // @ts-expect-error
+    const result = getElementDatasource(variableSet, constantVar, 'variable');
+    expect(result).toBeUndefined();
+  });
+
+  it('should return undefined for non-query variables', () => {
+    // Create a variable set with non-query variable types
+    const variableSet = new SceneVariableSet({
+      variables: [
+        // Use TextBoxVariable which is not a QueryVariable
+        new TextBoxVariable({
+          name: 'textVar',
+          value: 'text-value',
+        }),
+      ],
+    });
+
+    // Non-query variable - this is safe because getElementDatasource checks if it's a query variable
+    const textVar = variableSet.getByName('textVar');
+
+    // Call the function
+    // @ts-expect-error
+    const result = getElementDatasource(variableSet, textVar, 'variable');
+    expect(result).toBeUndefined();
+  });
+
+  it('should handle invalid input combinations', () => {
+    const vizPanel = new VizPanel({
+      key: 'panel-1',
+      pluginId: 'timeseries',
+    });
+
+    const variableSet = new SceneVariableSet({
+      variables: [
+        new QueryVariable({
+          name: 'A',
+        }),
+      ],
+    });
+
+    const variable = variableSet.getByName('A');
+    const query: SceneDataQuery = { refId: 'A' };
+
+    if (variable && sceneUtils.isQueryVariable(variable)) {
+      // Panel with variable
+      expect(getElementDatasource(vizPanel, variable, 'panel')).toBeUndefined();
+    }
+    // Variable set with query
+    expect(getElementDatasource(variableSet, query, 'variable')).toBeUndefined();
+  });
+});
+
 function getMinimalSceneState(body: DashboardLayoutManager): Partial<DashboardSceneState> {
   return {
     id: 1,

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

@@ -5,11 +5,14 @@ import { config } from '@grafana/runtime';
 import {
   behaviors,
   dataLayers,
+  QueryVariable,
   SceneDataQuery,
   SceneDataTransformer,
   SceneQueryRunner,
+  SceneVariables,
   SceneVariableSet,
   VizPanel,
+  sceneUtils,
 } from '@grafana/scenes';
 import { DataSourceRef } from '@grafana/schema';
 import { sortedDeepCloneWithoutNulls } from 'app/core/utils/object';
@@ -46,14 +49,9 @@ import { DashboardDataLayerSet } from '../scene/DashboardDataLayerSet';
 import { DashboardScene, DashboardSceneState } from '../scene/DashboardScene';
 import { PanelTimeRange } from '../scene/PanelTimeRange';
 import { dashboardSceneGraph } from '../utils/dashboardSceneGraph';
-import {
-  getDashboardSceneFor,
-  getLibraryPanelBehavior,
-  getPanelIdForVizPanel,
-  getQueryRunnerFor,
-  isLibraryPanel,
-} from '../utils/utils';
+import { getLibraryPanelBehavior, getPanelIdForVizPanel, getQueryRunnerFor, isLibraryPanel } from '../utils/utils';
 
+import { DSReferencesMapping } from './DashboardSceneSerializer';
 import { getLayout } from './layoutSerializers/utils';
 import { sceneVariablesSetToSchemaV2Variables } from './sceneVariablesSetToVariables';
 import { colorIdEnumToColorIdV2, transformCursorSynctoEnum } from './transformToV2TypesUtils';
@@ -72,6 +70,8 @@ export function transformSceneToSaveModelSchemaV2(scene: DashboardScene, isSnaps
   const controlsState = sceneDash.controls?.state;
   const refreshPicker = controlsState?.refreshPicker;
 
+  const dsReferencesMapping: DSReferencesMapping = scene.serializer.getDSReferencesMapping();
+
   const dashboardSchemaV2: DeepPartial<DashboardV2Spec> = {
     //dashboard settings
     title: sceneDash.title,
@@ -100,11 +100,11 @@ export function transformSceneToSaveModelSchemaV2(scene: DashboardScene, isSnaps
     // EOF time settings
 
     // variables
-    variables: getVariables(sceneDash),
+    variables: getVariables(sceneDash, dsReferencesMapping),
     // EOF variables
 
     // elements
-    elements: getElements(scene),
+    elements: getElements(scene, dsReferencesMapping),
     // EOF elements
 
     // annotations
@@ -147,14 +147,18 @@ function getLiveNow(state: DashboardSceneState) {
   return Boolean(liveNow);
 }
 
-function getElements(scene: DashboardScene) {
+function getElements(scene: DashboardScene, dsReferencesMapping: DSReferencesMapping) {
   const panels = scene.state.body.getVizPanels() ?? [];
-
-  const panelsArray = panels.map(vizPanelToSchemaV2);
+  const panelsArray = panels.map((vizPanel) => {
+    return vizPanelToSchemaV2(vizPanel, dsReferencesMapping);
+  });
   return createElements(panelsArray, scene);
 }
 
-export function vizPanelToSchemaV2(vizPanel: VizPanel): PanelKind | LibraryPanelKind {
+export function vizPanelToSchemaV2(
+  vizPanel: VizPanel,
+  dsReferencesMapping?: DSReferencesMapping
+): PanelKind | LibraryPanelKind {
   if (isLibraryPanel(vizPanel)) {
     const behavior = getLibraryPanelBehavior(vizPanel)!;
     const elementSpec: LibraryPanelKind = {
@@ -216,7 +220,7 @@ export function vizPanelToSchemaV2(vizPanel: VizPanel): PanelKind | LibraryPanel
       data: {
         kind: 'QueryGroup',
         spec: {
-          queries: getVizPanelQueries(vizPanel),
+          queries: getVizPanelQueries(vizPanel, dsReferencesMapping),
           transformations: getVizPanelTransformations(vizPanel),
           queryOptions: getVizPanelQueryOptions(vizPanel),
         },
@@ -242,14 +246,14 @@ function getPanelLinks(panel: VizPanel): DataLink[] {
   return [];
 }
 
-function getVizPanelQueries(vizPanel: VizPanel): PanelQueryKind[] {
+function getVizPanelQueries(vizPanel: VizPanel, dsReferencesMapping?: DSReferencesMapping): PanelQueryKind[] {
   const queries: PanelQueryKind[] = [];
   const queryRunner = getQueryRunnerFor(vizPanel);
   const vizPanelQueries = queryRunner?.state.queries;
-  const autoAssignedPanelDSRef = getAutoAssignedPanelDSRef(vizPanel);
+
   if (vizPanelQueries) {
     vizPanelQueries.forEach((query) => {
-      const queryDatasource = getPersistedDSForQuery(query, queryRunner, autoAssignedPanelDSRef);
+      const queryDatasource = getElementDatasource(vizPanel, query, 'panel', queryRunner, dsReferencesMapping);
       const dataQuery: DataQueryKind = {
         kind: getDataQueryKind(query),
         spec: omit(query, 'datasource', 'refId', 'hide'),
@@ -356,7 +360,7 @@ function createElements(panels: Element[], scene: DashboardScene): Record<string
   }, {});
 }
 
-function getVariables(oldDash: DashboardSceneState) {
+function getVariables(oldDash: DashboardSceneState, dsReferencesMapping?: DSReferencesMapping) {
   const variablesSet = oldDash.$variables;
 
   // variables is an array of all variables kind (union)
@@ -372,7 +376,7 @@ function getVariables(oldDash: DashboardSceneState) {
   > = [];
 
   if (variablesSet instanceof SceneVariableSet) {
-    variables = sceneVariablesSetToSchemaV2Variables(variablesSet);
+    variables = sceneVariablesSetToSchemaV2Variables(variablesSet, false, dsReferencesMapping);
   }
 
   return variables;
@@ -592,40 +596,125 @@ function validateRowsLayout(layout: unknown) {
   }
 }
 
-/**
- * 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();
+function getAutoAssignedDSRef(
+  element: VizPanel | SceneVariables,
+  type: 'panels' | 'variables',
+  elementMapReferences?: DSReferencesMapping
+): Set<string> {
+  if (!elementMapReferences) {
+    return new Set();
+  }
+  if (type === 'panels' && isVizPanel(element)) {
+    const elementKey = dashboardSceneGraph.getElementIdentifierForVizPanel(element);
+    return elementMapReferences.panels.get(elementKey) || new Set();
+  }
 
-  const panelQueries = elementMapReferences.panels.get(elementKey);
-  return panelQueries;
+  return elementMapReferences.variables;
 }
 
 /**
- * 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
+ * Determines if a data source reference should be persisted for a query or variable
  */
-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) {
+export function getPersistedDSFor<T extends SceneDataQuery | QueryVariable>(
+  element: T,
+  autoAssignedDsRef: Set<string>,
+  type: 'query' | 'variable',
+  context?: SceneQueryRunner
+): DataSourceRef | undefined {
+  // Get the element identifier - refId for queries, name for variables
+  const elementId = getElementIdentifier(element, type);
+
+  // If the element is in the auto-assigned set, it didn't have a datasource specified
+  if (autoAssignedDsRef?.has(elementId)) {
     return undefined;
   }
 
-  return query.datasource || queryRunner?.state?.datasource;
+  // Return appropriate datasource reference based on element type
+  if (type === 'query') {
+    if ('datasource' in element && element.datasource) {
+      // If element has its own datasource, use that
+      return element.datasource;
+    }
+
+    // For queries missing a datasource but not in auto-assigned set, use datasource from context (queryRunner)
+    return context?.state?.datasource;
+  }
+
+  if (type === 'variable' && 'state' in element && 'datasource' in element.state) {
+    return element.state.datasource || {};
+  }
+
+  return undefined;
+}
+
+/**
+ * Helper function to extract which identifier to use from a query or variable element
+ * @returns refId for queries, name for variables
+ * TODO: we will add annotations in the future
+ */
+function getElementIdentifier<T extends SceneDataQuery | QueryVariable>(
+  element: T,
+  type: 'query' | 'variable'
+): string {
+  // when is type query look for refId
+  if (type === 'query') {
+    return 'refId' in element ? element.refId : '';
+  }
+  // when is type variable look for the name of the variable
+  return 'state' in element && 'name' in element.state ? element.state.name : '';
+}
+
+function isVizPanel(element: VizPanel | SceneVariables): element is VizPanel {
+  // FIXME: is there another way to do this?
+  return 'pluginId' in element.state;
+}
+
+function isSceneVariables(element: VizPanel | SceneVariables): element is SceneVariables {
+  // Check for properties unique to SceneVariables but not in VizPanel
+  return !('pluginId' in element.state) && ('variables' in element.state || 'getValue' in element);
+}
+
+function isSceneDataQuery(query: SceneDataQuery | QueryVariable): query is SceneDataQuery {
+  return 'refId' in query && !('state' in query);
+}
+
+/**
+ * Get the persisted datasource for a query or variable
+ * When a query or variable 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
+ *
+ */
+export function getElementDatasource(
+  element: VizPanel | SceneVariables,
+  queryElement: SceneDataQuery | QueryVariable,
+  type: 'panel' | 'variable',
+  queryRunner?: SceneQueryRunner,
+  dsReferencesMapping?: DSReferencesMapping
+): DataSourceRef | undefined {
+  if (type === 'panel') {
+    if (!queryRunner || !isVizPanel(element) || !isSceneDataQuery(queryElement)) {
+      return undefined;
+    }
+    // Get datasource for panel query
+    const autoAssignedRefs = getAutoAssignedDSRef(element, 'panels', dsReferencesMapping);
+    return getPersistedDSFor(queryElement, autoAssignedRefs, 'query', queryRunner);
+  }
+
+  if (type === 'variable') {
+    if (!isSceneVariables(element) || isSceneDataQuery(queryElement)) {
+      return undefined;
+    }
+    // Get datasource for variable
+    if (!sceneUtils.isQueryVariable(queryElement)) {
+      return undefined;
+    }
+    const autoAssignedRefs = getAutoAssignedDSRef(element, 'variables', dsReferencesMapping);
+    // Important: Only return the datasource if it's not in auto-assigned refs
+    // and if the result would not be an empty object
+    const result = getPersistedDSFor(queryElement, autoAssignedRefs, 'variable');
+    return Object.keys(result || {}).length > 0 ? result : undefined;
+  }
+
+  return undefined;
 }