From ee38bbe03049279ad52679d760eb4dac4db7fd30 Mon Sep 17 00:00:00 2001
From: Joey <90795735+joey-grafana@users.noreply.github.com>
Date: Wed, 8 Mar 2023 13:29:28 +0000
Subject: [PATCH] Phlare: Allow variables in labelSelector (in query) (#64324)

applyTemplateVariables and tests
---
 .../datasource/phlare/datasource.test.ts      | 42 +++++++++++++++++++
 .../plugins/datasource/phlare/datasource.ts   | 22 ++++++++--
 2 files changed, 61 insertions(+), 3 deletions(-)

diff --git a/public/app/plugins/datasource/phlare/datasource.test.ts b/public/app/plugins/datasource/phlare/datasource.test.ts
index 10f6d6d1469..f11c4ebc9bb 100644
--- a/public/app/plugins/datasource/phlare/datasource.test.ts
+++ b/public/app/plugins/datasource/phlare/datasource.test.ts
@@ -1,5 +1,7 @@
 import { AbstractLabelOperator, DataSourceInstanceSettings, PluginMetaInfo, PluginType } from '@grafana/data';
+import { TemplateSrv } from 'app/features/templating/template_srv';
 
+import { defaultPhlareQueryType } from './dataquery.gen';
 import { PhlareDataSource } from './datasource';
 
 describe('Phlare data source', () => {
@@ -45,8 +47,48 @@ describe('Phlare data source', () => {
       ]);
     });
   });
+
+  describe('applyTemplateVariables', () => {
+    const interpolationVar = '$interpolationVar';
+    const interpolationText = 'interpolationText';
+    const noInterpolation = 'noInterpolation';
+
+    it('should not update labelSelector if there are no template variables', () => {
+      const templateSrv = new TemplateSrv();
+      templateSrv.replace = jest.fn((query: string): string => {
+        return query.replace(/\$interpolationVar/g, interpolationText);
+      });
+      ds = new PhlareDataSource(defaultSettings, templateSrv);
+      const query = ds.applyTemplateVariables(defaultQuery(`{${noInterpolation}}`), {});
+      expect(templateSrv.replace).toBeCalledTimes(1);
+      expect(query.labelSelector).toBe(`{${noInterpolation}}`);
+    });
+
+    it('should update labelSelector if there are template variables', () => {
+      const templateSrv = new TemplateSrv();
+      templateSrv.replace = jest.fn((query: string): string => {
+        return query.replace(/\$interpolationVar/g, interpolationText);
+      });
+      ds = new PhlareDataSource(defaultSettings, templateSrv);
+      const query = ds.applyTemplateVariables(defaultQuery(`{${interpolationVar}="${interpolationVar}"}`), {
+        interpolationVar: { text: interpolationText, value: interpolationText },
+      });
+      expect(templateSrv.replace).toBeCalledTimes(1);
+      expect(query.labelSelector).toBe(`{${interpolationText}="${interpolationText}"}`);
+    });
+  });
 });
 
+const defaultQuery = (query: string) => {
+  return {
+    refId: 'x',
+    groupBy: [],
+    labelSelector: query,
+    profileTypeId: '',
+    queryType: defaultPhlareQueryType,
+  };
+};
+
 const defaultSettings: DataSourceInstanceSettings = {
   id: 0,
   uid: 'phlare',
diff --git a/public/app/plugins/datasource/phlare/datasource.ts b/public/app/plugins/datasource/phlare/datasource.ts
index 95da2c417de..cb5d7c93c7f 100644
--- a/public/app/plugins/datasource/phlare/datasource.ts
+++ b/public/app/plugins/datasource/phlare/datasource.ts
@@ -1,8 +1,14 @@
 import Prism, { Grammar } from 'prismjs';
 import { Observable, of } from 'rxjs';
 
-import { AbstractQuery, DataQueryRequest, DataQueryResponse, DataSourceInstanceSettings } from '@grafana/data';
-import { DataSourceWithBackend } from '@grafana/runtime';
+import {
+  AbstractQuery,
+  DataQueryRequest,
+  DataQueryResponse,
+  DataSourceInstanceSettings,
+  ScopedVars,
+} from '@grafana/data';
+import { DataSourceWithBackend, getTemplateSrv, TemplateSrv } from '@grafana/runtime';
 
 import { extractLabelMatchers, toPromLikeExpr } from '../prometheus/language_utils';
 
@@ -10,7 +16,10 @@ import { normalizeQuery } from './QueryEditor/QueryEditor';
 import { PhlareDataSourceOptions, Query, ProfileTypeMessage, SeriesMessage } from './types';
 
 export class PhlareDataSource extends DataSourceWithBackend<Query, PhlareDataSourceOptions> {
-  constructor(instanceSettings: DataSourceInstanceSettings<PhlareDataSourceOptions>) {
+  constructor(
+    instanceSettings: DataSourceInstanceSettings<PhlareDataSourceOptions>,
+    private readonly templateSrv: TemplateSrv = getTemplateSrv()
+  ) {
     super(instanceSettings);
   }
 
@@ -49,6 +58,13 @@ export class PhlareDataSource extends DataSourceWithBackend<Query, PhlareDataSou
     return await super.getResource('labelNames');
   }
 
+  applyTemplateVariables(query: Query, scopedVars: ScopedVars): Query {
+    return {
+      ...query,
+      labelSelector: this.templateSrv.replace(query.labelSelector ?? '', scopedVars),
+    };
+  }
+
   async importFromAbstractQueries(abstractQueries: AbstractQuery[]): Promise<Query[]> {
     return abstractQueries.map((abstractQuery) => this.importFromAbstractQuery(abstractQuery));
   }