From b8618723d0a2c3cc153d65125909c0a937c80e37 Mon Sep 17 00:00:00 2001
From: Joey Tawadrous <90795735+joey-grafana@users.noreply.github.com>
Date: Tue, 3 Jan 2023 10:25:16 +0000
Subject: [PATCH] Phlare: Transition from LogQL/PromQL to Phlare should keep
 the query (#60217)

* Transition from LogQL/PromQL to Phlare keeps the query

* Tests
---
 .../datasource/phlare/datasource.test.ts      | 66 +++++++++++++++++++
 .../plugins/datasource/phlare/datasource.ts   | 59 ++++++++++++++++-
 2 files changed, 124 insertions(+), 1 deletion(-)
 create mode 100644 public/app/plugins/datasource/phlare/datasource.test.ts

diff --git a/public/app/plugins/datasource/phlare/datasource.test.ts b/public/app/plugins/datasource/phlare/datasource.test.ts
new file mode 100644
index 00000000000..10f6d6d1469
--- /dev/null
+++ b/public/app/plugins/datasource/phlare/datasource.test.ts
@@ -0,0 +1,66 @@
+import { AbstractLabelOperator, DataSourceInstanceSettings, PluginMetaInfo, PluginType } from '@grafana/data';
+
+import { PhlareDataSource } from './datasource';
+
+describe('Phlare data source', () => {
+  let ds: PhlareDataSource;
+  beforeEach(() => {
+    ds = new PhlareDataSource(defaultSettings);
+  });
+
+  describe('importing queries', () => {
+    it('keeps all labels and values', async () => {
+      const queries = await ds.importFromAbstractQueries([
+        {
+          refId: 'A',
+          labelMatchers: [
+            { name: 'foo', operator: AbstractLabelOperator.Equal, value: 'bar' },
+            { name: 'foo2', operator: AbstractLabelOperator.Equal, value: 'bar2' },
+          ],
+        },
+      ]);
+      expect(queries[0].labelSelector).toBe('{foo="bar", foo2="bar2"}');
+    });
+  });
+
+  describe('exporting queries', () => {
+    it('keeps all labels and values', async () => {
+      const queries = await ds.exportToAbstractQueries([
+        {
+          refId: 'A',
+          labelSelector: '{foo="bar", foo2="bar2"}',
+          queryType: 'both',
+          profileTypeId: '',
+          groupBy: [''],
+        },
+      ]);
+      expect(queries).toMatchObject([
+        {
+          refId: 'A',
+          labelMatchers: [
+            { name: 'foo', operator: AbstractLabelOperator.Equal, value: 'bar' },
+            { name: 'foo2', operator: AbstractLabelOperator.Equal, value: 'bar2' },
+          ],
+        },
+      ]);
+    });
+  });
+});
+
+const defaultSettings: DataSourceInstanceSettings = {
+  id: 0,
+  uid: 'phlare',
+  type: 'profiling',
+  name: 'phlare',
+  access: 'proxy',
+  meta: {
+    id: 'phlare',
+    name: 'phlare',
+    type: PluginType.datasource,
+    info: {} as PluginMetaInfo,
+    module: '',
+    baseUrl: '',
+  },
+  jsonData: {},
+  readOnly: false,
+};
diff --git a/public/app/plugins/datasource/phlare/datasource.ts b/public/app/plugins/datasource/phlare/datasource.ts
index aa80869c12c..95da2c417de 100644
--- a/public/app/plugins/datasource/phlare/datasource.ts
+++ b/public/app/plugins/datasource/phlare/datasource.ts
@@ -1,8 +1,11 @@
+import Prism, { Grammar } from 'prismjs';
 import { Observable, of } from 'rxjs';
 
-import { DataQueryRequest, DataQueryResponse, DataSourceInstanceSettings } from '@grafana/data';
+import { AbstractQuery, DataQueryRequest, DataQueryResponse, DataSourceInstanceSettings } from '@grafana/data';
 import { DataSourceWithBackend } from '@grafana/runtime';
 
+import { extractLabelMatchers, toPromLikeExpr } from '../prometheus/language_utils';
+
 import { normalizeQuery } from './QueryEditor/QueryEditor';
 import { PhlareDataSourceOptions, Query, ProfileTypeMessage, SeriesMessage } from './types';
 
@@ -45,4 +48,58 @@ export class PhlareDataSource extends DataSourceWithBackend<Query, PhlareDataSou
   async getLabelNames(): Promise<string[]> {
     return await super.getResource('labelNames');
   }
+
+  async importFromAbstractQueries(abstractQueries: AbstractQuery[]): Promise<Query[]> {
+    return abstractQueries.map((abstractQuery) => this.importFromAbstractQuery(abstractQuery));
+  }
+
+  importFromAbstractQuery(labelBasedQuery: AbstractQuery): Query {
+    return {
+      refId: labelBasedQuery.refId,
+      labelSelector: toPromLikeExpr(labelBasedQuery),
+      queryType: 'both',
+      profileTypeId: '',
+      groupBy: [''],
+    };
+  }
+
+  async exportToAbstractQueries(queries: Query[]): Promise<AbstractQuery[]> {
+    return queries.map((query) => this.exportToAbstractQuery(query));
+  }
+
+  exportToAbstractQuery(query: Query): AbstractQuery {
+    const phlareQuery = query.labelSelector;
+    if (!phlareQuery || phlareQuery.length === 0) {
+      return { refId: query.refId, labelMatchers: [] };
+    }
+    const tokens = Prism.tokenize(phlareQuery, grammar);
+    return {
+      refId: query.refId,
+      labelMatchers: extractLabelMatchers(tokens),
+    };
+  }
 }
+
+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: /[{}(),.]/,
+};