diff --git a/public/app/features/dashboard-scene/serialization/DashboardSceneSerializer.ts b/public/app/features/dashboard-scene/serialization/DashboardSceneSerializer.ts index a5c3ec6e4d4..ef61766b96d 100644 --- a/public/app/features/dashboard-scene/serialization/DashboardSceneSerializer.ts +++ b/public/app/features/dashboard-scene/serialization/DashboardSceneSerializer.ts @@ -60,7 +60,7 @@ interface DashboardTrackingInfo { settings_livenow?: boolean; } -interface DSReferencesMapping { +export interface DSReferencesMapping { panels: Map>; variables: Set; annotations: Set; @@ -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() { diff --git a/public/app/features/dashboard-scene/serialization/layoutSerializers/utils.ts b/public/app/features/dashboard-scene/serialization/layoutSerializers/utils.ts index 66bcaed53f7..d34169c680a 100644 --- a/public/app/features/dashboard-scene/serialization/layoutSerializers/utils.ts +++ b/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, diff --git a/public/app/features/dashboard-scene/serialization/sceneVariablesSetToVariables.ts b/public/app/features/dashboard-scene/serialization/sceneVariablesSetToVariables.ts index c7d86f4c910..c787153fc94 100644 --- a/public/app/features/dashboard-scene/serialization/sceneVariablesSetToVariables.ts +++ b/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, diff --git a/public/app/features/dashboard-scene/serialization/transformSaveModelSchemaV2ToScene.ts b/public/app/features/dashboard-scene/serialization/transformSaveModelSchemaV2ToScene.ts index ceb9ebd7eeb..babc4b08f73 100644 --- a/public/app/features/dashboard-scene/serialization/transformSaveModelSchemaV2ToScene.ts +++ b/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, diff --git a/public/app/features/dashboard-scene/serialization/transformSceneToSaveModelSchemaV2.test.ts b/public/app/features/dashboard-scene/serialization/transformSceneToSaveModelSchemaV2.test.ts index 403432cf8ba..72c0a46c034 100644 --- a/public/app/features/dashboard-scene/serialization/transformSceneToSaveModelSchemaV2.test.ts +++ b/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(['A'])]])), + variables: new Set(), + annotations: new Set(), + }; + + // 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(['A'])]])), + variables: new Set(['A']), + annotations: new Set(), + }; + + // 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 { return { id: 1, diff --git a/public/app/features/dashboard-scene/serialization/transformSceneToSaveModelSchemaV2.ts b/public/app/features/dashboard-scene/serialization/transformSceneToSaveModelSchemaV2.ts index b88dfa68719..22dbb958ebd 100644 --- a/public/app/features/dashboard-scene/serialization/transformSceneToSaveModelSchemaV2.ts +++ b/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 = { //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 = []; 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 { + 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 | 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( + element: T, + autoAssignedDsRef: Set, + 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( + 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; }