Add query migration handlers (#93735)

packages/grafana-runtime/src/components/QueryEditorWithMigration.test.tsx

@@ -0,0 +1,102 @@
+import { render, screen, waitFor } from '@testing-library/react';
+
+import { DataSourceInstanceSettings, QueryEditorProps } from '@grafana/data';
+import { DataQuery, DataSourceJsonData } from '@grafana/schema';
+
+import { config } from '../config';
+import { BackendSrv, BackendSrvRequest } from '../services';
+import { DataSourceWithBackend } from '../utils/DataSourceWithBackend';
+import { MigrationHandler } from '../utils/migrationHandler';
+
+import { QueryEditorWithMigration } from './QueryEditorWithMigration';
+
+const backendSrv = {
+  post<T = unknown>(url: string, data?: unknown, options?: Partial<BackendSrvRequest>): Promise<T> {
+    return mockDatasourcePost({ url, data, ...options });
+  },
+} as unknown as BackendSrv;
+
+jest.mock('../services', () => ({
+  ...jest.requireActual('../services'),
+  getBackendSrv: () => backendSrv,
+}));
+
+let mockDatasourcePost = jest.fn();
+
+interface MyQuery extends DataQuery {}
+
+class MyDataSource extends DataSourceWithBackend<MyQuery, DataSourceJsonData> implements MigrationHandler {
+  hasBackendMigration: boolean;
+
+  constructor(instanceSettings: DataSourceInstanceSettings<DataSourceJsonData>) {
+    super(instanceSettings);
+    this.hasBackendMigration = true;
+  }
+
+  shouldMigrate(query: DataQuery): boolean {
+    return true;
+  }
+}
+
+type Props = QueryEditorProps<MyDataSource, MyQuery, DataSourceJsonData>;
+
+function QueryEditor(props: Props) {
+  return <div>{JSON.stringify(props.query)}</div>;
+}
+
+function createMockDatasource(otherSettings?: Partial<DataSourceInstanceSettings<DataSourceJsonData>>) {
+  const settings = {
+    name: 'test',
+    id: 1234,
+    uid: 'abc',
+    type: 'dummy',
+    jsonData: {},
+    ...otherSettings,
+  } as DataSourceInstanceSettings<DataSourceJsonData>;
+
+  return new MyDataSource(settings);
+}
+
+describe('QueryEditorWithMigration', () => {
+  const originalFeatureToggles = config.featureToggles;
+  beforeEach(() => {
+    config.featureToggles = { ...originalFeatureToggles, grafanaAPIServerWithExperimentalAPIs: true };
+  });
+  afterEach(() => {
+    config.featureToggles = originalFeatureToggles;
+  });
+
+  it('should migrate a query', async () => {
+    const WithMigration = QueryEditorWithMigration(QueryEditor);
+    const ds = createMockDatasource();
+    const originalQuery = { refId: 'A', datasource: { type: 'dummy' }, foo: 'bar' };
+    const migratedQuery = { refId: 'A', datasource: { type: 'dummy' }, foobar: 'barfoo' };
+
+    mockDatasourcePost = jest.fn().mockImplementation((args: { url: string; data: unknown }) => {
+      expect(args.url).toBe('/apis/dummy.datasource.grafana.app/v0alpha1/namespaces/default/queryconvert');
+      expect(args.data).toMatchObject({ queries: [originalQuery] });
+      return Promise.resolve({ queries: [{ JSON: migratedQuery }] });
+    });
+
+    render(<WithMigration datasource={ds} query={originalQuery} onChange={jest.fn()} onRunQuery={jest.fn()} />);
+
+    await waitFor(() => {
+      // Check that migratedQuery is rendered
+      expect(screen.getByText(JSON.stringify(migratedQuery))).toBeInTheDocument();
+    });
+  });
+
+  it('should render a Skeleton while migrating', async () => {
+    const WithMigration = QueryEditorWithMigration(QueryEditor);
+    const ds = createMockDatasource();
+    const originalQuery = { refId: 'A', datasource: { type: 'dummy' }, foo: 'bar' };
+
+    mockDatasourcePost = jest.fn().mockImplementation(async (args: { url: string; data: unknown }) => {
+      await waitFor(() => {}, { timeout: 5000 });
+      return Promise.resolve({ queries: [{ JSON: originalQuery }] });
+    });
+
+    render(<WithMigration datasource={ds} query={originalQuery} onChange={jest.fn()} onRunQuery={jest.fn()} />);
+    expect(screen.getByTestId('react-loading-skeleton-testid')).toBeInTheDocument();
+  });
+});

packages/grafana-runtime/src/components/QueryEditorWithMigration.tsx

@@ -0,0 +1,44 @@
+import React, { useEffect, useState } from 'react';
+import Skeleton from 'react-loading-skeleton';
+
+import { DataSourceApi, DataSourceOptionsType, DataSourceQueryType, QueryEditorProps } from '@grafana/data';
+import { DataQuery, DataSourceJsonData } from '@grafana/schema';
+
+import { isMigrationHandler, migrateQuery } from '../utils/migrationHandler';
+
+/**
+ * @alpha Experimental: QueryEditorWithMigration is a higher order component that wraps the QueryEditor component
+ * and ensures that the query is migrated before being passed to the QueryEditor.
+ */
+export function QueryEditorWithMigration<
+  DSType extends DataSourceApi<TQuery, TOptions>,
+  TQuery extends DataQuery = DataSourceQueryType<DSType>,
+  TOptions extends DataSourceJsonData = DataSourceOptionsType<DSType>,
+>(QueryEditor: React.ComponentType<QueryEditorProps<DSType, TQuery, TOptions>>) {
+  const WithExtra = (props: QueryEditorProps<DSType, TQuery, TOptions>) => {
+    const [migrated, setMigrated] = useState(false);
+    const [query, setQuery] = useState(props.query);
+
+    useEffect(() => {
+      if (props.query && isMigrationHandler(props.datasource)) {
+        migrateQuery(props.datasource, props.query).then((migrated) => {
+          props.onChange(migrated);
+          setQuery(migrated);
+          setMigrated(true);
+        });
+      } else {
+        setMigrated(true);
+      }
+    }, []); // eslint-disable-line react-hooks/exhaustive-deps
+
+    useEffect(() => {
+      setQuery(props.query);
+    }, [props.query]);
+
+    if (!migrated) {
+      return <Skeleton containerTestId="react-loading-skeleton-testid" height={75} />;
+    }
+    return <QueryEditor {...props} query={query} />;
+  };
+  return WithExtra;
+}

packages/grafana-runtime/src/index.ts

@@ -54,3 +54,5 @@ export { setReturnToPreviousHook, useReturnToPrevious } from './utils/returnToPr
 export { setChromeHeaderHeightHook, useChromeHeaderHeight } from './utils/chromeHeaderHeight';
 export { type EmbeddedDashboardProps, EmbeddedDashboard, setEmbeddedDashboard } from './components/EmbeddedDashboard';
 export { hasPermission, hasPermissionInMetadata, hasAllPermissions, hasAnyPermission } from './utils/rbac';
+export { QueryEditorWithMigration } from './components/QueryEditorWithMigration';
+export { type MigrationHandler, isMigrationHandler, migrateQuery, migrateRequest } from './utils/migrationHandler';

packages/grafana-runtime/src/utils/migrationHandler.test.ts

@@ -0,0 +1,174 @@
+import { BackendSrv, BackendSrvRequest } from 'src/services';
+
+import { DataQueryRequest, DataSourceInstanceSettings } from '@grafana/data';
+import { DataQuery, DataSourceJsonData } from '@grafana/schema';
+
+import { config } from '../config';
+
+import { DataSourceWithBackend } from './DataSourceWithBackend';
+import { isMigrationHandler, migrateQuery, migrateRequest, MigrationHandler } from './migrationHandler';
+
+let mockDatasourcePost = jest.fn();
+
+interface MyQuery extends DataQuery {}
+
+class MyDataSourceWithoutMigration extends DataSourceWithBackend<MyQuery, DataSourceJsonData> {
+  constructor(instanceSettings: DataSourceInstanceSettings<DataSourceJsonData>) {
+    super(instanceSettings);
+  }
+}
+
+class MyDataSource extends DataSourceWithBackend<MyQuery, DataSourceJsonData> implements MigrationHandler {
+  hasBackendMigration: boolean;
+  constructor(instanceSettings: DataSourceInstanceSettings<DataSourceJsonData>) {
+    super(instanceSettings);
+    this.hasBackendMigration = true;
+  }
+
+  shouldMigrate(query: DataQuery): boolean {
+    return true;
+  }
+}
+
+const backendSrv = {
+  post<T = unknown>(url: string, data?: unknown, options?: Partial<BackendSrvRequest>): Promise<T> {
+    return mockDatasourcePost({ url, data, ...options });
+  },
+} as unknown as BackendSrv;
+
+jest.mock('../services', () => ({
+  ...jest.requireActual('../services'),
+  getBackendSrv: () => backendSrv,
+}));
+
+describe('query migration', () => {
+  // Configure config.featureToggles.grafanaAPIServerWithExperimentalAPIs
+  const originalFeatureToggles = config.featureToggles;
+  beforeEach(() => {
+    config.featureToggles = { ...originalFeatureToggles, grafanaAPIServerWithExperimentalAPIs: true };
+  });
+  afterEach(() => {
+    config.featureToggles = originalFeatureToggles;
+  });
+
+  describe('isMigrationHandler', () => {
+    it('returns true for a datasource with backend migration', () => {
+      const ds = createMockDatasource();
+      expect(isMigrationHandler(ds)).toBe(true);
+    });
+
+    it('returns false for a datasource without backend migration', () => {
+      const ds = new MyDataSourceWithoutMigration({} as DataSourceInstanceSettings<DataSourceJsonData>); // eslint-disable-line @typescript-eslint/no-explicit-any
+      expect(isMigrationHandler(ds)).toBe(false);
+    });
+  });
+
+  describe('migrateQuery', () => {
+    it('skips migration if the datasource does not support it', async () => {
+      const ds = createMockDatasource();
+      ds.hasBackendMigration = false;
+      const query = { refId: 'A', datasource: { type: 'dummy' } };
+
+      const result = await migrateQuery(ds, query);
+
+      expect(query).toEqual(result);
+      expect(mockDatasourcePost).not.toHaveBeenCalled();
+    });
+
+    it('skips migration if the query should not be migrated', async () => {
+      const ds = createMockDatasource();
+      ds.shouldMigrate = jest.fn().mockReturnValue(false);
+      const query = { refId: 'A', datasource: { type: 'dummy' } };
+
+      const result = await migrateQuery(ds, query);
+
+      expect(query).toEqual(result);
+      expect(mockDatasourcePost).not.toHaveBeenCalled();
+    });
+
+    it('check that migrateQuery works', async () => {
+      const ds = createMockDatasource();
+
+      const originalQuery = { refId: 'A', datasource: { type: 'dummy' }, foo: 'bar' };
+      const migratedQuery = { refId: 'A', datasource: { type: 'dummy' }, foobar: 'barfoo' };
+      mockDatasourcePost = jest.fn().mockImplementation((args: { url: string; data: unknown }) => {
+        expect(args.url).toBe('/apis/dummy.datasource.grafana.app/v0alpha1/namespaces/default/queryconvert');
+        expect(args.data).toMatchObject({ queries: [originalQuery] });
+        return Promise.resolve({ queries: [{ JSON: migratedQuery }] });
+      });
+
+      const result = await migrateQuery(ds, originalQuery);
+
+      expect(migratedQuery).toEqual(result);
+    });
+  });
+
+  describe('migrateRequest', () => {
+    it('skips migration if the datasource does not support it', async () => {
+      const ds = createMockDatasource();
+      ds.hasBackendMigration = false;
+      const request = {
+        targets: [{ refId: 'A', datasource: { type: 'dummy' } }],
+      } as unknown as DataQueryRequest<MyQuery>; // eslint-disable-line @typescript-eslint/no-explicit-any
+
+      const result = await migrateRequest(ds, request);
+
+      expect(request).toEqual(result);
+      expect(mockDatasourcePost).not.toHaveBeenCalled();
+    });
+
+    it('skips migration if none of the queries should be migrated', async () => {
+      const ds = createMockDatasource();
+      ds.shouldMigrate = jest.fn().mockReturnValue(false);
+      const request = {
+        targets: [{ refId: 'A', datasource: { type: 'dummy' } }],
+      } as unknown as DataQueryRequest<MyQuery>; // eslint-disable-line @typescript-eslint/no-explicit-any
+
+      const result = await migrateRequest(ds, request);
+
+      expect(request).toEqual(result);
+      expect(mockDatasourcePost).not.toHaveBeenCalled();
+    });
+
+    it('check that migrateRequest migrates a request', async () => {
+      const ds = createMockDatasource();
+
+      const originalRequest = {
+        targets: [
+          { refId: 'A', datasource: { type: 'dummy' }, foo: 'bar' },
+          { refId: 'A', datasource: { type: 'dummy' }, bar: 'foo' },
+        ],
+      } as unknown as DataQueryRequest<MyQuery>; // eslint-disable-line @typescript-eslint/no-explicit-any
+      const migratedRequest = {
+        targets: [
+          { refId: 'A', datasource: { type: 'dummy' }, foobar: 'foobar' },
+          { refId: 'A', datasource: { type: 'dummy' }, barfoo: 'barfoo' },
+        ],
+      };
+      mockDatasourcePost = jest.fn().mockImplementation((args: { url: string; data: unknown }) => {
+        expect(args.url).toBe('/apis/dummy.datasource.grafana.app/v0alpha1/namespaces/default/queryconvert');
+        expect(args.data).toMatchObject({ queries: originalRequest.targets });
+        return Promise.resolve({ queries: migratedRequest.targets.map((query) => ({ JSON: query })) });
+      });
+
+      const result = await migrateRequest(ds, originalRequest);
+
+      expect(migratedRequest).toEqual(result);
+    });
+  });
+});
+
+function createMockDatasource(otherSettings?: Partial<DataSourceInstanceSettings<DataSourceJsonData>>) {
+  const settings = {
+    name: 'test',
+    id: 1234,
+    uid: 'abc',
+    type: 'dummy',
+    jsonData: {},
+    ...otherSettings,
+  } as DataSourceInstanceSettings<DataSourceJsonData>;
+
+  mockDatasourcePost.mockReset();
+
+  return new MyDataSource(settings);
+}

packages/grafana-runtime/src/utils/migrationHandler.ts

@@ -0,0 +1,72 @@
+import { DataQueryRequest } from '@grafana/data';
+import { DataQuery } from '@grafana/schema';
+
+import { config } from '../config';
+import { getBackendSrv } from '../services';
+
+import { DataSourceWithBackend } from './DataSourceWithBackend';
+
+/**
+ * @alpha Experimental: Plugins implementing MigrationHandler interface will automatically have their queries migrated.
+ */
+export interface MigrationHandler {
+  hasBackendMigration: boolean;
+  shouldMigrate(query: DataQuery): boolean;
+}
+
+export function isMigrationHandler(object: unknown): object is MigrationHandler {
+  return object instanceof DataSourceWithBackend && 'hasBackendMigration' in object && 'shouldMigrate' in object;
+}
+
+async function postMigrateRequest<TQuery extends DataQuery>(queries: TQuery[]): Promise<TQuery[]> {
+  if (!(config.featureToggles.grafanaAPIServerWithExperimentalAPIs || config.featureToggles.datasourceAPIServers)) {
+    console.warn('migrateQuery is only available with the experimental API server');
+    return queries;
+  }
+
+  // Obtaining the GroupName from the plugin ID as done in the backend, this is temporary until we have a better way to obtain it
+  // https://github.com/grafana/grafana/blob/e013cd427cb0457177e11f19ebd30bc523b36c76/pkg/plugins/apiserver.go#L10
+  const dsnameURL = queries[0].datasource?.type?.replace(/^(grafana-)?(.*?)(-datasource)?$/, '$2');
+  const groupName = `${dsnameURL}.datasource.grafana.app`;
+  // Asuming apiVersion is v0alpha1, we'll need to obtain it from a trusted source
+  const apiVersion = 'v0alpha1';
+  const url = `/apis/${groupName}/${apiVersion}/namespaces/${config.namespace}/queryconvert`;
+  const request = {
+    queries: queries.map((query) => {
+      return {
+        ...query,
+        JSON: query, // JSON is not part of the type but it should be what holds the query
+      };
+    }),
+  };
+  const res = await getBackendSrv().post(url, request);
+  return res.queries.map((query: { JSON: TQuery }) => query.JSON);
+}
+
+/**
+ * @alpha Experimental: Calls migration endpoint with one query. Requires grafanaAPIServerWithExperimentalAPIs or datasourceAPIServers feature toggle.
+ */
+export async function migrateQuery<TQuery extends DataQuery>(
+  datasource: MigrationHandler,
+  query: TQuery
+): Promise<TQuery> {
+  if (!datasource.hasBackendMigration || !datasource.shouldMigrate(query)) {
+    return query;
+  }
+  const res = await postMigrateRequest([query]);
+  return res[0];
+}
+
+/**
+ * @alpha Experimental: Calls migration endpoint with multiple queries. Requires grafanaAPIServerWithExperimentalAPIs or datasourceAPIServers feature toggle.
+ */
+export async function migrateRequest<TQuery extends DataQuery>(
+  datasource: MigrationHandler,
+  request: DataQueryRequest<TQuery>
+): Promise<DataQueryRequest<TQuery>> {
+  if (!datasource.hasBackendMigration || !request.targets.some((query) => datasource.shouldMigrate(query))) {
+    return request;
+  }
+  const res = await postMigrateRequest(request.targets);
+  return { ...request, targets: res };
+}

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

@@ -20,7 +20,7 @@ import { Echo } from '../../../core/services/echo/Echo';
 import { createDashboardModelFixture } from '../../dashboard/state/__fixtures__/dashboardFixtures';
 
 import { getMockDataSource, TestQuery } from './__mocks__/mockDataSource';
-import { callQueryMethod, runRequest } from './runRequest';
+import { callQueryMethodWithMigration, runRequest } from './runRequest';
 
 jest.mock('app/core/services/backend_srv');
 
@@ -42,6 +42,14 @@ jest.mock('app/features/expressions/ExpressionDatasource', () => ({
   },
 }));
 
+let isMigrationHandlerMock = jest.fn().mockReturnValue(false);
+let migrateRequestMock = jest.fn();
+jest.mock('@grafana/runtime', () => ({
+  ...jest.requireActual('@grafana/runtime'),
+  isMigrationHandler: () => isMigrationHandlerMock(),
+  migrateRequest: () => migrateRequestMock(),
+}));
+
 class ScenarioCtx {
   ds!: DataSourceApi;
   request!: DataQueryRequest;
@@ -405,7 +413,7 @@ describe('runRequest', () => {
   });
 });
 
-describe('callQueryMethod', () => {
+describe('callQueryMethodWithMigration', () => {
   let request: DataQueryRequest<TestQuery>;
   let filterQuerySpy: jest.SpyInstance;
   let querySpy: jest.SpyInstance;
@@ -417,11 +425,13 @@ describe('callQueryMethod', () => {
     filterQuery,
     getDefaultQuery,
     queryFunction,
+    migrateRequest,
   }: {
     targets: TestQuery[];
     getDefaultQuery?: (app: CoreApp) => Partial<TestQuery>;
     filterQuery?: typeof ds.filterQuery;
     queryFunction?: typeof ds.query;
+    migrateRequest?: jest.Mock;
   }) => {
     request = {
       range: {
@@ -448,8 +458,12 @@ describe('callQueryMethod', () => {
       ds.getDefaultQuery = getDefaultQuery;
       defaultQuerySpy = jest.spyOn(ds, 'getDefaultQuery');
     }
+    if (migrateRequest) {
+      isMigrationHandlerMock = jest.fn().mockReturnValue(true);
+      migrateRequestMock = migrateRequest;
+    }
     querySpy = jest.spyOn(ds, 'query');
-    callQueryMethod(ds, request, queryFunction);
+    return callQueryMethodWithMigration(ds, request, queryFunction);
   };
 
   beforeEach(() => {
@@ -590,6 +604,48 @@ describe('callQueryMethod', () => {
       })
     );
   });
+
+  it('Should migrate a request if defined', (done) => {
+    const migrateRequest = jest.fn();
+    const res = setup({
+      targets: [
+        {
+          refId: 'A',
+          q: 'SUM(foo)',
+        },
+      ],
+      migrateRequest: migrateRequest.mockResolvedValue({
+        range: {
+          from: dateTime(),
+          to: dateTime(),
+          raw: { from: '1h', to: 'now' },
+        },
+        targets: [
+          {
+            refId: 'A',
+            qMigrated: 'SUM(foo)',
+          },
+        ],
+        requestId: '',
+        interval: '',
+        intervalMs: 0,
+        scopedVars: {},
+        timezone: '',
+        app: '',
+        startTime: 0,
+      }),
+    });
+    expect(migrateRequest).toHaveBeenCalledTimes(1);
+    res.subscribe((res) => {
+      expect(res).toBeDefined();
+      expect(querySpy).toHaveBeenCalledWith(
+        expect.objectContaining({
+          targets: [{ qMigrated: 'SUM(foo)', refId: 'A' }],
+        })
+      );
+      done();
+    });
+  });
 });
 
 const expectThatRangeHasNotMutated = (ctx: ScenarioCtx) => {

public/app/features/query/state/runRequest.ts

@@ -1,7 +1,7 @@
 // Libraries
 import { isString, map as isArray } from 'lodash';
 import { from, merge, Observable, of, timer } from 'rxjs';
-import { catchError, map, mapTo, share, takeUntil, tap } from 'rxjs/operators';
+import { catchError, map, mapTo, mergeMap, share, takeUntil, tap } from 'rxjs/operators';
 
 // Utils & Services
 // Types
@@ -18,7 +18,7 @@ import {
   PanelData,
   TimeRange,
 } from '@grafana/data';
-import { config, toDataQueryError } from '@grafana/runtime';
+import { config, isMigrationHandler, migrateRequest, toDataQueryError } from '@grafana/runtime';
 import { isExpressionReference } from '@grafana/runtime/src/utils/DataSourceWithBackend';
 import { backendSrv } from 'app/core/services/backend_srv';
 import { queryIsEmpty } from 'app/core/utils/query';
@@ -143,7 +143,7 @@ export function runRequest(
     return of(state.panelData);
   }
 
-  const dataObservable = callQueryMethod(datasource, request, queryFunction).pipe(
+  const dataObservable = callQueryMethodWithMigration(datasource, request, queryFunction).pipe(
     // Transform response packets into PanelData with merged results
     map((packet: DataQueryResponse) => {
       if (!isArray(packet.data)) {
@@ -186,6 +186,20 @@ export function runRequest(
   return merge(timer(200).pipe(mapTo(state.panelData), takeUntil(dataObservable)), dataObservable);
 }
 
+export function callQueryMethodWithMigration(
+  datasource: DataSourceApi,
+  request: DataQueryRequest,
+  queryFunction?: typeof datasource.query
+) {
+  if (isMigrationHandler(datasource)) {
+    const migratedRequestPromise = migrateRequest(datasource, request);
+    return from(migratedRequestPromise).pipe(
+      mergeMap((migratedRequest) => callQueryMethod(datasource, migratedRequest, queryFunction))
+    );
+  }
+  return callQueryMethod(datasource, request, queryFunction);
+}
+
 export function callQueryMethod(
   datasource: DataSourceApi,
   request: DataQueryRequest,