Pyroscope: Added app integration for datasource (#75789)

* feat: integrate pyroscope query editor with link extensions

- allows plugin app extensions to register links based on
  the query and selected datasource

* fix: remove not-as-designed busy wait mechanism

* Apply suggestions from code review

Co-authored-by: Marcus Andersson <marcus.andersson@grafana.com>

* fix: implement feedback

* fix: lint

---------

Co-authored-by: Marcus Andersson <marcus.andersson@grafana.com>

public/app/plugins/datasource/grafana-pyroscope-datasource/QueryEditor/QueryEditor.test.tsx

@@ -3,13 +3,20 @@ import userEvent from '@testing-library/user-event';
 import React from 'react';
 
 import { CoreApp, PluginType } from '@grafana/data';
+import { setPluginExtensionGetter } from '@grafana/runtime';
 
 import { PyroscopeDataSource } from '../datasource';
+import { mockFetchPyroscopeDatasourceSettings } from '../datasource.test';
 import { ProfileTypeMessage } from '../types';
 
 import { Props, QueryEditor } from './QueryEditor';
 
 describe('QueryEditor', () => {
+  beforeEach(() => {
+    setPluginExtensionGetter(() => ({ extensions: [] })); // No extensions
+    mockFetchPyroscopeDatasourceSettings();
+  });
+
   it('should render without error', async () => {
     setup();
 

public/app/plugins/datasource/grafana-pyroscope-datasource/QueryEditor/QueryEditor.tsx

@@ -12,6 +12,7 @@ import { EditorRow } from './EditorRow';
 import { EditorRows } from './EditorRows';
 import { LabelsEditor } from './LabelsEditor';
 import { ProfileTypesCascader, useProfileTypes } from './ProfileTypesCascader';
+import { PyroscopeQueryLinkExtensions } from './QueryLinkExtension';
 import { QueryOptions } from './QueryOptions';
 
 export type Props = QueryEditorProps<PyroscopeDataSource, Query, PyroscopeDataSourceOptions>;
@@ -57,6 +58,7 @@ export function QueryEditor(props: Props) {
           labels={labels}
           getLabelValues={getLabelValues}
         />
+        <PyroscopeQueryLinkExtensions {...props} />
       </EditorRow>
       <EditorRow>
         <QueryOptions query={query} onQueryChange={props.onChange} app={props.app} labels={labels} />

public/app/plugins/datasource/grafana-pyroscope-datasource/QueryEditor/QueryLinkExtension.test.tsx

@@ -0,0 +1,125 @@
+import { render, screen } from '@testing-library/react';
+import React from 'react';
+import { act } from 'react-dom/test-utils';
+
+import { PluginType, rangeUtil, PluginExtensionLink, PluginExtensionTypes } from '@grafana/data';
+import { getPluginLinkExtensions } from '@grafana/runtime';
+
+import { PyroscopeDataSource } from '../datasource';
+import { mockFetchPyroscopeDatasourceSettings } from '../datasource.test';
+
+import { Props, PyroscopeQueryLinkExtensions, resetPyroscopeQueryLinkExtensionsFetches } from './QueryLinkExtension';
+
+// Constants copied from `QueryLinkExtension.tsx`
+const EXTENSION_POINT_ID = 'plugins/grafana-pyroscope-datasource/query-links';
+
+jest.mock('@grafana/runtime', () => ({
+  ...jest.requireActual('@grafana/runtime'),
+  setPluginExtensionGetter: jest.fn(),
+  getPluginLinkExtensions: jest.fn(),
+}));
+
+const getPluginLinkExtensionsMock = jest.mocked(getPluginLinkExtensions);
+
+const defaultPyroscopeDataSourceSettings = {
+  uid: 'default-pyroscope',
+  url: 'http://pyroscope',
+  basicAuthUser: 'pyroscope_user',
+};
+
+describe('PyroscopeQueryLinkExtensions', () => {
+  const EXPECTED_BUTTON_LABEL = 'Profiles App';
+  const DEFAULT_EXTENSION_PATH = 'a/mock-path-app/fake-path';
+
+  function createExtension(overrides?: Partial<PluginExtensionLink>) {
+    return {
+      ...{
+        description: 'unremarkable-description',
+        extensionPointId: EXTENSION_POINT_ID,
+        title: EXPECTED_BUTTON_LABEL,
+        path: DEFAULT_EXTENSION_PATH,
+        type: PluginExtensionTypes.link,
+        category: 'unremarkable-category',
+        icon: 'heart',
+        onClick() {},
+        pluginId: 'mock-path-app',
+        id: `${Date.now()}}`,
+      },
+      ...overrides,
+    } as PluginExtensionLink;
+  }
+
+  beforeEach(() => {
+    resetPyroscopeQueryLinkExtensionsFetches();
+    mockFetchPyroscopeDatasourceSettings(defaultPyroscopeDataSourceSettings);
+
+    getPluginLinkExtensionsMock.mockRestore();
+    getPluginLinkExtensionsMock.mockReturnValue({ extensions: [] }); // Unless stated otherwise, no extensions
+  });
+
+  it('should render if extension present', async () => {
+    getPluginLinkExtensionsMock.mockReturnValue({ extensions: [createExtension()] }); // Default extension
+
+    await act(setup);
+    expect(await screen.findAllByText(EXPECTED_BUTTON_LABEL)).toBeDefined();
+  });
+
+  it('Should not render if no extension present', async () => {
+    await act(setup);
+    expect(screen.queryByText(EXPECTED_BUTTON_LABEL)).toBeNull();
+  });
+});
+
+function setupDs() {
+  const ds = new PyroscopeDataSource({
+    ...defaultPyroscopeDataSourceSettings,
+    name: 'test',
+    type: PluginType.datasource,
+    access: 'proxy',
+    id: 1,
+    jsonData: {},
+    meta: {
+      name: '',
+      id: '',
+      type: PluginType.datasource,
+      baseUrl: '',
+      info: {
+        author: {
+          name: '',
+        },
+        description: '',
+        links: [],
+        logos: {
+          large: '',
+          small: '',
+        },
+        screenshots: [],
+        updated: '',
+        version: '',
+      },
+      module: '',
+    },
+    readOnly: false,
+  });
+
+  return ds;
+}
+
+async function setup(options: { props: Partial<Props> } = { props: {} }) {
+  const utils = render(
+    <PyroscopeQueryLinkExtensions
+      query={{
+        queryType: 'both',
+        labelSelector: '',
+        profileTypeId: 'process_cpu:cpu',
+        refId: 'A',
+        maxNodes: 1000,
+        groupBy: [],
+      }}
+      datasource={setupDs()}
+      range={rangeUtil.convertRawToRange({ from: 'now-1h', to: 'now' })}
+      {...options.props}
+    />
+  );
+  return { ...utils };
+}

public/app/plugins/datasource/grafana-pyroscope-datasource/QueryEditor/QueryLinkExtension.tsx

@@ -0,0 +1,104 @@
+import { css } from '@emotion/css';
+import React from 'react';
+import { useAsync } from 'react-use';
+
+import { GrafanaTheme2, QueryEditorProps, TimeRange } from '@grafana/data';
+import { getBackendSrv, getPluginLinkExtensions } from '@grafana/runtime';
+import { LinkButton, useStyles2 } from '@grafana/ui';
+
+import { PyroscopeDataSource } from '../datasource';
+import { PyroscopeDataSourceOptions, Query } from '../types';
+
+const EXTENSION_POINT_ID = 'plugins/grafana-pyroscope-datasource/query-links';
+
+/** A subset of the datasource settings that are relevant for this integration */
+type PyroscopeDatasourceSettings = {
+  uid: string;
+  url: string;
+  type: string;
+  basicAuthUser: string;
+};
+
+/** The context object that will be shared with the link extension's configure function */
+type ExtensionQueryLinksContext = {
+  datasourceUid: string;
+  query: Query;
+  range?: TimeRange | undefined;
+  datasourceSettings?: PyroscopeDatasourceSettings;
+};
+
+/* Global promises to fetch pyroscope datasource settings by uid as encountered */
+const pyroscopeDatasourceSettingsByUid: Record<string, PyroscopeDatasourceSettings> = {};
+
+/* Reset promises for testing purposes */
+export function resetPyroscopeQueryLinkExtensionsFetches() {
+  Object.keys(pyroscopeDatasourceSettingsByUid).forEach((key) => delete pyroscopeDatasourceSettingsByUid[key]);
+}
+
+/** A subset of the `PyroscopeDataSource` `QueryEditorProps` */
+export type Props = Pick<
+  QueryEditorProps<PyroscopeDataSource, Query, PyroscopeDataSourceOptions>,
+  'datasource' | 'query' | 'range'
+>;
+
+export function PyroscopeQueryLinkExtensions(props: Props) {
+  const {
+    datasource: { uid: datasourceUid },
+    query,
+    range,
+  } = props;
+
+  const { value: datasourceSettings } = useAsync(async () => {
+    if (pyroscopeDatasourceSettingsByUid[datasourceUid]) {
+      return pyroscopeDatasourceSettingsByUid[datasourceUid];
+    }
+    const settings = await getBackendSrv().get<PyroscopeDatasourceSettings>(`/api/datasources/uid/${datasourceUid}`);
+    pyroscopeDatasourceSettingsByUid[datasourceUid] = settings;
+    return settings;
+  }, [datasourceUid]);
+
+  const context: ExtensionQueryLinksContext = {
+    datasourceUid,
+    query,
+    range,
+    datasourceSettings,
+  };
+
+  const { extensions } = getPluginLinkExtensions({
+    extensionPointId: EXTENSION_POINT_ID,
+    context,
+  });
+
+  const styles = useStyles2(getStyles);
+
+  if (extensions.length === 0) {
+    return null;
+  }
+
+  return (
+    <>
+      {extensions.map((extension) => (
+        <LinkButton
+          className={styles.linkButton}
+          key={`${extension.id}`}
+          variant="secondary"
+          icon={extension.icon || 'external-link-alt'}
+          tooltip={extension.description}
+          target="_blank"
+          href={extension.path}
+          onClick={extension.onClick}
+        >
+          {extension.title}
+        </LinkButton>
+      ))}
+    </>
+  );
+}
+
+function getStyles(theme: GrafanaTheme2) {
+  return {
+    linkButton: css({
+      marginLeft: theme.spacing(1),
+    }),
+  };
+}

public/app/plugins/datasource/grafana-pyroscope-datasource/datasource.test.ts

@@ -1,13 +1,43 @@
-import { AbstractLabelOperator, CoreApp, DataSourceInstanceSettings, PluginMetaInfo, PluginType } from '@grafana/data';
+import {
+  AbstractLabelOperator,
+  CoreApp,
+  DataSourceInstanceSettings,
+  PluginMetaInfo,
+  PluginType,
+  DataSourceJsonData,
+} from '@grafana/data';
+import { setPluginExtensionGetter, getBackendSrv, setBackendSrv } from '@grafana/runtime';
 import { TemplateSrv } from 'app/features/templating/template_srv';
 
 import { defaultPyroscopeQueryType } from './dataquery.gen';
 import { normalizeQuery, PyroscopeDataSource } from './datasource';
 import { Query } from './types';
 
+/** The datasource QueryEditor fetches datasource settings to send to the extension's `configure` method */
+export function mockFetchPyroscopeDatasourceSettings(
+  datasourceSettings?: Partial<DataSourceInstanceSettings<DataSourceJsonData>>
+) {
+  const settings = { ...defaultSettings, ...datasourceSettings };
+  const returnValues: Record<string, unknown> = {
+    [`/api/datasources/uid/${settings.uid}`]: settings,
+  };
+  setBackendSrv({
+    ...getBackendSrv(),
+    get: function <T>(path: string) {
+      const value = returnValues[path];
+      if (value) {
+        return Promise.resolve(value as T);
+      }
+      return Promise.reject({ message: 'reject' });
+    },
+  });
+}
+
 describe('Pyroscope data source', () => {
   let ds: PyroscopeDataSource;
   beforeEach(() => {
+    mockFetchPyroscopeDatasourceSettings();
+    setPluginExtensionGetter(() => ({ extensions: [] })); // No extensions
     ds = new PyroscopeDataSource(defaultSettings);
   });