Dashboard: SchemaV2 Fix Import mapping datasource (#104200)

* Dashboard: SchemaV2 Fix mapping ds for variables and annotations

* process annotations and vars ds inputs so they can be selected

* clean up so async works

* Add basic unit test for v2 dashboards mapping

* clean up tests

* linting

---------

Co-authored-by: Haris Rozajac <haris.rozajac12@gmail.com>

public/app/features/dashboard-scene/v2schema/ImportDashboardOverviewV2.tsx

@@ -39,8 +39,8 @@ export function ImportDashboardOverviewV2() {
       ...dashboard,
       title: form.dashboard.title,
       annotations: dashboard.annotations?.map((annotation: AnnotationQueryKind) => {
-        if (annotation.spec.datasource?.type) {
-          const dsType = annotation.spec.datasource.type;
+        if (annotation.spec.query?.kind) {
+          const dsType = annotation.spec.query.kind;
           if (form[`datasource-${dsType}` as keyof typeof form]) {
             const ds = form[`datasource-${dsType}` as keyof typeof form] as { uid: string; type: string };
             return {
@@ -59,8 +59,8 @@ export function ImportDashboardOverviewV2() {
       }),
       variables: dashboard.variables?.map((variable) => {
         if (variable.kind === 'QueryVariable') {
-          if (variable.spec.datasource?.type) {
-            const dsType = variable.spec.datasource.type;
+          if (variable.spec.query?.kind) {
+            const dsType = variable.spec.query.kind;
             if (form[`datasource-${dsType}` as keyof typeof form]) {
               const ds = form[`datasource-${dsType}` as keyof typeof form] as { uid: string; type: string };
               return {

public/app/features/manage-dashboards/state/actions.test.ts

@@ -2,6 +2,12 @@ import { thunkTester } from 'test/core/thunk/thunkTester';
 
 import { DataSourceInstanceSettings, ThresholdsMode } from '@grafana/data';
 import { defaultDashboard, FieldColorModeId } from '@grafana/schema';
+import {
+  DashboardV2Spec,
+  defaultDashboardV2Spec,
+  defaultPanelSpec,
+  defaultQueryVariableSpec,
+} from '@grafana/schema/dist/esm/schema/dashboard/v2alpha0';
 import { browseDashboardsAPI } from 'app/features/browse-dashboards/api/browseDashboardsAPI';
 import { getLibraryPanel } from 'app/features/library-panels/state/api';
 
@@ -10,7 +16,7 @@ import { LibraryElementDTO } from '../../library-panels/types';
 import { DashboardJson } from '../types';
 import { validateDashboardJson } from '../utils/validation';
 
-import { getLibraryPanelInputs, importDashboard, processDashboard } from './actions';
+import { getLibraryPanelInputs, importDashboard, processDashboard, processV2Datasources } from './actions';
 import { DataSourceInput, ImportDashboardDTO, initialImportDashboardState, InputType } from './reducers';
 
 jest.mock('app/features/library-panels/state/api');
@@ -18,6 +24,43 @@ const mocks = {
   getLibraryPanel: jest.mocked(getLibraryPanel),
 };
 
+jest.mock('@grafana/runtime', () => ({
+  ...jest.requireActual('@grafana/runtime'),
+  getDataSourceSrv: () => ({
+    ...jest.requireActual('@grafana/runtime').getDataSourceSrv(),
+    get: jest.fn().mockImplementation((dsType: { type: string }) => {
+      const dsList: {
+        [key: string]: {
+          uid: string;
+          name: string;
+          type: string;
+          meta: { id: string };
+        };
+      } = {
+        prometheus: {
+          uid: 'prom-uid',
+          name: 'prometheus',
+          type: 'prometheus',
+          meta: { id: 'prometheus' },
+        },
+        loki: {
+          uid: 'loki-uid',
+          name: 'Loki',
+          type: 'loki',
+          meta: { id: 'loki' },
+        },
+        grafana: {
+          uid: 'grafana-uid',
+          name: 'Grafana',
+          type: 'grafana',
+          meta: { id: 'grafana' },
+        },
+      };
+      return dsList[dsType.type];
+    }),
+  }),
+}));
+
 describe('importDashboard', () => {
   it('Should send data source uid', async () => {
     // note: the actual action returned is more complicated
@@ -755,3 +798,164 @@ describe('processDashboard', () => {
     expect(dsInputsForLibPanels).toHaveLength(0);
   });
 });
+
+describe('processV2Datasources', () => {
+  let panels: DashboardV2Spec['elements'];
+  let v2DashboardJson: DashboardV2Spec;
+
+  beforeEach(() => {
+    panels = {
+      'element-panel-a': {
+        kind: 'Panel',
+        spec: {
+          ...defaultPanelSpec(),
+          id: 1,
+          title: 'Panel A',
+          data: {
+            kind: 'QueryGroup',
+            spec: {
+              transformations: [],
+              queryOptions: {},
+              queries: [
+                {
+                  kind: 'PanelQuery',
+                  spec: {
+                    refId: 'A',
+                    hidden: false,
+                    query: {
+                      kind: 'prometheus',
+                      spec: {
+                        expr: 'access_evaluation_duration_count',
+                        range: true,
+                      },
+                    },
+                  },
+                },
+              ],
+            },
+          },
+        },
+      },
+    };
+    v2DashboardJson = {
+      ...defaultDashboardV2Spec(),
+      elements: {
+        ...panels,
+      },
+      variables: [
+        {
+          kind: 'QueryVariable',
+          spec: {
+            ...defaultQueryVariableSpec(),
+            name: 'var1',
+            query: {
+              kind: 'loki',
+              spec: {
+                expr: 'access_evaluation_duration_count',
+                range: true,
+              },
+            },
+          },
+        },
+      ],
+      annotations: [
+        {
+          kind: 'AnnotationQuery',
+          spec: {
+            name: 'annotation1',
+            enable: true,
+            hide: false,
+            iconColor: 'red',
+            query: {
+              kind: 'loki',
+              spec: {
+                expr: 'access_evaluation_duration_count',
+                range: true,
+              },
+            },
+          },
+        },
+      ],
+      layout: {
+        kind: 'GridLayout',
+        spec: {
+          items: [
+            {
+              kind: 'GridLayoutItem',
+              spec: {
+                x: 0,
+                y: 0,
+                width: 12,
+                height: 8,
+                element: {
+                  kind: 'ElementReference',
+                  name: 'element-panel-a',
+                },
+              },
+            },
+            {
+              kind: 'GridLayoutItem',
+              spec: {
+                x: 0,
+                y: 0,
+                width: 12,
+                height: 8,
+                element: {
+                  kind: 'ElementReference',
+                  name: 'element-panel-b',
+                },
+              },
+            },
+          ],
+        },
+      },
+    };
+  });
+  // should set the correct inputs for panels
+  it('Should extract datasource inputs from panel queries, variables and annotations', async () => {
+    // Execute the test using thunkTester
+    const dispatchedActions = await thunkTester({
+      thunk: processV2Datasources,
+      initialState: {
+        inputs: [
+          // for panels
+          {
+            name: 'Prometheus',
+            pluginId: 'prometheus',
+            type: InputType.DataSource,
+          },
+          // for variables and annotations
+          {
+            name: 'Loki',
+            pluginId: 'loki',
+            type: InputType.DataSource,
+          },
+        ],
+      },
+    })
+      .givenThunk(processV2Datasources)
+      .whenThunkIsDispatched(v2DashboardJson);
+
+    // Find the setInputs action in the dispatched actions
+    const setInputsAction = dispatchedActions.find((action) => action.type === 'manageDashboards/setInputs');
+    //
+    // Verify the action was dispatched
+    expect(setInputsAction).toBeDefined();
+
+    // Verify the datasource inputs were correctly extracted
+    expect(setInputsAction?.payload).toEqual(
+      expect.arrayContaining([
+        expect.objectContaining({
+          name: 'prometheus',
+          pluginId: 'prometheus',
+          type: InputType.DataSource,
+        }),
+        expect.objectContaining({
+          name: 'Loki',
+          pluginId: 'loki',
+          type: InputType.DataSource,
+        }),
+      ])
+    );
+  });
+});

public/app/features/manage-dashboards/state/actions.ts

@@ -1,6 +1,11 @@
 import { DataSourceInstanceSettings } from '@grafana/data';
 import { getBackendSrv, getDataSourceSrv, isFetchError } from '@grafana/runtime';
-import { Spec as DashboardV2Spec } from '@grafana/schema/dist/esm/schema/dashboard/v2alpha1/types.spec.gen';
+import {
+  Spec as DashboardV2Spec,
+  QueryVariableKind,
+  PanelQueryKind,
+  AnnotationQueryKind,
+} from '@grafana/schema/dist/esm/schema/dashboard/v2alpha1/types.spec.gen';
 import { notifyApp } from 'app/core/actions';
 import { createErrorNotification } from 'app/core/copy/appNotification';
 import { browseDashboardsAPI, ImportInputs } from 'app/features/browse-dashboards/api/browseDashboardsAPI';
@@ -62,7 +67,7 @@ export function importDashboardJson(dashboard: any): ThunkResult<void> {
 export function importDashboardV2Json(dashboard: DashboardV2Spec): ThunkResult<void> {
   return async (dispatch) => {
     dispatch(setJsonDashboard(dashboard));
-    dispatch(processV2Elements(dashboard));
+    dispatch(processV2Datasources(dashboard));
   };
 }
 
@@ -152,9 +157,9 @@ function processElements(dashboardJson?: { __elements?: Record<string, LibraryEl
   };
 }
 
-function processV2Elements(dashboard: DashboardV2Spec): ThunkResult<void> {
+export function processV2Datasources(dashboard: DashboardV2Spec): ThunkResult<void> {
   return async function (dispatch) {
-    const elements = dashboard.elements;
+    const { elements, variables, annotations } = dashboard;
     // get elements from dashboard
     // each element can only be a panel
     const inputs: Record<string, DataSourceInput> = {};
@@ -162,39 +167,23 @@ function processV2Elements(dashboard: DashboardV2Spec): ThunkResult<void> {
       if (element.kind !== 'Panel') {
         throw new Error('Only panels are currenlty supported in v2 dashboards');
       }
-
+      if (element.spec.data.spec.queries.length > 0) {
         for (const query of element.spec.data.spec.queries) {
-        const datasourceRef = query.spec.datasource;
-        if (!datasourceRef) {
-          let dataSourceInput: DataSourceInput | undefined;
-          const dsType = query.spec.query.kind;
-          const datasource = await getDatasourceSrv().get({ type: dsType });
-          if (!datasource) {
-            dataSourceInput = {
-              name: dsType,
-              label: dsType,
-              info: `No data sources of type ${dsType} found`,
-              value: '',
-              type: InputType.DataSource,
-              pluginId: dsType,
-            };
+          await processV2DatasourceInput(query.spec, inputs);
+        }
+      }
+    }
 
-            inputs[dsType] = dataSourceInput;
-          } else {
-            dataSourceInput = {
-              name: datasource.name,
-              label: datasource.name,
-              info: `Select a ${datasource.name} data source`,
-              value: datasource.uid,
-              type: InputType.DataSource,
-              pluginId: datasource.meta?.id,
-            };
+    for (const variable of variables) {
+      if (variable.kind === 'QueryVariable') {
+        await processV2DatasourceInput(variable.spec, inputs);
+      }
+    }
 
-            inputs[datasource.meta?.id] = dataSourceInput;
-          }
-        }
-      }
+    for (const annotation of annotations) {
+      await processV2DatasourceInput(annotation.spec, inputs);
     }
+
     dispatch(setInputs(Object.values(inputs)));
   };
 }
@@ -340,3 +329,38 @@ export async function searchFolders(
 export function getFolderByUid(uid: string): Promise<{ uid: string; title: string }> {
   return getBackendSrv().get(`/api/folders/${uid}`);
 }
+
+export async function processV2DatasourceInput(
+  obj: PanelQueryKind['spec'] | QueryVariableKind['spec'] | AnnotationQueryKind['spec'],
+  inputs: Record<string, DataSourceInput> = {}
+) {
+  const datasourceRef = obj?.datasource;
+  if (!datasourceRef && obj?.query) {
+    const dsType = obj.query.kind;
+    const datasource = await getDatasourceSrv().get({ type: dsType });
+    let dataSourceInput: DataSourceInput | undefined;
+    if (datasource) {
+      dataSourceInput = {
+        name: datasource.name,
+        label: datasource.name,
+        info: `Select a ${datasource.name} data source`,
+        value: datasource.uid,
+        type: InputType.DataSource,
+        pluginId: datasource.meta?.id,
+      };
+
+      inputs[datasource.meta?.id] = dataSourceInput;
+    } else {
+      dataSourceInput = {
+        name: dsType,
+        label: dsType,
+        info: `No data sources of type ${dsType} found`,
+        value: '',
+        type: InputType.DataSource,
+        pluginId: dsType,
+      };
+
+      inputs[dsType] = dataSourceInput;
+    }
+  }
+}