Prefix requestId with data source id (#27)

Co-authored-by: Kevin Yu <[email protected]>
grafana · Feb 15, 2024 · c78cb2d · c78cb2d
1 parent 1a7c93f
commit c78cb2d
Show file tree

Hide file tree

Showing 3 changed files with 84 additions and 10 deletions.
diff --git a/src/DatasourceWithAsyncBackend.test.ts b/src/DatasourceWithAsyncBackend.test.ts
@@ -1,12 +1,25 @@
-import { DataQuery, DataSourceInstanceSettings, PluginType, getDefaultTimeRange } from '@grafana/data';
+import {
+  DataQuery,
+  DataQueryRequest,
+  DataSourceInstanceSettings,
+  PluginType,
+  getDefaultTimeRange,
+} from '@grafana/data';
 import { DataSourceWithBackend } from '@grafana/runtime';
 import { DatasourceWithAsyncBackend } from './DatasourceWithAsyncBackend';
+import { RequestLoopOptions } from 'requestLooper';
 
 const queryMock = jest.fn().mockImplementation(() => Promise.resolve({ data: [] }));
 jest.spyOn(DataSourceWithBackend.prototype, 'query').mockImplementation(queryMock);
 
+const getRequestLooperMock = jest.fn();
+jest.mock('./requestLooper.ts', () => ({
+  ...jest.requireActual('./requestLooper.ts'),
+  getRequestLooper: (req: DataQueryRequest, options: RequestLoopOptions) => getRequestLooperMock(req, options),
+}));
+
 const defaultInstanceSettings: DataSourceInstanceSettings<{}> = {
-  id: 1,
+  id: 12,
   uid: 'test',
   type: 'test',
   name: 'test',
@@ -33,6 +46,7 @@ const defaultInstanceSettings: DataSourceInstanceSettings<{}> = {
   },
   access: 'direct',
   jsonData: {},
+  readOnly: false,
 };
 const defaultQuery = { refId: 'refId-1' };
 const defaultQuery2 = { refId: 'refId-2' };
@@ -106,4 +120,17 @@ describe('DatasourceWithAsyncBackend', () => {
     expect(queryMock).toHaveBeenCalledTimes(1);
     expect(queryMock).toHaveBeenCalledWith(defaultRequest);
   });
+
+  it('uses the datasource id for the request id', () => {
+    const ds = setupDatasourceWithAsyncBackend({ asyncQueryDataSupport: true });
+    expect(getRequestLooperMock).not.toHaveBeenCalled();
+    ds.doSingle(defaultQuery, defaultRequest);
+    expect(getRequestLooperMock).toHaveBeenCalledTimes(1);
+    const expectedRequest = {
+      ...defaultRequest,
+      targets: [defaultQuery],
+      requestId: '12_100',
+    };
+    expect(getRequestLooperMock).toHaveBeenCalledWith(expectedRequest, expect.anything());
+  });
 });
diff --git a/src/DatasourceWithAsyncBackend.ts b/src/DatasourceWithAsyncBackend.ts
@@ -6,9 +6,15 @@ import {
   DataSourceInstanceSettings,
   DataSourceJsonData,
 } from '@grafana/data';
-import { BackendDataSourceResponse, DataSourceWithBackend, config, getBackendSrv, toDataQueryResponse } from '@grafana/runtime';
+import {
+  BackendDataSourceResponse,
+  DataSourceWithBackend,
+  config,
+  getBackendSrv,
+  toDataQueryResponse,
+} from '@grafana/runtime';
 import { merge, Observable, of } from 'rxjs';
-import { catchError,map } from 'rxjs/operators';
+import { catchError, map } from 'rxjs/operators';
 import { getRequestLooper } from './requestLooper';
 
 export interface CustomMeta {
@@ -38,9 +44,11 @@ export class DatasourceWithAsyncBackend<
   private runningQueries: { [hash: string]: RunningQueryInfo } = {};
   private requestCounter = 100;
   private asyncQueryDataSupport: boolean;
+  private requestIdPrefix: number;
 
   constructor(instanceSettings: DataSourceInstanceSettings<TOptions>, asyncQueryDataSupport = false) {
     super(instanceSettings);
+    this.requestIdPrefix = instanceSettings.id;
     this.asyncQueryDataSupport = asyncQueryDataSupport;
   }
 
@@ -85,7 +93,7 @@ export class DatasourceWithAsyncBackend<
     let allData: DataFrame[] = [];
 
     return getRequestLooper(
-      { ...request, targets: [target], requestId: `aws_ts_${this.requestCounter++}` },
+      { ...request, targets: [target], requestId: `${this.requestIdPrefix}_${this.requestCounter++}` },
       {
         /**
          * Additional query to execute if the current query is still in a running state
@@ -136,7 +144,8 @@ export class DatasourceWithAsyncBackend<
           };
 
           let headers = {};
-          const cachingDisabled = !config.featureToggles.useCachingService || !config.featureToggles.awsAsyncQueryCaching
+          const cachingDisabled =
+            !config.featureToggles.useCachingService || !config.featureToggles.awsAsyncQueryCaching;
           if (cachingDisabled && isRunning(status)) {
             // bypass query caching for Grafana Enterprise to
             // prevent an infinite loop
@@ -152,9 +161,12 @@ export class DatasourceWithAsyncBackend<
 
           return getBackendSrv()
             .fetch<BackendDataSourceResponse>(options)
-            .pipe(map((result) => ({ data: toDataQueryResponse(result).data })), catchError((err) => {
-              return of(toDataQueryResponse(err));
-            }));
+            .pipe(
+              map((result) => ({ data: toDataQueryResponse(result).data })),
+              catchError((err) => {
+                return of(toDataQueryResponse(err));
+              })
+            );
         },
 
         /**

diff --git a/src/requestLooper.test.ts b/src/requestLooper.test.ts
@@ -1,6 +1,7 @@
 import { RequestLoopOptions, getRequestLooper } from 'requestLooper';
 import { TestScheduler } from 'rxjs/testing';
-import { getDefaultTimeRange, DataQuery, LoadingState } from '@grafana/data';
+import { of } from 'rxjs';
+import { getDefaultTimeRange, DataQuery, LoadingState, DataQueryRequest } from '@grafana/data';
 
 const mockQuery: DataQuery = {
   refId: '',
@@ -90,4 +91,38 @@ describe('requestLooper', () => {
     scheduler.flush();
     expect(onCancel).toBeCalledTimes(1);
   });
+
+  it('increments the request id', (done) => {
+    let requestIds: string[] = [];
+
+    const queryMock = jest.fn().mockImplementation((req: DataQueryRequest) => {
+      requestIds.push(req.requestId);
+      return of({ data: [], state: LoadingState.Loading, meta });
+    });
+
+    const opt: RequestLoopOptions = {
+      getNextQuery: jest
+        .fn()
+        .mockImplementationOnce(() => ({ ...mockQuery, queryId: 'queryId' }))
+        .mockImplementationOnce(() => ({ ...mockQuery, queryId: 'queryId' }))
+        .mockImplementationOnce(() => undefined),
+      query: (req) => queryMock(req),
+      onCancel: jest.fn(),
+      process: jest.fn().mockImplementation(() => []),
+      shouldCancel: jest.fn().mockImplementation(() => false),
+    };
+
+    const looper = getRequestLooper(request, opt);
+
+    looper.subscribe({
+      next: () => {},
+      complete: () => {
+        expect(requestIds).toHaveLength(3);
+        expect(requestIds[0]).toBe(request.requestId);
+        expect(requestIds[1]).toBe(request.requestId + '.2');
+        expect(requestIds[2]).toBe(request.requestId + '.3');
+        done();
+      },
+    });
+  });
 });