Pyroscope: Fix template variable support (#84477)

pkg/tsdb/grafana-pyroscope-datasource/pyroscopeClient.go

@@ -256,6 +256,10 @@ func (c *PyroscopeClient) LabelNames(ctx context.Context, labelSelector string,
 		return nil, fmt.Errorf("error sending LabelNames request %v", err)
 	}
 
+	if resp.Msg.Names == nil {
+		return []string{}, nil
+	}
+
 	var filtered []string
 	for _, label := range resp.Msg.Names {
 		if !isPrivateLabel(label) {
@@ -281,6 +285,9 @@ func (c *PyroscopeClient) LabelValues(ctx context.Context, label string, labelSe
 		span.SetStatus(codes.Error, err.Error())
 		return nil, err
 	}
+	if resp.Msg.Names == nil {
+		return []string{}, nil
+	}
 	return resp.Msg.Names, nil
 }
 

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

@@ -30,23 +30,16 @@ export function VariableQueryEditor(props: QueryEditorProps<PyroscopeDataSource,
             onChange={(value) => {
               if (value.value! === 'profileType') {
                 props.onChange({
-                  ...props.query,
                   type: value.value!,
+                  refId: props.query.refId,
                 });
               }
-              if (value.value! === 'label') {
+              if (value.value! === 'label' || value.value! === 'labelValue') {
                 props.onChange({
-                  ...props.query,
                   type: value.value!,
-                  profileTypeId: '',
-                });
-              }
-              if (value.value! === 'labelValue') {
-                props.onChange({
-                  ...props.query,
-                  type: value.value!,
-                  profileTypeId: '',
-                  labelName: '',
+                  refId: props.query.refId,
+                  // Make sure we keep already selected values if they make sense for the variable type
+                  profileTypeId: props.query.type !== 'profileType' ? props.query.profileTypeId : '',
                 });
               }
             }}
@@ -98,7 +91,13 @@ function LabelRow(props: {
   const [labels, setLabels] = useState<string[]>();
   useEffect(() => {
     (async () => {
-      setLabels(await props.datasource.getLabelNames((props.profileTypeId || '') + '{}', props.from, props.to));
+      setLabels(
+        await props.datasource.getLabelNames(
+          props.profileTypeId ? getProfileTypeLabel(props.profileTypeId) : '{}',
+          props.from,
+          props.to
+        )
+      );
     })();
   }, [props.datasource, props.profileTypeId, props.to, props.from]);
 
@@ -156,3 +155,7 @@ function ProfileTypeRow(props: {
     </InlineFieldRow>
   );
 }
+
+export function getProfileTypeLabel(type: string) {
+  return `{__profile_type__="${type}"}`;
+}

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

@@ -24,7 +24,11 @@ describe('VariableSupport', () => {
       vs.query(getDefaultRequest({ type: 'label', profileTypeId: 'profile:type:3', refId: 'A' }))
     );
     expect(resp.data).toEqual([{ text: 'foo' }, { text: 'bar' }, { text: 'baz' }]);
-    expect(mock.getLabelNames).toBeCalledWith('profile:type:3{}', expect.any(Number), expect.any(Number));
+    expect(mock.getLabelNames).toBeCalledWith(
+      '{__profile_type__="profile:type:3"}',
+      expect.any(Number),
+      expect.any(Number)
+    );
   });
 
   it('should query label values', async function () {
@@ -34,7 +38,12 @@ describe('VariableSupport', () => {
       vs.query(getDefaultRequest({ type: 'labelValue', labelName: 'foo', profileTypeId: 'profile:type:3', refId: 'A' }))
     );
     expect(resp.data).toEqual([{ text: 'val1' }, { text: 'val2' }, { text: 'val3' }]);
-    expect(mock.getLabelValues).toBeCalledWith('profile:type:3{}', 'foo', expect.any(Number), expect.any(Number));
+    expect(mock.getLabelValues).toBeCalledWith(
+      '{__profile_type__="profile:type:3"}',
+      'foo',
+      expect.any(Number),
+      expect.any(Number)
+    );
   });
 });
 

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

@@ -2,7 +2,7 @@ import { from, map, Observable, of } from 'rxjs';
 
 import { CustomVariableSupport, DataQueryRequest, DataQueryResponse, MetricFindValue } from '@grafana/data';
 
-import { VariableQueryEditor } from './VariableQueryEditor';
+import { getProfileTypeLabel, VariableQueryEditor } from './VariableQueryEditor';
 import { PyroscopeDataSource } from './datasource';
 import { ProfileTypeMessage, VariableQuery } from './types';
 
@@ -34,7 +34,7 @@ export class VariableSupport extends CustomVariableSupport<PyroscopeDataSource>
       }
       return from(
         this.dataAPI.getLabelNames(
-          request.targets[0].profileTypeId + '{}',
+          getProfileTypeLabel(request.targets[0].profileTypeId),
           request.range.from.valueOf(),
           request.range.to.valueOf()
         )
@@ -51,7 +51,7 @@ export class VariableSupport extends CustomVariableSupport<PyroscopeDataSource>
       }
       return from(
         this.dataAPI.getLabelValues(
-          request.targets[0].profileTypeId + '{}',
+          getProfileTypeLabel(request.targets[0].profileTypeId),
           request.targets[0].labelName,
           request.range.from.valueOf(),
           request.range.to.valueOf()