From 42662667d15ffc4f4be1411a0ea1df0c696f3471 Mon Sep 17 00:00:00 2001
From: Brendan O'Handley <brendan.ohandley@grafana.com>
Date: Wed, 30 Aug 2023 09:52:23 -0400
Subject: [PATCH] Prometheus: Query variable editor persist query type on
 clicking run query (#74045)

run query to save query type selection
---
 .../components/VariableQueryEditor.test.tsx   | 15 ++++------
 .../components/VariableQueryEditor.tsx        |  4 +--
 .../migrations/variableMigration.ts           | 28 +++++++++++--------
 3 files changed, 24 insertions(+), 23 deletions(-)

diff --git a/public/app/plugins/datasource/prometheus/components/VariableQueryEditor.test.tsx b/public/app/plugins/datasource/prometheus/components/VariableQueryEditor.test.tsx
index df6c715954c..8a722907550 100644
--- a/public/app/plugins/datasource/prometheus/components/VariableQueryEditor.test.tsx
+++ b/public/app/plugins/datasource/prometheus/components/VariableQueryEditor.test.tsx
@@ -170,7 +170,7 @@ describe('PromVariableQueryEditor', () => {
     });
   });
 
-  test('Calls onChange for label_names() query', async () => {
+  test('Calls onChange for label_names, label_values, metrics, and query result queries', async () => {
     const onChange = jest.fn();
 
     props.query = {
@@ -181,21 +181,18 @@ describe('PromVariableQueryEditor', () => {
     render(<PromVariableQueryEditor {...props} onChange={onChange} />);
 
     await selectOptionInTest(screen.getByLabelText('Query type'), 'Label names');
+    await selectOptionInTest(screen.getByLabelText('Query type'), 'Label values');
+    await selectOptionInTest(screen.getByLabelText('Query type'), 'Metrics');
+    await selectOptionInTest(screen.getByLabelText('Query type'), 'Query result');
 
-    expect(onChange).toHaveBeenCalledWith({
-      query: 'label_names()',
-      refId,
-    });
+    expect(onChange).toHaveBeenCalledTimes(4);
   });
 
-  test('Does not call onChange for other queries', async () => {
+  test('Does not call onChange for series query', async () => {
     const onChange = jest.fn();
 
     render(<PromVariableQueryEditor {...props} onChange={onChange} />);
 
-    await selectOptionInTest(screen.getByLabelText('Query type'), 'Label values');
-    await selectOptionInTest(screen.getByLabelText('Query type'), 'Metrics');
-    await selectOptionInTest(screen.getByLabelText('Query type'), 'Query result');
     await selectOptionInTest(screen.getByLabelText('Query type'), 'Series query');
 
     expect(onChange).not.toHaveBeenCalled();
diff --git a/public/app/plugins/datasource/prometheus/components/VariableQueryEditor.tsx b/public/app/plugins/datasource/prometheus/components/VariableQueryEditor.tsx
index 33cc68f2320..c448b4058df 100644
--- a/public/app/plugins/datasource/prometheus/components/VariableQueryEditor.tsx
+++ b/public/app/plugins/datasource/prometheus/components/VariableQueryEditor.tsx
@@ -135,8 +135,8 @@ export const PromVariableQueryEditor = ({ onChange, query, datasource }: Props)
   /** Call onchange for label names query type change */
   const onQueryTypeChange = (newType: SelectableValue<QueryType>) => {
     setQryType(newType.value);
-    if (newType.value === QueryType.LabelNames) {
-      onChangeWithVariableString({ qryType: newType.value });
+    if (newType.value !== QueryType.SeriesQuery) {
+      onChangeWithVariableString({ qryType: newType.value ?? 0 });
     }
   };
 
diff --git a/public/app/plugins/datasource/prometheus/migrations/variableMigration.ts b/public/app/plugins/datasource/prometheus/migrations/variableMigration.ts
index 260d00a84d3..ec0df3dc5da 100644
--- a/public/app/plugins/datasource/prometheus/migrations/variableMigration.ts
+++ b/public/app/plugins/datasource/prometheus/migrations/variableMigration.ts
@@ -38,11 +38,11 @@ export function migrateVariableQueryToEditor(rawQuery: string | PromVariableQuer
     };
   }
 
-  const labelValues = rawQuery.match(PrometheusLabelValuesRegex);
-
-  if (labelValues) {
-    const label = labelValues[2];
-    const metric = labelValues[1];
+  const labelValuesCheck = rawQuery.match(/^label_values\(/);
+  if (labelValuesCheck) {
+    const labelValues = rawQuery.match(PrometheusLabelValuesRegex);
+    const label = labelValues ? labelValues[2] : '';
+    const metric = labelValues ? labelValues[1] : '';
 
     if (metric) {
       const visQuery = buildVisualQueryFromString(metric);
@@ -62,26 +62,30 @@ export function migrateVariableQueryToEditor(rawQuery: string | PromVariableQuer
     }
   }
 
-  const metricNames = rawQuery.match(PrometheusMetricNamesRegex);
-  if (metricNames) {
+  const metricNamesCheck = rawQuery.match(/^metrics\(/);
+  if (metricNamesCheck) {
+    const metricNames = rawQuery.match(PrometheusMetricNamesRegex);
+    const metric = metricNames ? metricNames[1] : '';
     return {
       ...queryBase,
       qryType: QueryType.MetricNames,
-      metric: metricNames[1],
+      metric,
     };
   }
 
-  const queryResult = rawQuery.match(PrometheusQueryResultRegex);
-  if (queryResult) {
+  const queryResultCheck = rawQuery.match(/^query_result\(/);
+  if (queryResultCheck) {
+    const queryResult = rawQuery.match(PrometheusQueryResultRegex);
+    const varQuery = queryResult ? queryResult[1] : '';
     return {
       ...queryBase,
       qryType: QueryType.VarQueryResult,
-      varQuery: queryResult[1],
+      varQuery,
     };
   }
 
   // seriesQuery does not have a function and no regex above
-  if (!labelNames && !labelValues && !metricNames && !queryResult) {
+  if (!labelNames && !labelValuesCheck && !metricNamesCheck && !queryResultCheck) {
     return {
       ...queryBase,
       qryType: QueryType.SeriesQuery,