feat(FlameGraph): Add missing export menu

.prettierignore

@@ -10,4 +10,4 @@ e2e/test-reports
 e2e/test-results
 
 # the PprofRequest class uses decorators - FIXME or don't use them
-PprofRequest.ts
+PprofRequest*.ts

src/pages/ProfilesExplorerView/components/SceneAiPanel/infrastructure/useFetchDotProfiles.ts

@@ -1,9 +1,9 @@
 import { TimeRange } from '@grafana/data';
 import { useQuery } from '@tanstack/react-query';
 
+import { ProfileApiClient } from '../../../infrastructure/profiles/ProfileApiClient';
 import { DataSourceProxyClientBuilder } from '../../../infrastructure/series/http/DataSourceProxyClientBuilder';
 import { cleanupDotResponse } from './cleanupDotResponse';
-import { ProfileApiClient } from './ProfileApiClient';
 
 export type FetchParams = Array<{
   query: string;
@@ -26,7 +26,9 @@ export function useFetchDotProfiles(dataSourceUid: string, fetchParams: FetchPar
       // TODO: pass a signal options to properly abort all in-flight requests
       return Promise.all(
         fetchParams.map(({ query, timeRange }) =>
-          profileApiClient.get({ query, timeRange, format: 'dot', maxNodes: MAX_NODES }).then(cleanupDotResponse)
+          profileApiClient
+            .get({ query, timeRange, format: 'dot', maxNodes: MAX_NODES })
+            .then((response) => cleanupDotResponse(response as string))
         )
       );
     },

src/pages/ProfilesExplorerView/components/SceneExploreServiceFlameGraph/SceneFlameGraph.tsx

@@ -22,11 +22,13 @@ import { buildFlameGraphQueryRunner } from '../../infrastructure/flame-graph/bui
 import { PYROSCOPE_DATA_SOURCE } from '../../infrastructure/pyroscope-data-sources';
 import { AIButton } from '../SceneAiPanel/components/AiButton/AIButton';
 import { SceneAiPanel } from '../SceneAiPanel/SceneAiPanel';
+import { SceneExportMenu } from './components/SceneExportMenu';
 
 interface SceneFlameGraphState extends SceneObjectState {
   $data: SceneQueryRunner;
   lastTimeRange?: TimeRange;
   aiPanel: SceneAiPanel;
+  exportMenu: SceneExportMenu;
 }
 
 // I've tried to use a SplitLayout for the body without any success (left: flame graph, right: explain flame graph content)
@@ -41,6 +43,7 @@ export class SceneFlameGraph extends SceneObjectBase<SceneFlameGraphState> {
       }),
       lastTimeRange: undefined,
       aiPanel: new SceneAiPanel(),
+      exportMenu: new SceneExportMenu(),
     });
 
     this.addActivationHandler(this.onActivate.bind(this));
@@ -90,7 +93,7 @@ export class SceneFlameGraph extends SceneObjectBase<SceneFlameGraphState> {
 
     const [maxNodes] = useMaxNodesFromUrl();
     const { settings, error: isFetchingSettingsError } = useFetchPluginSettings();
-    const { $data, aiPanel, lastTimeRange } = this.useState();
+    const { $data, lastTimeRange, aiPanel, exportMenu } = this.useState();
 
     if (isFetchingSettingsError) {
       displayWarning([
@@ -126,6 +129,11 @@ export class SceneFlameGraph extends SceneObjectBase<SceneFlameGraphState> {
           panel: aiPanel,
           fetchParams: [{ query, timeRange: lastTimeRange }],
         },
+        export: {
+          menu: exportMenu,
+          query,
+          timeRange: lastTimeRange,
+        },
       },
       actions: {
         getTheme,
@@ -179,6 +187,13 @@ export class SceneFlameGraph extends SceneObjectBase<SceneFlameGraphState> {
             disableCollapsing={!data.settings?.collapsedFlamegraphs}
             getTheme={actions.getTheme as any}
             getExtraContextMenuButtons={gitHubIntegration.actions.getExtraFlameGraphMenuItems}
+            extraHeaderElements={
+              <data.export.menu.Component
+                model={data.export.menu}
+                query={data.export.query}
+                timeRange={data.export.timeRange}
+              />
+            }
           />
         </Panel>
 

src/pages/ProfilesExplorerView/components/SceneExploreServiceFlameGraph/components/SceneExportMenu.tsx

@@ -0,0 +1,194 @@
+import { TimeRange } from '@grafana/data';
+import { reportInteraction } from '@grafana/runtime';
+import { SceneComponentProps, sceneGraph, SceneObjectBase, SceneObjectState } from '@grafana/scenes';
+import { Button, Dropdown, Menu } from '@grafana/ui';
+import { displayError } from '@shared/domain/displayStatus';
+import { useMaxNodesFromUrl } from '@shared/domain/url-params/useMaxNodesFromUrl';
+import { DEFAULT_SETTINGS } from '@shared/infrastructure/settings/PluginSettings';
+import { useFetchPluginSettings } from '@shared/infrastructure/settings/useFetchPluginSettings';
+import { DomainHookReturnValue } from '@shared/types/DomainHookReturnValue';
+import { FlamebearerProfile } from '@shared/types/FlamebearerProfile';
+import 'compression-streams-polyfill';
+import saveAs from 'file-saver';
+import React from 'react';
+
+import { ProfilesDataSourceVariable } from '../../../domain/variables/ProfilesDataSourceVariable';
+import { ProfileApiClient } from '../../../infrastructure/profiles/ProfileApiClient';
+import { DataSourceProxyClientBuilder } from '../../../infrastructure/series/http/DataSourceProxyClientBuilder';
+import { getExportFilename } from './domain/getExportFilename';
+import { flamegraphDotComApiClient } from './infrastructure/flamegraphDotComApiClient';
+import { PprofApiClient } from './infrastructure/PprofApiClient';
+
+interface SceneExportMenuState extends SceneObjectState {}
+
+type ExtraProps = {
+  query: string;
+  timeRange: TimeRange;
+};
+
+export class SceneExportMenu extends SceneObjectBase<SceneExportMenuState> {
+  constructor() {
+    super({ key: 'export-flame-graph-menu' });
+  }
+
+  async fetchFlamebearerProfile({
+    dataSourceUid,
+    query,
+    timeRange,
+    maxNodes,
+  }: ExtraProps & { dataSourceUid: string; maxNodes: number | null }): Promise<FlamebearerProfile | null> {
+    const profileApiClient = DataSourceProxyClientBuilder.build(dataSourceUid, ProfileApiClient) as ProfileApiClient;
+
+    let profile;
+
+    try {
+      profile = await profileApiClient.get({
+        query,
+        timeRange,
+        format: 'json',
+        maxNodes: maxNodes || DEFAULT_SETTINGS.maxNodes,
+      });
+    } catch (error) {
+      displayError(error, ['Error while loading flamebearer profile data!', (error as Error).message]);
+      return null;
+    }
+
+    return profile as FlamebearerProfile;
+  }
+
+  async fetchPprofProfile({
+    dataSourceUid,
+    query,
+    timeRange,
+    maxNodes,
+  }: ExtraProps & { dataSourceUid: string; maxNodes: number | null }): Promise<Blob | null> {
+    const pprofApiClient = DataSourceProxyClientBuilder.build(dataSourceUid, PprofApiClient) as PprofApiClient;
+
+    let profile;
+
+    try {
+      const blob = await pprofApiClient.selectMergeProfile({
+        query,
+        timeRange,
+        maxNodes: maxNodes || DEFAULT_SETTINGS.maxNodes,
+      });
+      profile = await new Response(blob.stream().pipeThrough(new CompressionStream('gzip'))).blob();
+    } catch (error) {
+      displayError(error, ['Failed to export to pprof!', (error as Error).message]);
+      return null;
+    }
+
+    return profile;
+  }
+
+  useSceneExportMenu = ({ query, timeRange }: ExtraProps): DomainHookReturnValue => {
+    const dataSourceUid = sceneGraph.findByKeyAndType(this, 'dataSource', ProfilesDataSourceVariable).useState()
+      .value as string;
+
+    const [maxNodes] = useMaxNodesFromUrl();
+    const { settings } = useFetchPluginSettings();
+
+    const downloadPng = () => {
+      reportInteraction('g_pyroscope_export_profile', { format: 'png' });
+
+      const filename = `${getExportFilename(query, timeRange)}.png`;
+
+      (document.querySelector('canvas[data-testid="flameGraph"]') as HTMLCanvasElement).toBlob((blob) => {
+        if (!blob) {
+          const error = new Error('Error while creating the image, no blob.');
+          displayError(error, ['Failed to export to png!', error.message]);
+          return;
+        }
+
+        saveAs(blob, filename);
+      }, 'image/png');
+    };
+
+    const downloadJson = async () => {
+      reportInteraction('g_pyroscope_export_profile', { format: 'json' });
+
+      const profile = await this.fetchFlamebearerProfile({ dataSourceUid, query, timeRange, maxNodes });
+      if (!profile) {
+        return;
+      }
+
+      const filename = `${getExportFilename(query, timeRange)}.json`;
+      const data = `data:text/json;charset=utf-8,${encodeURIComponent(JSON.stringify(profile))}`;
+
+      saveAs(data, filename);
+    };
+
+    const downloadPprof = async () => {
+      reportInteraction('g_pyroscope_export_profile', { format: 'pprof' });
+
+      const profile = await this.fetchPprofProfile({ dataSourceUid, query, timeRange, maxNodes });
+      if (!profile) {
+        return;
+      }
+
+      const filename = `${getExportFilename(query, timeRange)}.pb.gz`;
+
+      saveAs(profile, filename);
+    };
+
+    const uploadToFlamegraphDotCom = async () => {
+      reportInteraction('g_pyroscope_export_profile', { format: 'flamegraph.com' });
+
+      const profile = await this.fetchFlamebearerProfile({ dataSourceUid, query, timeRange, maxNodes });
+      if (!profile) {
+        return;
+      }
+
+      try {
+        const response = await flamegraphDotComApiClient.upload(getExportFilename(query, timeRange), profile);
+
+        if (!response.url) {
+          throw new Error('Empty URL received.');
+        }
+
+        const dlLink = document.createElement('a');
+        dlLink.target = '_blank';
+        dlLink.href = response.url;
+        document.body.appendChild(dlLink);
+        dlLink.click();
+        document.body.removeChild(dlLink);
+      } catch (error) {
+        displayError(error, ['Failed to export to flamegraph.com!', (error as Error).message]);
+        return;
+      }
+    };
+
+    return {
+      data: {
+        shouldDisplayFlamegraphDotCom: Boolean(settings?.enableFlameGraphDotComExport),
+      },
+      actions: {
+        downloadPng,
+        downloadJson,
+        downloadPprof,
+        uploadToFlamegraphDotCom,
+      },
+    };
+  };
+
+  static Component = ({ model, query, timeRange }: SceneComponentProps<SceneExportMenu> & ExtraProps) => {
+    const { data, actions } = model.useSceneExportMenu({ query, timeRange });
+
+    return (
+      <Dropdown
+        overlay={
+          <Menu>
+            <Menu.Item label="png" onClick={actions.downloadPng} />
+            <Menu.Item label="json" onClick={actions.downloadJson} />
+            <Menu.Item label="pprof" onClick={actions.downloadPprof} />
+            {data.shouldDisplayFlamegraphDotCom && (
+              <Menu.Item label="flamegraph.com" onClick={actions.uploadToFlamegraphDotCom} />
+            )}
+          </Menu>
+        }
+      >
+        <Button icon={'download-alt'} size={'sm'} variant={'secondary'} fill={'outline'} aria-label="Export data" />
+      </Dropdown>
+    );
+  };
+}

src/pages/ProfilesExplorerView/components/SceneExploreServiceFlameGraph/components/domain/__tests__/getExportFilename.spec.ts

@@ -0,0 +1,17 @@
+import { translatePyroscopeTimeRangeToGrafana } from '@shared/domain/translation';
+
+import { getExportFilename } from '../getExportFilename';
+
+describe('getFileName(timeRange, appName)', () => {
+  it('computes the correct filename', () => {
+    const timeRange = translatePyroscopeTimeRangeToGrafana('1708210800', '1708297200'); // 2024-02-18 - 2024-02-19
+    const filename = getExportFilename(
+      'process_cpu:cpu:nanoseconds:cpu:nanoseconds{service_name="alerting-ops/grafana",}',
+      timeRange
+    );
+
+    expect(filename).toBe(
+      'alerting-ops-grafana_process_cpu:cpu:nanoseconds:cpu:nanoseconds_2024-02-17_2300-to-2024-02-18_2300'
+    );
+  });
+});

src/pages/ProfilesExplorerView/components/SceneExploreServiceFlameGraph/components/domain/getExportFilename.ts

@@ -0,0 +1,8 @@
+import { TimeRange } from '@grafana/data';
+import { parseQuery } from '@shared/domain/url-params/parseQuery';
+
+export function getExportFilename(query: string, timeRange: TimeRange) {
+  const { serviceId, profileMetricId } = parseQuery(query);
+  const dateString = `${timeRange.from.format('YYYY-MM-DD_HHmm')}-to-${timeRange.to.format('YYYY-MM-DD_HHmm')}`;
+  return `${serviceId.replace(/\//g, '-')}_${profileMetricId}_${dateString}`;
+}

src/pages/ProfilesExplorerView/components/SceneExploreServiceFlameGraph/components/infrastructure/PprofApiClient.ts

@@ -0,0 +1,34 @@
+import { TimeRange } from '@grafana/data';
+import { parseQuery } from '@shared/domain/url-params/parseQuery';
+
+import { DataSourceProxyClient } from '../../../../infrastructure/series/http/DataSourceProxyClient';
+import { PprofRequest } from './PprofRequest';
+
+type SelectMergeProfileParams = {
+  query: string;
+  timeRange: TimeRange;
+  maxNodes: number;
+};
+
+export class PprofApiClient extends DataSourceProxyClient {
+  static buildPprofRequest(query: string, timeRange: TimeRange, maxNodes: number): Uint8Array {
+    const { profileMetricId, labelsSelector } = parseQuery(query);
+
+    const start = timeRange.from.unix() * 1000;
+    const end = timeRange.to.unix() * 1000;
+
+    const message = new PprofRequest(profileMetricId, labelsSelector, start, end, maxNodes);
+
+    return PprofRequest.encode(message).finish();
+  }
+
+  async selectMergeProfile({ query, timeRange, maxNodes }: SelectMergeProfileParams): Promise<Blob> {
+    const response = await this.fetch('/querier.v1.QuerierService/SelectMergeProfile', {
+      method: 'POST',
+      headers: { 'content-type': 'application/proto' },
+      body: new Blob([PprofApiClient.buildPprofRequest(query, timeRange, maxNodes)]),
+    });
+
+    return response.blob();
+  }
+}

src/shared/components/FlameGraph/components/infrastructure/PprofRequest.ts → src/pages/ProfilesExplorerView/components/SceneExploreServiceFlameGraph/components/infrastructure/PprofRequest.ts

@@ -5,13 +5,15 @@ export class PprofRequest extends Message<PprofRequest> {
         profile_typeID: string,
         label_selector: string,
         start: number,
-        end: number
+        end: number,
+        max_nodes: number
     ) {
         super();
         this.profile_typeID = profile_typeID;
         this.label_selector = label_selector;
         this.start = start;
         this.end = end;
+        this.max_nodes = max_nodes;
     }
 
     @Field.d(1, 'string')
@@ -25,4 +27,7 @@ export class PprofRequest extends Message<PprofRequest> {
 
     @Field.d(4, 'int64')
     end: number;
+
+    @Field.d(5, 'int64')
+    max_nodes: number;
 }