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 };
+}