Prometheus: Run annotation queries trough backend (#41059)

* Prometheus: Run annotation queries trough backend

* Be more explicit on what type of query we are running

* Fix typing, remove unused variables

public/app/plugins/datasource/prometheus/datasource.test.ts

@@ -871,26 +871,7 @@ describe('PrometheusDatasource', () => {
       },
     };
 
-    const response = {
-      status: 'success',
-      data: {
-        data: {
-          resultType: 'matrix',
-          result: [
-            {
-              metric: {
-                __name__: 'ALERTS',
-                alertname: 'InstanceDown',
-                alertstate: 'firing',
-                instance: 'testinstance',
-                job: 'testjob',
-              },
-              values: [[123, '1']],
-            },
-          ],
-        },
-      },
-    };
+    const response = createAnnotationResponse();
 
     describe('when time series query is cancelled', () => {
       it('should return empty results', async () => {
@@ -919,7 +900,7 @@ describe('PrometheusDatasource', () => {
         expect(results[0].tags).toContain('testjob');
         expect(results[0].title).toBe('InstanceDown');
         expect(results[0].text).toBe('testinstance');
-        expect(results[0].time).toBe(123 * 1000);
+        expect(results[0].time).toBe(123);
       });
     });
 
@@ -934,7 +915,7 @@ describe('PrometheusDatasource', () => {
       });
 
       it('should return annotation list', () => {
-        expect(results[0].time).toEqual(1);
+        expect(results[0].time).toEqual(456);
       });
     });
 
@@ -953,7 +934,7 @@ describe('PrometheusDatasource', () => {
         };
         ds.annotationQuery(query);
         const req = fetchMock.mock.calls[0][0];
-        expect(req.url).toContain('step=60');
+        expect(req.data.queries[0].interval).toBe('60s');
       });
 
       it('should use default step for short range when annotation step is empty string', () => {
@@ -970,7 +951,7 @@ describe('PrometheusDatasource', () => {
         };
         ds.annotationQuery(query);
         const req = fetchMock.mock.calls[0][0];
-        expect(req.url).toContain('step=60');
+        expect(req.data.queries[0].interval).toBe('60s');
       });
 
       it('should use custom step for short range', () => {
@@ -988,42 +969,7 @@ describe('PrometheusDatasource', () => {
         };
         ds.annotationQuery(query);
         const req = fetchMock.mock.calls[0][0];
-        expect(req.url).toContain('step=10');
-      });
-
-      it('should use custom step for short range', () => {
-        const annotation = {
-          ...options.annotation,
-          step: '10s',
-        };
-        const query = {
-          ...options,
-          annotation,
-          range: {
-            from: time({ seconds: 63 }),
-            to: time({ seconds: 123 }),
-          },
-        };
-        ds.annotationQuery(query);
-        const req = fetchMock.mock.calls[0][0];
-        expect(req.url).toContain('step=10');
-      });
-
-      it('should use dynamic step on long ranges if no option was given', () => {
-        const query = {
-          ...options,
-          range: {
-            from: time({ seconds: 63 }),
-            to: time({ hours: 24 * 30, seconds: 63 }),
-          },
-        };
-        ds.annotationQuery(query);
-        const req = fetchMock.mock.calls[0][0];
-        // Range in seconds: (to - from) / 1000
-        // Max_datapoints: 11000
-        // Step: range / max_datapoints
-        const step = 236;
-        expect(req.url).toContain(`step=${step}`);
+        expect(req.data.queries[0].interval).toBe('10s');
       });
     });
 
@@ -1041,21 +987,9 @@ describe('PrometheusDatasource', () => {
         },
       };
 
-      async function runAnnotationQuery(resultValues: Array<[number, string]>) {
-        const response = {
-          status: 'success',
-          data: {
-            data: {
-              resultType: 'matrix',
-              result: [
-                {
-                  metric: { __name__: 'test', job: 'testjob' },
-                  values: resultValues,
-                },
-              ],
-            },
-          },
-        };
+      async function runAnnotationQuery(data: number[][]) {
+        let response = createAnnotationResponse();
+        response.data.results['X'].frames[0].data.values = data;
 
         options.annotation.useValueForTime = false;
         fetchMock.mockImplementation(() => of(response));
@@ -1065,14 +999,8 @@ describe('PrometheusDatasource', () => {
 
       it('should handle gaps and inactive values', async () => {
         const results = await runAnnotationQuery([
-          [2 * 60, '1'],
-          [3 * 60, '1'],
-          // gap
-          [5 * 60, '1'],
-          [6 * 60, '1'],
-          [7 * 60, '1'],
-          [8 * 60, '0'], // false --> create new block
-          [9 * 60, '1'],
+          [2 * 60000, 3 * 60000, 5 * 60000, 6 * 60000, 7 * 60000, 8 * 60000, 9 * 60000],
+          [1, 1, 1, 1, 1, 0, 1],
         ]);
         expect(results.map((result) => [result.time, result.timeEnd])).toEqual([
           [120000, 180000],
@@ -1083,44 +1011,27 @@ describe('PrometheusDatasource', () => {
 
       it('should handle single region', async () => {
         const results = await runAnnotationQuery([
-          [2 * 60, '1'],
-          [3 * 60, '1'],
+          [2 * 60000, 3 * 60000],
+          [1, 1],
         ]);
         expect(results.map((result) => [result.time, result.timeEnd])).toEqual([[120000, 180000]]);
       });
 
       it('should handle 0 active regions', async () => {
         const results = await runAnnotationQuery([
-          [2 * 60, '0'],
-          [3 * 60, '0'],
-          [5 * 60, '0'],
+          [2 * 60000, 3 * 60000, 5 * 60000],
+          [0, 0, 0],
         ]);
         expect(results.length).toBe(0);
       });
 
       it('should handle single active value', async () => {
-        const results = await runAnnotationQuery([[2 * 60, '1']]);
+        const results = await runAnnotationQuery([[2 * 60000], [1]]);
         expect(results.map((result) => [result.time, result.timeEnd])).toEqual([[120000, 120000]]);
       });
     });
   });
 
-  describe('createAnnotationQueryOptions', () => {
-    it.each`
-      options                                | expected
-      ${{}}                                  | ${{ interval: '60s' }}
-      ${{ annotation: {} }}                  | ${{ annotation: {}, interval: '60s' }}
-      ${{ annotation: { step: undefined } }} | ${{ annotation: { step: undefined }, interval: '60s' }}
-      ${{ annotation: { step: null } }}      | ${{ annotation: { step: null }, interval: '60s' }}
-      ${{ annotation: { step: '' } }}        | ${{ annotation: { step: '' }, interval: '60s' }}
-      ${{ annotation: { step: 0 } }}         | ${{ annotation: { step: 0 }, interval: '60s' }}
-      ${{ annotation: { step: 5 } }}         | ${{ annotation: { step: 5 }, interval: '60s' }}
-      ${{ annotation: { step: '5m' } }}      | ${{ annotation: { step: '5m' }, interval: '5m' }}
-    `("when called with options: '$options'", ({ options, expected }) => {
-      expect(ds.createAnnotationQueryOptions(options)).toEqual(expected);
-    });
-  });
-
   describe('When resultFormat is table and instant = true', () => {
     let results: any;
     const query = {
@@ -2251,3 +2162,50 @@ function createDefaultPromResponse() {
     },
   };
 }
+
+function createAnnotationResponse() {
+  const response = {
+    data: {
+      results: {
+        X: {
+          frames: [
+            {
+              schema: {
+                name: 'bar',
+                refId: 'X',
+                fields: [
+                  {
+                    name: 'Time',
+                    type: 'time',
+                    typeInfo: {
+                      frame: 'time.Time',
+                    },
+                  },
+                  {
+                    name: 'Value',
+                    type: 'number',
+                    typeInfo: {
+                      frame: 'float64',
+                    },
+                    labels: {
+                      __name__: 'ALERTS',
+                      alertname: 'InstanceDown',
+                      alertstate: 'firing',
+                      instance: 'testinstance',
+                      job: 'testjob',
+                    },
+                  },
+                ],
+              },
+              data: {
+                values: [[123], [456]],
+              },
+            },
+          ],
+        },
+      },
+    },
+  };
+
+  return { ...response };
+}

public/app/plugins/datasource/prometheus/datasource.ts

@@ -15,8 +15,17 @@ import {
   rangeUtil,
   ScopedVars,
   TimeRange,
+  DataFrame,
 } from '@grafana/data';
-import { BackendSrvRequest, FetchError, FetchResponse, getBackendSrv, DataSourceWithBackend } from '@grafana/runtime';
+import {
+  BackendSrvRequest,
+  FetchError,
+  FetchResponse,
+  getBackendSrv,
+  DataSourceWithBackend,
+  BackendDataSourceResponse,
+  toDataQueryResponse,
+} from '@grafana/runtime';
 
 import { safeStringifyValue } from 'app/core/utils/explore';
 import { getTimeSrv, TimeSrv } from 'app/features/dashboard/services/TimeSrv';
@@ -28,7 +37,6 @@ import { getInitHints, getQueryHints } from './query_hints';
 import { getOriginalMetricName, renderTemplate, transform, transformV2 } from './result_transformer';
 import {
   ExemplarTraceIdDestination,
-  isFetchErrorResponse,
   PromDataErrorResponse,
   PromDataSuccessResponse,
   PromExemplarData,
@@ -51,6 +59,7 @@ export class PrometheusDatasource extends DataSourceWithBackend<PromQuery, PromO
   editorSrc: string;
   ruleMappings: { [index: string]: string };
   url: string;
+  id: number;
   directUrl: string;
   access: 'direct' | 'proxy';
   basicAuth: any;
@@ -74,6 +83,7 @@ export class PrometheusDatasource extends DataSourceWithBackend<PromQuery, PromO
 
     this.type = 'prometheus';
     this.editorSrc = 'app/features/prometheus/partials/query.editor.html';
+    this.id = instanceSettings.id;
     this.url = instanceSettings.url!;
     this.access = instanceSettings.access;
     this.basicAuth = instanceSettings.basicAuth;
@@ -645,106 +655,126 @@ export class PrometheusDatasource extends DataSourceWithBackend<PromQuery, PromO
     };
   }
 
-  createAnnotationQueryOptions = (options: any): DataQueryRequest<PromQuery> => {
-    const annotation = options.annotation;
-    const interval =
-      annotation && annotation.step && typeof annotation.step === 'string'
-        ? annotation.step
-        : ANNOTATION_QUERY_STEP_DEFAULT;
-    return {
-      ...options,
-      interval,
-    };
-  };
-
   async annotationQuery(options: any): Promise<AnnotationEvent[]> {
     const annotation = options.annotation;
-    const { expr = '', tagKeys = '', titleFormat = '', textFormat = '' } = annotation;
+    const { expr = '' } = annotation;
 
     if (!expr) {
       return Promise.resolve([]);
     }
 
-    const start = this.getPrometheusTime(options.range.from, false);
-    const end = this.getPrometheusTime(options.range.to, true);
-    const queryOptions = this.createAnnotationQueryOptions(options);
-
-    // Unsetting min interval for accurate event resolution
-    const minStep = '1s';
+    const step = options.annotation.step || ANNOTATION_QUERY_STEP_DEFAULT;
     const queryModel = {
       expr,
-      interval: minStep,
+      range: true,
+      instant: false,
+      exemplar: false,
+      interval: step,
+      queryType: PromQueryType.timeSeriesQuery,
       refId: 'X',
-      requestId: `prom-query-${annotation.name}`,
+      datasourceId: this.id,
     };
 
-    const query = this.createQuery(queryModel, queryOptions, start, end);
-    const response = await lastValueFrom(this.performTimeSeriesQuery(query, query.start, query.end));
-    const eventList: AnnotationEvent[] = [];
-    const splitKeys = tagKeys.split(',');
+    return await lastValueFrom(
+      getBackendSrv()
+        .fetch<BackendDataSourceResponse>({
+          url: '/api/ds/query',
+          method: 'POST',
+          data: {
+            from: (this.getPrometheusTime(options.range.from, false) * 1000).toString(),
+            to: (this.getPrometheusTime(options.range.to, true) * 1000).toString(),
+            queries: [queryModel],
+          },
+          requestId: `prom-query-${annotation.name}`,
+        })
+        .pipe(
+          map((rsp: FetchResponse<BackendDataSourceResponse>) => {
+            return this.processsAnnotationResponse(options, rsp.data);
+          })
+        )
+    );
+  }
 
-    if (isFetchErrorResponse(response) && response.cancelled) {
+  processsAnnotationResponse = (options: any, data: BackendDataSourceResponse) => {
+    const frames: DataFrame[] = toDataQueryResponse({ data: data }).data;
+    if (!frames || !frames.length) {
       return [];
     }
 
-    const step = Math.floor(query.step ?? 15) * 1000;
+    const annotation = options.annotation;
+    const { tagKeys = '', titleFormat = '', textFormat = '' } = annotation;
 
-    response?.data?.data?.result?.forEach((series) => {
-      const tags = Object.entries(series.metric)
-        .filter(([k]) => splitKeys.includes(k))
-        .map(([_k, v]: [string, string]) => v);
+    const step = rangeUtil.intervalToSeconds(annotation.step || ANNOTATION_QUERY_STEP_DEFAULT) * 1000;
+    const tagKeysArray = tagKeys.split(',');
+    const frame = frames[0];
+    const timeField = frame.fields[0];
+    const valueField = frame.fields[1];
+    const labels = valueField?.labels || {};
 
-      series.values.forEach((value: any[]) => {
-        let timestampValue;
-        // rewrite timeseries to a common format
-        if (annotation.useValueForTime) {
-          timestampValue = Math.floor(parseFloat(value[1]));
-          value[1] = 1;
-        } else {
-          timestampValue = Math.floor(parseFloat(value[0])) * 1000;
-        }
-        value[0] = timestampValue;
-      });
+    const tags = Object.keys(labels)
+      .filter((label) => tagKeysArray.includes(label))
+      .map((label) => labels[label]);
 
-      const activeValues = series.values.filter((value) => parseFloat(value[1]) >= 1);
-      const activeValuesTimestamps = activeValues.map((value) => value[0]);
+    const timeValueTuple: Array<[number, number]> = [];
 
-      // Instead of creating singular annotation for each active event we group events into region if they are less
-      // then `step` apart.
-      let latestEvent: AnnotationEvent | null = null;
+    let idx = 0;
+    valueField.values.toArray().forEach((value: string) => {
+      let timeStampValue: number;
+      let valueValue: number;
+      const time = timeField.values.get(idx);
 
-      for (const timestamp of activeValuesTimestamps) {
-        // We already have event `open` and we have new event that is inside the `step` so we just update the end.
-        if (latestEvent && (latestEvent.timeEnd ?? 0) + step >= timestamp) {
-          latestEvent.timeEnd = timestamp;
-          continue;
-        }
-
-        // Event exists but new one is outside of the `step` so we "finish" the current region.
-        if (latestEvent) {
-          eventList.push(latestEvent);
-        }
-
-        // We start a new region.
-        latestEvent = {
-          time: timestamp,
-          timeEnd: timestamp,
-          annotation,
-          title: renderTemplate(titleFormat, series.metric),
-          tags,
-          text: renderTemplate(textFormat, series.metric),
-        };
+      // If we want to use value as a time, we use value as timeStampValue and valueValue will be 1
+      if (options.annotation.useValueForTime) {
+        timeStampValue = Math.floor(parseFloat(value));
+        valueValue = 1;
+      } else {
+        timeStampValue = Math.floor(parseFloat(time));
+        valueValue = parseFloat(value);
       }
 
-      if (latestEvent) {
-        // finish up last point if we have one
-        latestEvent.timeEnd = activeValuesTimestamps[activeValuesTimestamps.length - 1];
-        eventList.push(latestEvent);
-      }
+      idx++;
+      timeValueTuple.push([timeStampValue, valueValue]);
     });
 
+    const activeValues = timeValueTuple.filter((value) => value[1] >= 1);
+    const activeValuesTimestamps = activeValues.map((value) => value[0]);
+
+    // Instead of creating singular annotation for each active event we group events into region if they are less
+    // or equal to `step` apart.
+    const eventList: AnnotationEvent[] = [];
+    let latestEvent: AnnotationEvent | null = null;
+
+    for (const timestamp of activeValuesTimestamps) {
+      // We already have event `open` and we have new event that is inside the `step` so we just update the end.
+      if (latestEvent && (latestEvent.timeEnd ?? 0) + step >= timestamp) {
+        latestEvent.timeEnd = timestamp;
+        continue;
+      }
+
+      // Event exists but new one is outside of the `step` so we add it to eventList.
+      if (latestEvent) {
+        eventList.push(latestEvent);
+      }
+
+      // We start a new region.
+      latestEvent = {
+        time: timestamp,
+        timeEnd: timestamp,
+        annotation,
+        title: renderTemplate(titleFormat, labels),
+        tags,
+        text: renderTemplate(textFormat, labels),
+      };
+    }
+
+    if (latestEvent) {
+      // Finish up last point if we have one
+      latestEvent.timeEnd = activeValuesTimestamps[activeValuesTimestamps.length - 1];
+      eventList.push(latestEvent);
+    }
+
     return eventList;
-  }
+  };
 
   getExemplars(query: PromQueryRequest) {
     const url = '/api/v1/query_exemplars';

public/app/plugins/datasource/prometheus/types.ts

@@ -1,5 +1,4 @@
 import { DataQuery, DataSourceJsonData, QueryResultMeta, ScopedVars } from '@grafana/data';
-import { FetchError } from '@grafana/runtime';
 
 export interface PromQuery extends DataQuery {
   expr: string;
@@ -114,10 +113,6 @@ export interface PromMetric {
   [index: string]: any;
 }
 
-export function isFetchErrorResponse(response: any): response is FetchError {
-  return 'cancelled' in response;
-}
-
 export function isMatrixData(result: MatrixOrVectorResult): result is PromMatrixData['result'][0] {
   return 'values' in result;
 }