From f5b4fc58e369ca252c28ebefd734b6f294464105 Mon Sep 17 00:00:00 2001
From: Sven Grossmann <sven.grossmann@grafana.com>
Date: Fri, 26 Jul 2024 11:55:04 +0200
Subject: [PATCH] Loki: Also replace `step` with vars (#91031)

---
 .../datasource/loki/datasource.test.ts        | 50 ++++++++++++++++++-
 .../app/plugins/datasource/loki/datasource.ts |  6 ++-
 2 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/public/app/plugins/datasource/loki/datasource.test.ts b/public/app/plugins/datasource/loki/datasource.test.ts
index 292290b1f0a..7fb44ddaf39 100644
--- a/public/app/plugins/datasource/loki/datasource.test.ts
+++ b/public/app/plugins/datasource/loki/datasource.test.ts
@@ -1875,7 +1875,7 @@ describe('applyTemplateVariables', () => {
         { expr: 'rate({job="grafana"}[$__interval]) + rate({job="grafana"}[$__interval_ms])', refId: 'A' },
         scopedVars
       );
-      expect(templateSrvMock.replace).toHaveBeenCalledTimes(2);
+      expect(templateSrvMock.replace).toHaveBeenCalledTimes(3);
       // Interpolated legend
       expect(templateSrvMock.replace).toHaveBeenCalledWith(
         undefined,
@@ -1910,7 +1910,7 @@ describe('applyTemplateVariables', () => {
         },
         scopedVars
       );
-      expect(templateSrvMock.replace).toHaveBeenCalledTimes(2);
+      expect(templateSrvMock.replace).toHaveBeenCalledTimes(3);
       // Interpolated legend
       expect(templateSrvMock.replace).toHaveBeenCalledWith(
         undefined,
@@ -1931,6 +1931,52 @@ describe('applyTemplateVariables', () => {
         expect.any(Function)
       );
     });
+
+    it('should replace step', () => {
+      const templateSrvMock = {
+        getAdhocFilters: jest.fn().mockImplementation((query: string) => query),
+        replace: jest.fn((a: string | undefined, ...rest: unknown[]) => a?.replace('$testVariable', 'foo')),
+      } as unknown as TemplateSrv;
+
+      const ds = createLokiDatasource(templateSrvMock);
+      ds.addAdHocFilters = jest.fn().mockImplementation((query: string) => query);
+      const replacedQuery = ds.applyTemplateVariables(
+        {
+          expr: 'rate({job="grafana"}[$__range]) + rate({job="grafana"}[$__range_ms]) + rate({job="grafana"}[$__range_s])',
+          refId: 'A',
+          step: '$testVariable',
+        },
+        scopedVars
+      );
+      expect(replacedQuery).toEqual(
+        expect.objectContaining({
+          step: 'foo',
+        })
+      );
+    });
+
+    it('should replace legendFormat', () => {
+      const templateSrvMock = {
+        getAdhocFilters: jest.fn().mockImplementation((query: string) => query),
+        replace: jest.fn((a: string | undefined, ...rest: unknown[]) => a?.replace('$testVariable', 'foo')),
+      } as unknown as TemplateSrv;
+
+      const ds = createLokiDatasource(templateSrvMock);
+      ds.addAdHocFilters = jest.fn().mockImplementation((query: string) => query);
+      const replacedQuery = ds.applyTemplateVariables(
+        {
+          expr: 'rate({job="grafana"}[$__range]) + rate({job="grafana"}[$__range_ms]) + rate({job="grafana"}[$__range_s])',
+          refId: 'A',
+          legendFormat: '$testVariable',
+        },
+        scopedVars
+      );
+      expect(replacedQuery).toEqual(
+        expect.objectContaining({
+          legendFormat: 'foo',
+        })
+      );
+    });
   });
 
   describe('getStatsTimeRange', () => {
diff --git a/public/app/plugins/datasource/loki/datasource.ts b/public/app/plugins/datasource/loki/datasource.ts
index c4cc01c9357..80a95dda5a5 100644
--- a/public/app/plugins/datasource/loki/datasource.ts
+++ b/public/app/plugins/datasource/loki/datasource.ts
@@ -1106,10 +1106,14 @@ export class LokiDatasource
       adhocFilters
     );
 
+    const step = this.templateSrv.replace(target.step, variables);
+    const legendFormat = this.templateSrv.replace(target.legendFormat, variables);
+
     return {
       ...target,
-      legendFormat: this.templateSrv.replace(target.legendFormat, rest),
       expr: exprWithAdHoc,
+      step,
+      legendFormat,
     };
   }