From b307ff0edc857b09bef91a8440cd14f3ac22c096 Mon Sep 17 00:00:00 2001
From: kay delaney <45561153+kaydelaney@users.noreply.github.com>
Date: Tue, 5 May 2020 13:19:52 +0100
Subject: [PATCH] Datasource/Loki: Fixes issue where cached log values weren't
 merged when labels were refreshed (#24101)

* Datasource/Loki: Fixes issue where cached log values weren't merged when labels were refreshed

Closes #24087

(cherry picked from commit 79a084392f7c844954cbeaa0834a9d4a400a7dd9)
---
 .../LokiExploreQueryEditor.test.tsx.snap      |  1 +
 .../datasource/loki/language_provider.test.ts | 34 +++++++++++++++++++
 .../datasource/loki/language_provider.ts      | 25 ++++++++------
 3 files changed, 50 insertions(+), 10 deletions(-)

diff --git a/public/app/plugins/datasource/loki/components/__snapshots__/LokiExploreQueryEditor.test.tsx.snap b/public/app/plugins/datasource/loki/components/__snapshots__/LokiExploreQueryEditor.test.tsx.snap
index badd8b9155c..3b5b08c7817 100644
--- a/public/app/plugins/datasource/loki/components/__snapshots__/LokiExploreQueryEditor.test.tsx.snap
+++ b/public/app/plugins/datasource/loki/components/__snapshots__/LokiExploreQueryEditor.test.tsx.snap
@@ -54,6 +54,7 @@ exports[`LokiExploreQueryEditor should render component 1`] = `
   datasource={
     Object {
       "languageProvider": LokiLanguageProvider {
+        "addLabelValuesToOptions": [Function],
         "cleanText": [Function],
         "datasource": [Circular],
         "fetchSeriesLabels": [Function],
diff --git a/public/app/plugins/datasource/loki/language_provider.test.ts b/public/app/plugins/datasource/loki/language_provider.test.ts
index 60f689c00e3..65f4ed07487 100644
--- a/public/app/plugins/datasource/loki/language_provider.test.ts
+++ b/public/app/plugins/datasource/loki/language_provider.test.ts
@@ -144,6 +144,40 @@ describe('Language completion provider', () => {
       ]);
     });
   });
+
+  describe('label values', () => {
+    it('should fetch label values if not cached', async () => {
+      const absoluteRange: AbsoluteTimeRange = {
+        from: 0,
+        to: 5000,
+      };
+
+      const datasource = makeMockLokiDatasource({ testkey: ['label1_val1', 'label1_val2'], label2: [] });
+      const provider = await getLanguageProvider(datasource);
+      const requestSpy = jest.spyOn(provider, 'request');
+      const labelValues = await provider.fetchLabelValues('testkey', absoluteRange);
+      expect(requestSpy).toHaveBeenCalled();
+      expect(labelValues).toEqual(['label1_val1', 'label1_val2']);
+    });
+
+    it('should return cached values', async () => {
+      const absoluteRange: AbsoluteTimeRange = {
+        from: 0,
+        to: 5000,
+      };
+
+      const datasource = makeMockLokiDatasource({ testkey: ['label1_val1', 'label1_val2'], label2: [] });
+      const provider = await getLanguageProvider(datasource);
+      const requestSpy = jest.spyOn(provider, 'request');
+      const labelValues = await provider.fetchLabelValues('testkey', absoluteRange);
+      expect(requestSpy).toHaveBeenCalledTimes(1);
+      expect(labelValues).toEqual(['label1_val1', 'label1_val2']);
+
+      const nextLabelValues = await provider.fetchLabelValues('testkey', absoluteRange);
+      expect(requestSpy).toHaveBeenCalledTimes(1);
+      expect(nextLabelValues).toEqual(['label1_val1', 'label1_val2']);
+    });
+  });
 });
 
 describe('Request URL', () => {
diff --git a/public/app/plugins/datasource/loki/language_provider.ts b/public/app/plugins/datasource/loki/language_provider.ts
index 939125a51c0..002f4dc3a8b 100644
--- a/public/app/plugins/datasource/loki/language_provider.ts
+++ b/public/app/plugins/datasource/loki/language_provider.ts
@@ -449,6 +449,7 @@ export default class LokiLanguageProvider extends LanguageProvider {
 
     const cacheKey = this.generateCacheKey(url, start, end, key);
     const params = { start, end };
+
     let value = this.labelsCache.get(cacheKey);
     if (!value) {
       try {
@@ -459,20 +460,24 @@ export default class LokiLanguageProvider extends LanguageProvider {
         value = values;
         this.labelsCache.set(cacheKey, value);
 
-        // Add to label options
-        this.logLabelOptions = this.logLabelOptions.map(keyOption => {
-          if (keyOption.value === key) {
-            return {
-              ...keyOption,
-              children: values.map(value => ({ label: value, value })),
-            };
-          }
-          return keyOption;
-        });
+        this.logLabelOptions = this.addLabelValuesToOptions(key, values);
       } catch (e) {
         console.error(e);
       }
+    } else {
+      this.logLabelOptions = this.addLabelValuesToOptions(key, value);
     }
     return value;
   }
+
+  private addLabelValuesToOptions = (labelKey: string, values: string[]) => {
+    return this.logLabelOptions.map(keyOption =>
+      keyOption.value === labelKey
+        ? {
+            ...keyOption,
+            children: values.map(value => ({ label: value, value })),
+          }
+        : keyOption
+    );
+  };
 }