DataSourceWithBackend: Add plugin id to the request headers (#58082)

packages/grafana-runtime/src/utils/DataSourceWithBackend.test.ts

@@ -32,7 +32,10 @@ jest.mock('../services', () => ({
   getBackendSrv: () => backendSrv,
   getDataSourceSrv: () => {
     return {
-      getInstanceSettings: (ref?: DataSourceRef) => ({ type: ref?.type ?? '?', uid: ref?.uid ?? '?' }),
+      getInstanceSettings: (ref?: DataSourceRef) => ({
+        type: ref?.type ?? '<mocktype>',
+        uid: ref?.uid ?? '<mockuid>',
+      }),
     };
   },
 }));
@@ -43,6 +46,8 @@ describe('DataSourceWithBackend', () => {
       maxDataPoints: 10,
       intervalMs: 5000,
       targets: [{ refId: 'A' }, { refId: 'B', datasource: { type: 'sample' } }],
+      dashboardUID: 'dashA',
+      panelId: 123,
     } as DataQueryRequest);
 
     const args = mock.calls[0][0];
@@ -65,7 +70,7 @@ describe('DataSourceWithBackend', () => {
             Object {
               "datasource": Object {
                 "type": "sample",
-                "uid": "?",
+                "uid": "<mockuid>",
               },
               "datasourceId": undefined,
               "intervalMs": 5000,
@@ -74,6 +79,12 @@ describe('DataSourceWithBackend', () => {
             },
           ],
         },
+        "headers": Object {
+          "X-Dashboard-Uid": "dashA",
+          "X-Datasource-Uid": "abc, <mockuid>",
+          "X-Panel-Id": "123",
+          "X-Plugin-Id": "dummy, sample",
+        },
         "hideFromInspector": false,
         "method": "POST",
         "requestId": undefined,
@@ -88,6 +99,8 @@ describe('DataSourceWithBackend', () => {
       intervalMs: 5000,
       targets: [{ refId: 'A' }, { refId: 'B', datasource: { type: 'sample' } }],
       hideFromInspector: true,
+      dashboardUID: 'dashA',
+      panelId: 123,
     } as DataQueryRequest);
 
     const args = mock.calls[0][0];
@@ -110,7 +123,7 @@ describe('DataSourceWithBackend', () => {
             Object {
               "datasource": Object {
                 "type": "sample",
-                "uid": "?",
+                "uid": "<mockuid>",
               },
               "datasourceId": undefined,
               "intervalMs": 5000,
@@ -119,6 +132,12 @@ describe('DataSourceWithBackend', () => {
             },
           ],
         },
+        "headers": Object {
+          "X-Dashboard-Uid": "dashA",
+          "X-Datasource-Uid": "abc, <mockuid>",
+          "X-Panel-Id": "123",
+          "X-Plugin-Id": "dummy, sample",
+        },
         "hideFromInspector": true,
         "method": "POST",
         "requestId": undefined,

packages/grafana-runtime/src/utils/DataSourceWithBackend.ts

@@ -71,6 +71,14 @@ export enum HealthStatus {
   Error = 'ERROR',
 }
 
+// Internal for now
+enum PluginRequestHeaders {
+  PluginID = 'X-Plugin-Id', // can be used for routing
+  DatasourceUID = 'X-Datasource-Uid', // can be used for routing/ load balancing
+  DashboardUID = 'X-Dashboard-Uid', // mainly useful for debuging slow queries
+  PanelID = 'X-Panel-Id', // mainly useful for debuging slow queries
+}
+
 /**
  * Describes the details in the payload returned when checking the health of a data source
  * plugin.
@@ -119,11 +127,15 @@ class DataSourceWithBackend<
       targets = targets.filter((q) => this.filterQuery!(q));
     }
 
+    let hasExpr = false;
+    const pluginIDs = new Set<string>();
+    const dsUIDs = new Set<string>();
     const queries = targets.map((q) => {
       let datasource = this.getRef();
       let datasourceId = this.id;
 
       if (isExpressionReference(q.datasource)) {
+        hasExpr = true;
         return {
           ...q,
           datasource: ExpressionDatasourceRef,
@@ -140,7 +152,12 @@ class DataSourceWithBackend<
         datasource = ds.rawRef ?? getDataSourceRef(ds);
         datasourceId = ds.id;
       }
-
+      if (datasource.type?.length) {
+        pluginIDs.add(datasource.type);
+      }
+      if (datasource.uid?.length) {
+        dsUIDs.add(datasource.uid);
+      }
       return {
         ...this.applyTemplateVariables(q, request.scopedVars),
         datasource,
@@ -170,13 +187,28 @@ class DataSourceWithBackend<
       });
     }
 
+    let url = '/api/ds/query';
+    if (hasExpr) {
+      url += '?expression=true';
+    }
+
+    const headers: Record<string, string> = {};
+    headers[PluginRequestHeaders.PluginID] = Array.from(pluginIDs).join(', ');
+    headers[PluginRequestHeaders.DatasourceUID] = Array.from(dsUIDs).join(', ');
+    if (request.dashboardUID) {
+      headers[PluginRequestHeaders.DashboardUID] = request.dashboardUID;
+    }
+    if (request.panelId) {
+      headers[PluginRequestHeaders.PanelID] = `${request.panelId}`;
+    }
     return getBackendSrv()
       .fetch<BackendDataSourceResponse>({
-        url: '/api/ds/query',
+        url,
         method: 'POST',
         data: body,
         requestId,
         hideFromInspector,
+        headers,
       })
       .pipe(
         switchMap((raw) => {
@@ -193,6 +225,14 @@ class DataSourceWithBackend<
       );
   }
 
+  /** Get request headers with plugin ID+UID set */
+  protected getRequestHeaders(): Record<string, string> {
+    const headers: Record<string, string> = {};
+    headers[PluginRequestHeaders.PluginID] = this.type;
+    headers[PluginRequestHeaders.DatasourceUID] = this.uid;
+    return headers;
+  }
+
   /**
    * Apply template variables for explore
    */
@@ -221,23 +261,43 @@ class DataSourceWithBackend<
   /**
    * Make a GET request to the datasource resource path
    */
-  async getResource(
+  async getResource<T = any>(
     path: string,
     params?: BackendSrvRequest['params'],
     options?: Partial<BackendSrvRequest>
-  ): Promise<any> {
-    return getBackendSrv().get(`/api/datasources/${this.id}/resources/${path}`, params, options?.requestId, options);
+  ): Promise<T> {
+    const headers = this.getRequestHeaders();
+    const result = await lastValueFrom(
+      getBackendSrv().fetch<T>({
+        ...options,
+        method: 'GET',
+        headers: options?.headers ? { ...options.headers, ...headers } : headers,
+        params: params ?? options?.params,
+        url: `/api/datasources/${this.id}/resources/${path}`,
+      })
+    );
+    return result.data;
   }
 
   /**
    * Send a POST request to the datasource resource path
    */
-  async postResource(
+  async postResource<T = any>(
     path: string,
     data?: BackendSrvRequest['data'],
     options?: Partial<BackendSrvRequest>
-  ): Promise<any> {
-    return getBackendSrv().post(`/api/datasources/${this.id}/resources/${path}`, { ...data }, options);
+  ): Promise<T> {
+    const headers = this.getRequestHeaders();
+    const result = await lastValueFrom(
+      getBackendSrv().fetch<T>({
+        ...options,
+        method: 'GET',
+        headers: options?.headers ? { ...options.headers, ...headers } : headers,
+        data: data ?? { ...data },
+        url: `/api/datasources/${this.id}/resources/${path}`,
+      })
+    );
+    return result.data;
   }
 
   /**
@@ -249,6 +309,7 @@ class DataSourceWithBackend<
         method: 'GET',
         url: `/api/datasources/${this.id}/health`,
         showErrorAlert: false,
+        headers: this.getRequestHeaders(),
       })
     )
       .then((v: FetchResponse) => v.data as HealthCheckResult)