From 58c5e5c9dc6ddf42141ce4b55ceb6631e7f42f36 Mon Sep 17 00:00:00 2001
From: Joey <90795735+joey-grafana@users.noreply.github.com>
Date: Fri, 8 Dec 2023 09:44:51 +0000
Subject: [PATCH] Tempo: Update the embedded flame graph labelSelector (#79061)

* Update the embedded flame graph labelSelector

* Update labelSelector
---
 .../SpanDetail/SpanFlameGraph.tsx             | 26 +++++++++++++------
 .../explore/TraceView/createSpanLink.tsx      |  4 +--
 2 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/public/app/features/explore/TraceView/components/TraceTimelineViewer/SpanDetail/SpanFlameGraph.tsx b/public/app/features/explore/TraceView/components/TraceTimelineViewer/SpanDetail/SpanFlameGraph.tsx
index f2bbf3b700e..dbe28555c26 100644
--- a/public/app/features/explore/TraceView/components/TraceTimelineViewer/SpanDetail/SpanFlameGraph.tsx
+++ b/public/app/features/explore/TraceView/components/TraceTimelineViewer/SpanDetail/SpanFlameGraph.tsx
@@ -21,7 +21,7 @@ import { PyroscopeQueryType } from 'app/plugins/datasource/grafana-pyroscope-dat
 import { PyroscopeDataSource } from 'app/plugins/datasource/grafana-pyroscope-datasource/datasource';
 import { Query } from 'app/plugins/datasource/grafana-pyroscope-datasource/types';
 
-import { pyroscopeProfileIdTagKey } from '../../../createSpanLink';
+import { defaultProfilingKeys, getFormattedTags, pyroscopeProfileIdTagKey } from '../../../createSpanLink';
 import { TraceSpan } from '../../types/trace';
 
 import { TraceFlameGraphs } from '.';
@@ -75,8 +75,20 @@ export default function SpanFlameGraph(props: SpanFlameGraphProps) {
   const queryFlameGraph = useCallback(
     async (
       profilesDataSourceSettings: DataSourceInstanceSettings<DataSourceJsonData>,
-      traceToProfilesOptions: TraceToProfilesOptions
+      traceToProfilesOptions: TraceToProfilesOptions,
+      span: TraceSpan
     ) => {
+      let labelSelector = '{}';
+      if (traceToProfilesOptions.customQuery && traceToProfilesOptions.query) {
+        labelSelector = traceToProfilesOptions.query;
+      } else {
+        const tags =
+          traceToProfilesOptions.tags && traceToProfilesOptions.tags.length > 0
+            ? traceToProfilesOptions.tags
+            : defaultProfilingKeys;
+        labelSelector = `{${getFormattedTags(span, tags)}}`;
+      }
+
       const request = {
         requestId: 'span-flamegraph-requestId',
         interval: '2s',
@@ -88,7 +100,7 @@ export default function SpanFlameGraph(props: SpanFlameGraphProps) {
         startTime: span.startTime,
         targets: [
           {
-            labelSelector: '{}',
+            labelSelector,
             groupBy: [],
             profileTypeId: traceToProfilesOptions.profileTypeId ?? '',
             queryType: 'profile' as PyroscopeQueryType,
@@ -107,7 +119,7 @@ export default function SpanFlameGraph(props: SpanFlameGraphProps) {
         setTraceFlameGraphs({ ...traceFlameGraphs, [profileTagValue]: flameGraph });
       }
     },
-    [getTimeRangeForProfile, profileTagValue, setTraceFlameGraphs, span.startTime, timeZone, traceFlameGraphs]
+    [getTimeRangeForProfile, profileTagValue, setTraceFlameGraphs, timeZone, traceFlameGraphs]
   );
 
   useEffect(() => {
@@ -117,18 +129,16 @@ export default function SpanFlameGraph(props: SpanFlameGraphProps) {
         profilesDataSourceSettings = getDatasourceSrv().getInstanceSettings(traceToProfilesOptions.datasourceUid);
       }
       if (traceToProfilesOptions && profilesDataSourceSettings) {
-        queryFlameGraph(profilesDataSourceSettings, traceToProfilesOptions);
+        queryFlameGraph(profilesDataSourceSettings, traceToProfilesOptions, span);
       }
     }
   }, [
     setTraceFlameGraphs,
-    span.tags,
+    span,
     traceFlameGraphs,
     traceToProfilesOptions,
     getTimeRangeForProfile,
-    span.startTime,
     timeZone,
-    span.spanID,
     queryFlameGraph,
     profileTagValue,
   ]);
diff --git a/public/app/features/explore/TraceView/createSpanLink.tsx b/public/app/features/explore/TraceView/createSpanLink.tsx
index 02e60b0d95c..80704a0fb7b 100644
--- a/public/app/features/explore/TraceView/createSpanLink.tsx
+++ b/public/app/features/explore/TraceView/createSpanLink.tsx
@@ -134,7 +134,7 @@ const formatDefaultKeys = (keys: string[]) => {
   }));
 };
 const defaultKeys = formatDefaultKeys(['cluster', 'hostname', 'namespace', 'pod', 'service.name', 'service.namespace']);
-const defaultProfilingKeys = formatDefaultKeys(['service.name', 'service.namespace']);
+export const defaultProfilingKeys = formatDefaultKeys(['service.name', 'service.namespace']);
 export const pyroscopeProfileIdTagKey = 'pyroscope.profile.id';
 
 function legacyCreateSpanLinkFactory(
@@ -499,7 +499,7 @@ function getQueryForFalconLogScale(span: TraceSpan, options: TraceToLogsOptionsV
  * Creates a string representing all the tags already formatted for use in the query. The tags are filtered so that
  * only intersection of tags that exist in a span and tags that you want are serialized into the string.
  */
-function getFormattedTags(
+export function getFormattedTags(
   span: TraceSpan,
   tags: TraceToLogsTag[],
   { labelValueSign = '=', joinBy = ', ' }: { labelValueSign?: string; joinBy?: string } = {}