Explore: Add error handling function when retrieving datasource (#55055)

* Use datasource function with error handling

* Check datasource exists when validating query, add test

* Add mock to test that now requires it

* Add comment explaining verbose validity logic

public/app/core/utils/explore.ts

@@ -323,11 +323,25 @@ export async function ensureQueries(
         refId = getNextRefIdChar(allQueries);
       }
 
-      allQueries.push({
-        ...query,
-        refId,
-        key,
-      });
+      // if a query has a datasource, validate it and only add it if valid
+      // if a query doesn't have a datasource, do not worry about it at this step
+      let validDS = true;
+      if (query.datasource) {
+        try {
+          await getDataSourceSrv().get(query.datasource.uid);
+        } catch {
+          console.error(`One of the queries has a datasource that is no longer available and was removed.`);
+          validDS = false;
+        }
+      }
+
+      if (validDS) {
+        allQueries.push({
+          ...query,
+          refId,
+          key,
+        });
+      }
     }
     return allQueries;
   }

public/app/features/explore/ExplorePaneContainer.tsx

@@ -28,6 +28,7 @@ import Explore from './Explore';
 import { initializeExplore, refreshExplore } from './state/explorePane';
 import { lastSavedUrl, cleanupPaneAction, stateSave } from './state/main';
 import { importQueries } from './state/query';
+import { loadAndInitDatasource } from './state/utils';
 
 const getStyles = (theme: GrafanaTheme2) => {
   return {
@@ -70,7 +71,7 @@ class ExplorePaneContainerUnconnected extends React.PureComponent<Props> {
   }
 
   async componentDidMount() {
-    const { initialized, exploreId, initialDatasource, initialQueries, initialRange, panelsState } = this.props;
+    const { initialized, exploreId, initialDatasource, initialQueries, initialRange, panelsState, orgId } = this.props;
     const width = this.el?.offsetWidth ?? 0;
     // initialize the whole explore first time we mount and if browser history contains a change in datasource
     if (!initialized) {
@@ -81,8 +82,8 @@ class ExplorePaneContainerUnconnected extends React.PureComponent<Props> {
         const isDSMixed =
           initialDatasource === MIXED_DATASOURCE_NAME || initialDatasource.uid === MIXED_DATASOURCE_NAME;
         if (!isDSMixed) {
-          const datasource = await getDatasourceSrv().get(initialDatasource);
-          queriesDatasourceOverride = datasource.getRef();
+          const { instance } = await loadAndInitDatasource(orgId, initialDatasource);
+          queriesDatasourceOverride = instance.getRef();
         }
       }
 
@@ -173,6 +174,7 @@ function mapStateToProps(state: StoreState, props: OwnProps) {
     initialQueries: queries,
     initialRange,
     panelsState,
+    orgId: state.user.orgId,
   };
 }
 

public/app/features/explore/Wrapper.test.tsx

@@ -464,6 +464,35 @@ describe('Wrapper', () => {
         'orgId=1&left={"datasource":"loki-uid","queries":[{"refId":"A","datasource":{"type":"logs","uid":"loki-uid"}}],"range":{"from":"now-1h","to":"now"}}'
       );
     });
+
+    it('Datasource in root not found and no queries changes to default', async () => {
+      setupExplore({
+        urlParams: 'orgId=1&left={"datasource":"asdasdasd","range":{"from":"now-1h","to":"now"}}',
+        prevUsedDatasource: { orgId: 1, datasource: 'elastic' },
+      });
+      await waitForExplore();
+      const urlParams = decodeURIComponent(locationService.getSearch().toString());
+      expect(urlParams).toBe(
+        'orgId=1&left={"datasource":"loki-uid","queries":[{"refId":"A","datasource":{"type":"logs","uid":"loki-uid"}}],"range":{"from":"now-1h","to":"now"}}'
+      );
+    });
+
+    it('Datasource root is mixed and there are two queries, one with datasource not found, only one query remains with root datasource as that datasource', async () => {
+      const consoleErrorSpy = jest.spyOn(console, 'error').mockImplementation();
+      setupExplore({
+        urlParams:
+          'orgId=1&left={"datasource":"-- Mixed --","queries":[{"refId":"A","datasource":{"type":"asdf","uid":"asdf"}},{"refId":"B","datasource":{"type":"logs","uid":"elastic-uid"}}],"range":{"from":"now-1h","to":"now"}}',
+        prevUsedDatasource: { orgId: 1, datasource: 'elastic' },
+      });
+      await waitForExplore();
+      const urlParams = decodeURIComponent(locationService.getSearch().toString());
+      expect(urlParams).toBe(
+        'orgId=1&left={"datasource":"elastic-uid","queries":[{"refId":"B","datasource":{"type":"logs","uid":"elastic-uid"}}],"range":{"from":"now-1h","to":"now"}}'
+      );
+      expect(consoleErrorSpy).toBeCalledTimes(1);
+
+      consoleErrorSpy.mockRestore();
+    });
   });
 
   it('removes `from` and `to` parameters from url when first mounted', async () => {

public/app/features/explore/state/query.test.ts

@@ -40,6 +40,25 @@ import { makeExplorePaneState } from './utils';
 
 const { testRange, defaultInitialState } = createDefaultInitialState();
 
+const datasources: DataSourceApi[] = [
+  {
+    name: 'testDs',
+    type: 'postgres',
+    uid: 'ds1',
+    getRef: () => {
+      return { type: 'postgres', uid: 'ds1' };
+    },
+  } as DataSourceApi<DataQuery, DataSourceJsonData, {}>,
+  {
+    name: 'testDs2',
+    type: 'postgres',
+    uid: 'ds2',
+    getRef: () => {
+      return { type: 'postgres', uid: 'ds2' };
+    },
+  } as DataSourceApi<DataQuery, DataSourceJsonData, {}>,
+];
+
 jest.mock('app/features/dashboard/services/TimeSrv', () => ({
   ...jest.requireActual('app/features/dashboard/services/TimeSrv'),
   getTimeSrv: () => ({
@@ -53,6 +72,11 @@ jest.mock('@grafana/runtime', () => ({
   getTemplateSrv: () => ({
     updateTimeRange: jest.fn(),
   }),
+  getDataSourceSrv: () => {
+    return {
+      get: (uid?: string) => datasources.find((ds) => ds.uid === uid) || datasources[0],
+    };
+  },
 }));
 
 function setupQueryResponse(state: StoreState) {
@@ -154,25 +178,6 @@ describe('running queries', () => {
 describe('importing queries', () => {
   describe('when importing queries between the same type of data source', () => {
     it('remove datasource property from all of the queries', async () => {
-      const datasources: DataSourceApi[] = [
-        {
-          name: 'testDs',
-          type: 'postgres',
-          uid: 'ds1',
-          getRef: () => {
-            return { type: 'postgres', uid: 'ds1' };
-          },
-        } as DataSourceApi<DataQuery, DataSourceJsonData, {}>,
-        {
-          name: 'testDs2',
-          type: 'postgres',
-          uid: 'ds2',
-          getRef: () => {
-            return { type: 'postgres', uid: 'ds2' };
-          },
-        } as DataSourceApi<DataQuery, DataSourceJsonData, {}>,
-      ];
-
       const { dispatch, getState }: { dispatch: ThunkDispatch; getState: () => StoreState } = configureStore({
         ...(defaultInitialState as any),
         explore: {