Tempo: Add filtering for service graph query (#41162)

* Add filter based on AdHocFilter element

* Add tests

* Cancel layout in case we have have new data or we unmount node graph

* Fix typing

* Fix test

public/app/plugins/datasource/prometheus/datasource.ts

@@ -300,7 +300,7 @@ export class PrometheusDatasource extends DataSourceWithBackend<PromQuery, PromO
   };
 
   shouldRunExemplarQuery(target: PromQuery): boolean {
-    /* We want to run exemplar query only for histogram metrics: 
+    /* We want to run exemplar query only for histogram metrics:
     1. If we haven't processd histogram metrics yet, we need to check if expr includes "_bucket" which means that it is probably histogram metric (can rarely lead to false positive).
     2. If we have processed histogram metrics, check if it is part of query expr.
     */
@@ -790,7 +790,7 @@ export class PrometheusDatasource extends DataSourceWithBackend<PromQuery, PromO
     return result?.data?.data?.map((value: any) => ({ text: value })) ?? [];
   }
 
-  async getTagValues(options: any = {}) {
+  async getTagValues(options: { key?: string } = {}) {
     const result = await this.metadataRequest(`/api/v1/label/${options.key}/values`);
     return result?.data?.data?.map((value: any) => ({ text: value })) ?? [];
   }

public/app/plugins/datasource/tempo/QueryEditor/NativeSearch.tsx

@@ -12,13 +12,13 @@ import {
   Alert,
   useStyles2,
 } from '@grafana/ui';
-import { tokenizer } from './syntax';
+import { tokenizer } from '../syntax';
 import Prism from 'prismjs';
 import { Node } from 'slate';
 import { css } from '@emotion/css';
 import { GrafanaTheme2, isValidGoDuration, SelectableValue } from '@grafana/data';
-import TempoLanguageProvider from './language_provider';
-import { TempoDatasource, TempoQuery } from './datasource';
+import TempoLanguageProvider from '../language_provider';
+import { TempoDatasource, TempoQuery } from '../datasource';
 import { debounce } from 'lodash';
 import { dispatch } from 'app/store/store';
 import { notifyApp } from 'app/core/actions';

public/app/plugins/datasource/tempo/QueryEditor/QueryField.tsx

@@ -1,6 +1,6 @@
 import { css } from '@emotion/css';
-import { DataSourceApi, QueryEditorProps, SelectableValue } from '@grafana/data';
-import { config, getDataSourceSrv } from '@grafana/runtime';
+import { QueryEditorProps, SelectableValue } from '@grafana/data';
+import { config } from '@grafana/runtime';
 import {
   Badge,
   FileDropzone,
@@ -14,13 +14,15 @@ import {
 } from '@grafana/ui';
 import { TraceToLogsOptions } from 'app/core/components/TraceToLogsSettings';
 import React from 'react';
-import { LokiQueryField } from '../loki/components/LokiQueryField';
-import { LokiQuery } from '../loki/types';
-import { TempoDatasource, TempoQuery, TempoQueryType } from './datasource';
-import LokiDatasource from '../loki/datasource';
-import { PrometheusDatasource } from '../prometheus/datasource';
+import { LokiQueryField } from '../../loki/components/LokiQueryField';
+import { LokiQuery } from '../../loki/types';
+import { TempoDatasource, TempoQuery, TempoQueryType } from '../datasource';
+import LokiDatasource from '../../loki/datasource';
+import { PrometheusDatasource } from '../../prometheus/datasource';
 import useAsync from 'react-use/lib/useAsync';
 import NativeSearch from './NativeSearch';
+import { getDS } from './utils';
+import { ServiceGraphSection } from './ServiceGraphSection';
 
 interface Props extends QueryEditorProps<TempoDatasource, TempoQuery>, Themeable2 {}
 
@@ -188,36 +190,14 @@ class TempoQueryFieldComponent extends React.PureComponent<Props, State> {
             </InlineField>
           </InlineFieldRow>
         )}
-        {query.queryType === 'serviceMap' && <ServiceGraphSection graphDatasourceUid={graphDatasourceUid} />}
+        {query.queryType === 'serviceMap' && (
+          <ServiceGraphSection graphDatasourceUid={graphDatasourceUid} query={query} onChange={onChange} />
+        )}
       </>
     );
   }
 }
 
-function ServiceGraphSection({ graphDatasourceUid }: { graphDatasourceUid?: string }) {
-  const dsState = useAsync(() => getDS(graphDatasourceUid), [graphDatasourceUid]);
-  if (dsState.loading) {
-    return null;
-  }
-
-  const ds = dsState.value as LokiDatasource;
-
-  if (!graphDatasourceUid) {
-    return <div className="text-warning">Please set up a service graph datasource in the datasource settings.</div>;
-  }
-
-  if (graphDatasourceUid && !ds) {
-    return (
-      <div className="text-warning">
-        Service graph datasource is configured but the data source no longer exists. Please configure existing data
-        source to use the service graph functionality.
-      </div>
-    );
-  }
-
-  return null;
-}
-
 interface SearchSectionProps {
   linkedDatasourceUid?: string;
   onChange: (value: LokiQuery) => void;
@@ -264,18 +244,4 @@ function SearchSection({ linkedDatasourceUid, onChange, onRunQuery, query }: Sea
   return null;
 }
 
-async function getDS(uid?: string): Promise<DataSourceApi | undefined> {
-  if (!uid) {
-    return undefined;
-  }
-
-  const dsSrv = getDataSourceSrv();
-  try {
-    return await dsSrv.get(uid);
-  } catch (error) {
-    console.error('Failed to load data source', error);
-    return undefined;
-  }
-}
-
 export const TempoQueryField = withTheme2(TempoQueryFieldComponent);

public/app/plugins/datasource/tempo/QueryEditor/ServiceGraphSection.tsx

@@ -0,0 +1,87 @@
+import React from 'react';
+import useAsync from 'react-use/lib/useAsync';
+import { getDS } from './utils';
+import { InlineField, InlineFieldRow } from '@grafana/ui';
+import { AdHocVariableFilter } from '../../../../features/variables/types';
+import { TempoQuery } from '../datasource';
+import { AdHocFilter } from '../../../../features/variables/adhoc/picker/AdHocFilter';
+import { PrometheusDatasource } from '../../prometheus/datasource';
+
+export function ServiceGraphSection({
+  graphDatasourceUid,
+  query,
+  onChange,
+}: {
+  graphDatasourceUid?: string;
+  query: TempoQuery;
+  onChange: (value: TempoQuery) => void;
+}) {
+  const dsState = useAsync(() => getDS(graphDatasourceUid), [graphDatasourceUid]);
+  if (dsState.loading) {
+    return null;
+  }
+
+  const ds = dsState.value as PrometheusDatasource;
+
+  if (!graphDatasourceUid) {
+    return <div className="text-warning">Please set up a service graph datasource in the datasource settings.</div>;
+  }
+
+  if (graphDatasourceUid && !ds) {
+    return (
+      <div className="text-warning">
+        Service graph datasource is configured but the data source no longer exists. Please configure existing data
+        source to use the service graph functionality.
+      </div>
+    );
+  }
+  const filters = queryToFilter(query.serviceMapQuery || '');
+
+  return (
+    <div>
+      <InlineFieldRow>
+        <InlineField label="Filter" labelWidth={14} grow>
+          <AdHocFilter
+            datasource={{ uid: graphDatasourceUid }}
+            filters={filters}
+            addFilter={(filter: AdHocVariableFilter) => {
+              onChange({
+                ...query,
+                serviceMapQuery: filtersToQuery([...filters, filter]),
+              });
+            }}
+            removeFilter={(index: number) => {
+              const newFilters = [...filters];
+              newFilters.splice(index, 1);
+              onChange({ ...query, serviceMapQuery: filtersToQuery(newFilters) });
+            }}
+            changeFilter={(index: number, filter: AdHocVariableFilter) => {
+              const newFilters = [...filters];
+              newFilters.splice(index, 1, filter);
+              onChange({ ...query, serviceMapQuery: filtersToQuery(newFilters) });
+            }}
+          />
+        </InlineField>
+      </InlineFieldRow>
+    </div>
+  );
+}
+
+function queryToFilter(query: string): AdHocVariableFilter[] {
+  let match;
+  let filters: AdHocVariableFilter[] = [];
+  const re = /([\w_]+)(=|!=|<|>|=~|!~)"(.*?)"/g;
+  while ((match = re.exec(query)) !== null) {
+    filters.push({
+      key: match[1],
+      operator: match[2],
+      value: match[3],
+      condition: '',
+    });
+  }
+  return filters;
+}
+
+function filtersToQuery(filters: AdHocVariableFilter[]): string {
+  return `{${filters.map((f) => `${f.key}${f.operator}"${f.value}"`).join(',')}}`;
+}

public/app/plugins/datasource/tempo/QueryEditor/utils.ts

@@ -0,0 +1,16 @@
+import { DataSourceApi } from '@grafana/data';
+import { getDataSourceSrv } from '@grafana/runtime';
+
+export async function getDS(uid?: string): Promise<DataSourceApi | undefined> {
+  if (!uid) {
+    return undefined;
+  }
+
+  const dsSrv = getDataSourceSrv();
+  try {
+    return await dsSrv.get(uid);
+  } catch (error) {
+    console.error('Failed to load data source', error);
+    return undefined;
+  }
+}

public/app/plugins/datasource/tempo/datasource.ts

@@ -54,6 +54,7 @@ export type TempoQuery = {
   minDuration?: string;
   maxDuration?: string;
   limit?: number;
+  serviceMapQuery?: string;
 } & DataQuery;
 
 export const DEFAULT_LIMIT = 20;
@@ -268,6 +269,16 @@ export class TempoDatasource extends DataSourceWithBackend<TempoQuery, TempoJson
     const tagsQueryObject = tagsQuery.reduce((tagQuery, item) => ({ ...tagQuery, ...item }), {});
     return { ...tagsQueryObject, ...tempoQuery };
   }
+
+  async getServiceGraphLabels() {
+    const ds = await getDatasourceSrv().get(this.serviceMap!.datasourceUid);
+    return ds.getTagKeys!();
+  }
+
+  async getServiceGraphLabelValues(key: string) {
+    const ds = await getDatasourceSrv().get(this.serviceMap!.datasourceUid);
+    return ds.getTagValues!({ key });
+  }
 }
 
 function queryServiceMapPrometheus(request: DataQueryRequest<PromQuery>, datasourceUid: string) {
@@ -324,7 +335,9 @@ function makePromServiceMapRequest(options: DataQueryRequest<TempoQuery>): DataQ
     targets: serviceMapMetrics.map((metric) => {
       return {
         refId: metric,
-        expr: `delta(${metric}[$__range])`,
+        // options.targets[0] is not correct here, but not sure what should happen if you have multiple queries for
+        // service map at the same time anyway
+        expr: `delta(${metric}${options.targets[0].serviceMapQuery || ''}[$__range])`,
         instant: true,
       };
     }),

public/app/plugins/datasource/tempo/module.ts

@@ -2,7 +2,7 @@ import { DataSourcePlugin } from '@grafana/data';
 import CheatSheet from './CheatSheet';
 import { ConfigEditor } from './configuration/ConfigEditor';
 import { TempoDatasource } from './datasource';
-import { TempoQueryField } from './QueryField';
+import { TempoQueryField } from './QueryEditor/QueryField';
 
 export const plugin = new DataSourcePlugin(TempoDatasource)
   .setConfigEditor(ConfigEditor)