Dashboards: SchemaV2 - Fix stateless queries for mixed ds (#103885)

* Dashboards: SchemaV2 - Fix stateless queries for mixed ds

* Add uni test for new utils function

* Add extra condition to validate we have datasources configured in the grafanaBootData

* Refactor code

* remove unnecessary test

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

@@ -0,0 +1,132 @@
+import { PanelQueryKind } from '@grafana/schema/dist/esm/schema/dashboard/v2alpha1/types.spec.gen';
+
+import { getRuntimePanelDataSource } from './utils';
+
+// Mock the config needed for the function
+jest.mock('@grafana/runtime', () => ({
+  ...jest.requireActual('@grafana/runtime'),
+  config: {
+    ...jest.requireActual('@grafana/runtime').config,
+    bootData: {
+      settings: {
+        defaultDatasource: 'default-ds-grafana',
+        datasources: {
+          'default-ds-grafana': {
+            uid: 'default-ds-uid',
+            name: 'Default DS',
+            meta: { id: 'default-ds-grafana' },
+            type: 'datasource',
+          },
+          prometheus: {
+            uid: 'prometheus-uid',
+            name: 'Prometheus',
+            meta: { id: 'prometheus' },
+            type: 'datasource',
+          },
+          loki: {
+            uid: 'loki-uid',
+            name: 'Loki',
+            meta: { id: 'loki' },
+            type: 'datasource',
+          },
+        },
+      },
+    },
+  },
+}));
+
+describe('getRuntimePanelDataSource', () => {
+  it('should return the datasource when it is specified in the query', () => {
+    const query: PanelQueryKind = {
+      kind: 'PanelQuery',
+      spec: {
+        refId: 'A',
+        hidden: false,
+        datasource: {
+          uid: 'test-ds-uid',
+          type: 'test-ds-type',
+        },
+        query: {
+          kind: 'prometheus',
+          spec: {},
+        },
+      },
+    };
+
+    const result = getRuntimePanelDataSource(query);
+
+    expect(result).toEqual({
+      uid: 'test-ds-uid',
+      type: 'test-ds-type',
+    });
+  });
+
+  it('should infer datasource based on query kind when datasource is not specified', () => {
+    const query: PanelQueryKind = {
+      kind: 'PanelQuery',
+      spec: {
+        refId: 'A',
+        hidden: false,
+        datasource: undefined,
+        query: {
+          kind: 'prometheus',
+          spec: {},
+        },
+      },
+    };
+
+    const result = getRuntimePanelDataSource(query);
+
+    expect(result).toEqual({
+      uid: 'prometheus-uid',
+      type: 'prometheus',
+    });
+  });
+
+  it('should use default datasource when no datasource is specified and query kind does not match any available datasource', () => {
+    const query: PanelQueryKind = {
+      kind: 'PanelQuery',
+      spec: {
+        refId: 'A',
+        hidden: false,
+        datasource: undefined,
+        query: {
+          kind: 'unknown-type',
+          spec: {},
+        },
+      },
+    };
+
+    const result = getRuntimePanelDataSource(query);
+
+    expect(result).toEqual({
+      uid: 'default-ds-uid',
+      type: 'default-ds-grafana',
+    });
+  });
+
+  it('should handle the case when datasource uid is empty string', () => {
+    const query: PanelQueryKind = {
+      kind: 'PanelQuery',
+      spec: {
+        refId: 'A',
+        hidden: false,
+        datasource: {
+          uid: '',
+          type: 'test-ds-type',
+        },
+        query: {
+          kind: 'prometheus',
+          spec: {},
+        },
+      },
+    };
+
+    const result = getRuntimePanelDataSource(query);
+
+    expect(result).toEqual({
+      uid: 'prometheus-uid',
+      type: 'prometheus',
+    });
+  });
+});

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

@@ -186,11 +186,7 @@ function getPanelDataSource(panel: PanelKind): DataSourceRef | undefined {
   panel.spec.data.spec.queries.forEach((query) => {
     if (!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];
+        datasource = getRuntimePanelDataSource(query);
       } else {
         datasource = query.spec.datasource;
       }
@@ -203,26 +199,53 @@ function getPanelDataSource(panel: PanelKind): DataSourceRef | undefined {
 }
 
 export function getRuntimeVariableDataSource(variable: QueryVariableKind): DataSourceRef | undefined {
-  let datasource: DataSourceRef | undefined = undefined;
+  return getDataSourceForQuery(variable.spec.datasource, variable.spec.query.kind);
+}
 
-  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;
-    }
+export function getRuntimePanelDataSource(query: PanelQueryKind): DataSourceRef | undefined {
+  return getDataSourceForQuery(query.spec.datasource, query.spec.query.kind);
+}
+
+/**
+ * @param querySpecDS - The datasource specified in the query
+ * @param queryKind - The kind of query being performed
+ * @returns The resolved DataSourceRef
+ */
+function getDataSourceForQuery(
+  querySpecDS: DataSourceRef | undefined | null,
+  queryKind: string
+): DataSourceRef | undefined {
+  // If datasource is specified and has a uid, use it
+  if (querySpecDS?.uid) {
+    return querySpecDS;
   }
-  return datasource;
+
+  // Otherwise try to infer datasource based on query kind (kind = ds type)
+  const defaultDatasource = config.bootData.settings.defaultDatasource;
+  const dsList = config.bootData.settings.datasources;
+
+  // Look up by query type/kind
+  const bestGuess = dsList && Object.values(dsList).find((ds) => ds.meta.id === queryKind);
+
+  if (bestGuess) {
+    return { uid: bestGuess.uid, type: bestGuess.meta.id };
+  } else if (dsList && dsList[defaultDatasource]) {
+    // In the datasource list from bootData "id" is the type and the uid could be uid or the name
+    // in cases like grafana, dashboard or mixed datasource
+    return {
+      uid: dsList[defaultDatasource].uid || dsList[defaultDatasource].name,
+      type: dsList[defaultDatasource].meta.id,
+    };
+  }
+
+  // If we don't find a default datasource, return undefined
+  return undefined;
 }
 
 function panelQueryKindToSceneQuery(query: PanelQueryKind): SceneDataQuery {
   return {
     refId: query.spec.refId,
-    datasource: query.spec.datasource,
+    datasource: getRuntimePanelDataSource(query),
     hide: query.spec.hidden,
     ...query.spec.query.spec,
   };