Pyroscope: Add adhoc filters support (#85601)

* Add adhoc filters support

* Add tests

* refactor tests

* Add comment

* Removed empty param docs

packages/grafana-data/src/types/time.ts

@@ -86,3 +86,20 @@ export function getDefaultRelativeTimeRange(): RelativeTimeRange {
     to: 0,
   };
 }
+
+/**
+ * Simple helper to quickly create a TimeRange object either from string representations of a dateTime or directly
+ * DateTime objects.
+ */
+export function makeTimeRange(from: DateTime | string, to: DateTime | string): TimeRange {
+  const fromDateTime = typeof from === 'string' ? dateTime(from) : from;
+  const toDateTime = typeof to === 'string' ? dateTime(to) : to;
+  return {
+    from: fromDateTime,
+    to: toDateTime,
+    raw: {
+      from: fromDateTime,
+      to: toDateTime,
+    },
+  };
+}

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

@@ -5,27 +5,14 @@ import {
   PluginMetaInfo,
   PluginType,
   DataSourceJsonData,
+  makeTimeRange,
 } from '@grafana/data';
-import { setPluginExtensionsHook, getBackendSrv, setBackendSrv, getTemplateSrv } from '@grafana/runtime';
+import { setPluginExtensionsHook, getBackendSrv, setBackendSrv, TemplateSrv } from '@grafana/runtime';
 
 import { defaultPyroscopeQueryType } from './dataquery.gen';
 import { normalizeQuery, PyroscopeDataSource } from './datasource';
 import { Query } from './types';
 
-jest.mock('@grafana/runtime', () => {
-  const actual = jest.requireActual('@grafana/runtime');
-  return {
-    ...actual,
-    getTemplateSrv: () => {
-      return {
-        replace: (query: string): string => {
-          return query.replace(/\$var/g, 'interpolated');
-        },
-      };
-    },
-  };
-});
-
 /** The datasource QueryEditor fetches datasource settings to send to the extension's `configure` method */
 export function mockFetchPyroscopeDatasourceSettings(
   datasourceSettings?: Partial<DataSourceInstanceSettings<DataSourceJsonData>>
@@ -46,16 +33,21 @@ export function mockFetchPyroscopeDatasourceSettings(
   });
 }
 
-describe('Pyroscope data source', () => {
-  let ds: PyroscopeDataSource;
-  beforeEach(() => {
-    mockFetchPyroscopeDatasourceSettings();
-    setPluginExtensionsHook(() => ({ extensions: [], isLoading: false })); // No extensions
-    ds = new PyroscopeDataSource(defaultSettings);
-  });
+function setupDatasource() {
+  mockFetchPyroscopeDatasourceSettings();
+  setPluginExtensionsHook(() => ({ extensions: [], isLoading: false })); // No extensions
+  const templateSrv = {
+    replace: (query: string): string => {
+      return query.replace(/\$var/g, 'interpolated');
+    },
+  } as unknown as TemplateSrv;
+  return new PyroscopeDataSource(defaultSettings, templateSrv);
+}
 
+describe('Pyroscope data source', () => {
   describe('importing queries', () => {
     it('keeps all labels and values', async () => {
+      const ds = setupDatasource();
       const queries = await ds.importFromAbstractQueries([
         {
           refId: 'A',
@@ -71,6 +63,7 @@ describe('Pyroscope data source', () => {
 
   describe('exporting queries', () => {
     it('keeps all labels and values', async () => {
+      const ds = setupDatasource();
       const queries = await ds.exportToAbstractQueries([
         {
           refId: 'A',
@@ -93,10 +86,8 @@ describe('Pyroscope data source', () => {
   });
 
   describe('applyTemplateVariables', () => {
-    const templateSrv = getTemplateSrv();
-
     it('should not update labelSelector if there are no template variables', () => {
-      ds = new PyroscopeDataSource(defaultSettings, templateSrv);
+      const ds = setupDatasource();
       const query = ds.applyTemplateVariables(defaultQuery({ labelSelector: `no var`, profileTypeId: 'no var' }), {});
       expect(query).toMatchObject({
         labelSelector: `no var`,
@@ -105,7 +96,7 @@ describe('Pyroscope data source', () => {
     });
 
     it('should update labelSelector if there are template variables', () => {
-      ds = new PyroscopeDataSource(defaultSettings, templateSrv);
+      const ds = setupDatasource();
       const query = ds.applyTemplateVariables(
         defaultQuery({ labelSelector: `{$var="$var"}`, profileTypeId: '$var' }),
         {}
@@ -113,6 +104,30 @@ describe('Pyroscope data source', () => {
       expect(query).toMatchObject({ labelSelector: `{interpolated="interpolated"}`, profileTypeId: 'interpolated' });
     });
   });
+
+  it('implements ad hoc variable support for keys', async () => {
+    const ds = setupDatasource();
+    jest.spyOn(ds, 'getResource').mockImplementationOnce(async (cb) => ['foo', 'bar', 'baz']);
+    const keys = await ds.getTagKeys({
+      filters: [],
+      timeRange: makeTimeRange('2024-01-01T00:00:00', '2024-01-01T01:00:00'),
+    });
+    expect(keys).toEqual(['foo', 'bar', 'baz'].map((v) => ({ text: v })));
+  });
+
+  it('implements ad hoc variable support for values', async () => {
+    const ds = setupDatasource();
+    jest.spyOn(ds, 'getResource').mockImplementationOnce(async (path, params) => {
+      expect(params?.label).toEqual('foo');
+      return ['xyz', 'tuv'];
+    });
+    const keys = await ds.getTagValues({
+      key: 'foo',
+      filters: [],
+      timeRange: makeTimeRange('2024-01-01T00:00:00', '2024-01-01T01:00:00'),
+    });
+    expect(keys).toEqual(['xyz', 'tuv'].map((v) => ({ text: v })));
+  });
 });
 
 describe('normalizeQuery', () => {

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

@@ -1,12 +1,16 @@
-import Prism, { Grammar } from 'prismjs';
+import Prism from 'prismjs';
 import { Observable, of } from 'rxjs';
 
 import {
   AbstractQuery,
+  AdHocVariableFilter,
   CoreApp,
   DataQueryRequest,
   DataQueryResponse,
+  DataSourceGetTagKeysOptions,
+  DataSourceGetTagValuesOptions,
   DataSourceInstanceSettings,
+  MetricFindValue,
   ScopedVars,
 } from '@grafana/data';
 import { DataSourceWithBackend, getTemplateSrv, TemplateSrv } from '@grafana/runtime';
@@ -14,7 +18,7 @@ import { DataSourceWithBackend, getTemplateSrv, TemplateSrv } from '@grafana/run
 import { VariableSupport } from './VariableSupport';
 import { defaultGrafanaPyroscopeDataQuery, defaultPyroscopeQueryType } from './dataquery.gen';
 import { PyroscopeDataSourceOptions, Query, ProfileTypeMessage } from './types';
-import { extractLabelMatchers, toPromLikeExpr } from './utils';
+import { addLabelToQuery, extractLabelMatchers, grammar, toPromLikeExpr } from './utils';
 
 export class PyroscopeDataSource extends DataSourceWithBackend<Query, PyroscopeDataSourceOptions> {
   constructor(
@@ -71,10 +75,37 @@ export class PyroscopeDataSource extends DataSourceWithBackend<Query, PyroscopeD
     });
   }
 
-  applyTemplateVariables(query: Query, scopedVars: ScopedVars): Query {
+  // By implementing getTagKeys and getTagValues we add ad-hoc filters functionality
+  async getTagKeys(options: DataSourceGetTagKeysOptions<Query>): Promise<MetricFindValue[]> {
+    const data = this.adhocFilterData(options);
+    const labels = await this.getLabelNames(data.query, data.from, data.to);
+    return labels.map((label) => ({ text: label }));
+  }
+
+  // By implementing getTagKeys and getTagValues we add ad-hoc filters functionality
+  async getTagValues(options: DataSourceGetTagValuesOptions<Query>): Promise<MetricFindValue[]> {
+    const data = this.adhocFilterData(options);
+    const labels = await this.getLabelValues(data.query, options.key, data.from, data.to);
+    return labels.map((label) => ({ text: label }));
+  }
+
+  private adhocFilterData(options: DataSourceGetTagKeysOptions<Query> | DataSourceGetTagValuesOptions<Query>) {
+    const from = options.timeRange?.from.valueOf() ?? Date.now() - 1000 * 60 * 60 * 24;
+    const to = options.timeRange?.to.valueOf() ?? Date.now();
+    const query = '{' + options.filters.map((f) => `${f.key}${f.operator}"${f.value}"`).join(',') + '}';
+    return { from, to, query };
+  }
+
+  applyTemplateVariables(query: Query, scopedVars: ScopedVars, filters?: AdHocVariableFilter[]): Query {
+    let labelSelector = this.templateSrv.replace(query.labelSelector ?? '', scopedVars);
+    if (filters && labelSelector) {
+      for (const filter of filters) {
+        labelSelector = addLabelToQuery(labelSelector, filter.key, filter.value, filter.operator);
+      }
+    }
     return {
       ...query,
-      labelSelector: this.templateSrv.replace(query.labelSelector ?? '', scopedVars),
+      labelSelector,
       profileTypeId: this.templateSrv.replace(query.profileTypeId ?? '', scopedVars),
     };
   }
@@ -86,7 +117,7 @@ export class PyroscopeDataSource extends DataSourceWithBackend<Query, PyroscopeD
   importFromAbstractQuery(labelBasedQuery: AbstractQuery): Query {
     return {
       refId: labelBasedQuery.refId,
-      labelSelector: toPromLikeExpr(labelBasedQuery),
+      labelSelector: toPromLikeExpr(labelBasedQuery.labelMatchers),
       queryType: 'both',
       profileTypeId: '',
       groupBy: [],
@@ -128,27 +159,3 @@ export function normalizeQuery(query: Query, app?: CoreApp | string) {
   }
   return normalized;
 }
-
-const grammar: Grammar = {
-  'context-labels': {
-    pattern: /\{[^}]*(?=}?)/,
-    greedy: true,
-    inside: {
-      comment: {
-        pattern: /#.*/,
-      },
-      'label-key': {
-        pattern: /[a-zA-Z_]\w*(?=\s*(=|!=|=~|!~))/,
-        alias: 'attr-name',
-        greedy: true,
-      },
-      'label-value': {
-        pattern: /"(?:\\.|[^\\"])*"/,
-        greedy: true,
-        alias: 'attr-value',
-      },
-      punctuation: /[{]/,
-    },
-  },
-  punctuation: /[{}(),.]/,
-};

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

@@ -1,7 +1,7 @@
 import { invert } from 'lodash';
-import { Token } from 'prismjs';
+import Prism, { Grammar, Token } from 'prismjs';
 
-import { AbstractLabelMatcher, AbstractLabelOperator, AbstractQuery } from '@grafana/data';
+import { AbstractLabelMatcher, AbstractLabelOperator } from '@grafana/data';
 
 export function extractLabelMatchers(tokens: Array<string | Token>): AbstractLabelMatcher[] {
   const labelMatchers: AbstractLabelMatcher[] = [];
@@ -47,8 +47,8 @@ export function extractLabelMatchers(tokens: Array<string | Token>): AbstractLab
   return labelMatchers;
 }
 
-export function toPromLikeExpr(labelBasedQuery: AbstractQuery): string {
-  const expr = labelBasedQuery.labelMatchers
+export function toPromLikeExpr(labelMatchers: AbstractLabelMatcher[]): string {
+  const expr = labelMatchers
     .map((selector: AbstractLabelMatcher) => {
       const operator = ToPromLikeMap[selector.operator];
       if (operator) {
@@ -82,3 +82,52 @@ const ToPromLikeMap: Record<AbstractLabelOperator, string> = invert(FromPromLike
   AbstractLabelOperator,
   string
 >;
+
+/**
+ * Modifies query, adding a new label=value pair to it while preserving other parts of the query. This operates on a
+ * string representation of the query which needs to be parsed and then rendered to string again.
+ */
+export function addLabelToQuery(query: string, key: string, value: string | number, operator = '='): string {
+  if (!key || !value) {
+    throw new Error('Need label to add to query.');
+  }
+
+  const tokens = Prism.tokenize(query, grammar);
+  let labels = extractLabelMatchers(tokens);
+
+  // If we already have such label in the query, remove it and we will replace it. If we didn't we would end up
+  // with query like `a=b,a=c` which won't return anything. Replacing also seems more meaningful here than just
+  // ignoring the filter and keeping the old value.
+  labels = labels.filter((l) => l.name !== key);
+  labels.push({
+    name: key,
+    value: value.toString(),
+    operator: FromPromLikeMap[operator] ?? AbstractLabelOperator.Equal,
+  });
+
+  return toPromLikeExpr(labels);
+}
+
+export const grammar: Grammar = {
+  'context-labels': {
+    pattern: /\{[^}]*(?=}?)/,
+    greedy: true,
+    inside: {
+      comment: {
+        pattern: /#.*/,
+      },
+      'label-key': {
+        pattern: /[a-zA-Z_]\w*(?=\s*(=|!=|=~|!~))/,
+        alias: 'attr-name',
+        greedy: true,
+      },
+      'label-value': {
+        pattern: /"(?:\\.|[^\\"])*"/,
+        greedy: true,
+        alias: 'attr-value',
+      },
+      punctuation: /[{]/,
+    },
+  },
+  punctuation: /[{}(),.]/,
+};