diff --git a/public/app/plugins/datasource/mixed/MixedDataSource.test.ts b/public/app/plugins/datasource/mixed/MixedDataSource.test.ts index 5fdec6cd058..f4d2554f895 100644 --- a/public/app/plugins/datasource/mixed/MixedDataSource.test.ts +++ b/public/app/plugins/datasource/mixed/MixedDataSource.test.ts @@ -2,7 +2,11 @@ import { lastValueFrom } from 'rxjs'; import { getQueryOptions } from 'test/helpers/getQueryOptions'; import { DatasourceSrvMock, MockObservableDataSourceApi } from 'test/mocks/datasource_srv'; -import { DataQueryRequest, DataSourceInstanceSettings, LoadingState } from '@grafana/data'; +import { DataQueryRequest, DataSourceInstanceSettings, DataSourceRef, LoadingState } from '@grafana/data'; +import { DataSourceSrv, setDataSourceSrv, setTemplateSrv } from '@grafana/runtime'; +import { CustomVariable, SceneFlexLayout, SceneVariableSet } from '@grafana/scenes'; + +import { TemplateSrv } from '../../../features/templating/template_srv'; import { MIXED_DATASOURCE_NAME } from './MixedDataSource'; import { MixedDatasource } from './module'; @@ -21,13 +25,19 @@ const datasourceSrv = new DatasourceSrvMock(defaultDS, { ]), }); -const getDataSourceSrvMock = jest.fn().mockReturnValue(datasourceSrv); -jest.mock('@grafana/runtime', () => ({ - ...jest.requireActual('@grafana/runtime'), - getDataSourceSrv: () => getDataSourceSrvMock(), -})); - describe('MixedDatasource', () => { + beforeEach(() => { + jest.clearAllMocks(); + setDataSourceSrv({ + ...datasourceSrv, + get: (uid: DataSourceRef) => datasourceSrv.get(uid), + getInstanceSettings: jest.fn().mockReturnValue({ meta: {} }), + getList: jest.fn(), + reload: jest.fn(), + }); + setTemplateSrv(new TemplateSrv()); + }); + describe('with no errors', () => { it('direct query should return results', async () => { const ds = new MixedDatasource({} as DataSourceInstanceSettings); @@ -83,6 +93,100 @@ describe('MixedDatasource', () => { }); }); + describe('with multi template variable', () => { + beforeAll(() => { + setDataSourceSrv({ + getInstanceSettings() { + return {}; + }, + } as DataSourceSrv); + }); + + const scene = new SceneFlexLayout({ + children: [], + $variables: new SceneVariableSet({ + variables: [new CustomVariable({ name: 'ds', value: ['B', 'C'] })], + }), + }); + + it('should run query for each datasource when there is a multi value template variable', async () => { + const ds = new MixedDatasource({} as DataSourceInstanceSettings); + + const request = { + targets: [{ refId: 'AA', datasource: { uid: '$ds' } }], + scopedVars: { + __sceneObject: { value: scene }, + }, + } as unknown as DataQueryRequest; + + await expect(ds.query(request)).toEmitValuesWith((results) => { + expect(results).toHaveLength(2); + expect(results[0].key).toBe('mixed-0-'); + expect(results[0].state).toBe(LoadingState.Loading); + expect(results[1].key).toBe('mixed-1-'); + expect(results[1].state).toBe(LoadingState.Done); + }); + }); + + it('should run query for picked datasource and template variable datasource', async () => { + const ds = new MixedDatasource({} as DataSourceInstanceSettings); + const request = { + targets: [ + { refId: 'AA', datasource: { uid: '$ds' } }, + { refId: 'BB', datasource: { uid: 'Loki' } }, + ], + scopedVars: { + __sceneObject: { value: scene }, + }, + } as unknown as DataQueryRequest; + + await expect(ds.query(request)).toEmitValuesWith((results) => { + expect(results).toHaveLength(4); + expect(results[0].key).toBe('mixed-0-'); + expect(results[0].state).toBe(LoadingState.Loading); + expect(results[1].key).toBe('mixed-1-'); + expect(results[1].state).toBe(LoadingState.Loading); + expect(results[2].key).toBe('mixed-2-A'); + expect(results[2].state).toBe(LoadingState.Loading); + expect(results[3].key).toBe('mixed-2-B'); + expect(results[3].state).toBe(LoadingState.Done); + }); + }); + }); + + describe('with single value template variable', () => { + beforeAll(() => { + setDataSourceSrv({ + getInstanceSettings() { + return {}; + }, + } as DataSourceSrv); + }); + + const scene = new SceneFlexLayout({ + children: [], + $variables: new SceneVariableSet({ + variables: [new CustomVariable({ name: 'ds', value: 'B' })], + }), + }); + + it('should run query for correct datasource', async () => { + const ds = new MixedDatasource({} as DataSourceInstanceSettings); + + const request = { + targets: [{ refId: 'AA', datasource: { uid: '$ds' } }], + scopedVars: { + __sceneObject: { value: scene }, + }, + } as unknown as DataQueryRequest; + + await expect(ds.query(request)).toEmitValuesWith((results) => { + expect(results).toHaveLength(1); + expect(results[0].data).toEqual(['BBBB']); + }); + }); + }); + it('should return both query results from the same data source', async () => { const ds = new MixedDatasource({} as DataSourceInstanceSettings); const request = { diff --git a/public/app/plugins/datasource/mixed/MixedDataSource.ts b/public/app/plugins/datasource/mixed/MixedDataSource.ts index 2466cdcbda9..8bcb35eca40 100644 --- a/public/app/plugins/datasource/mixed/MixedDataSource.ts +++ b/public/app/plugins/datasource/mixed/MixedDataSource.ts @@ -10,8 +10,10 @@ import { DataSourceApi, DataSourceInstanceSettings, LoadingState, + ScopedVars, } from '@grafana/data'; -import { getDataSourceSrv, toDataQueryError } from '@grafana/runtime'; +import { getDataSourceSrv, getTemplateSrv, toDataQueryError } from '@grafana/runtime'; +import { CustomFormatterVariable } from '@grafana/scenes'; export const MIXED_DATASOURCE_NAME = '-- Mixed --'; @@ -19,7 +21,8 @@ export const mixedRequestId = (queryIdx: number, requestId?: string) => `mixed-$ export interface BatchedQueries { datasource: Promise; - targets: DataQuery[]; + queries: DataQuery[]; + scopedVars: ScopedVars; } export class MixedDatasource extends DataSourceApi { @@ -39,23 +42,71 @@ export class MixedDatasource extends DataSourceApi { // Build groups of queries to run in parallel const sets: { [key: string]: DataQuery[] } = groupBy(queries, 'datasource.uid'); - const mixed: BatchedQueries[] = []; + const batches: BatchedQueries[] = []; for (const key in sets) { - const targets = sets[key]; - - mixed.push({ - datasource: getDataSourceSrv().get(targets[0].datasource, request.scopedVars), - targets, - }); + batches.push(...this.getBatchesForQueries(sets[key], request)); } // Missing UIDs? - if (!mixed.length) { + if (!batches.length) { return of({ data: [] }); // nothing } - return this.batchQueries(mixed, request); + return this.batchQueries(batches, request); + } + + /** + * Almost always returns a single batch for each set of queries. + * Unless the query is using a multi value variable. + */ + private getBatchesForQueries(queries: DataQuery[], request: DataQueryRequest) { + const dsRef = queries[0].datasource; + const batches: BatchedQueries[] = []; + + // Using the templateSrv.replace function here with a custom formatter as that is the cleanest way + // to access the raw value or value array of a variable. + const datasourceUid = getTemplateSrv().replace( + dsRef?.uid, + request.scopedVars, + (value: string | string[], variable: CustomFormatterVariable) => { + // If it's not a data source variable, or single value + if (!Array.isArray(value)) { + return value; + } + + for (const uid of value) { + if (uid === 'default') { + continue; + } + + const dsSettings = getDataSourceSrv().getInstanceSettings(uid); + + batches.push({ + datasource: getDataSourceSrv().get(uid), + queries: cloneDeep(queries), + scopedVars: { + ...request.scopedVars, + [variable.name]: { value: uid, text: dsSettings?.name }, + }, + }); + } + + return ''; + } + ); + + if (datasourceUid !== '') { + batches.push({ + datasource: getDataSourceSrv().get(datasourceUid), + queries: cloneDeep(queries), + scopedVars: { + ...request.scopedVars, + }, + }); + } + + return batches; } batchQueries(mixed: BatchedQueries[], request: DataQueryRequest): Observable { @@ -64,7 +115,8 @@ export class MixedDatasource extends DataSourceApi { mergeMap((api: DataSourceApi) => { const dsRequest = cloneDeep(request); dsRequest.requestId = mixedRequestId(i, dsRequest.requestId); - dsRequest.targets = query.targets; + dsRequest.targets = query.queries; + dsRequest.scopedVars = query.scopedVars; return from(api.query(dsRequest)).pipe( map((response) => { @@ -102,7 +154,7 @@ export class MixedDatasource extends DataSourceApi { } private isQueryable(query: BatchedQueries): boolean { - return query && Array.isArray(query.targets) && query.targets.length > 0; + return query && Array.isArray(query.queries) && query.queries.length > 0; } private finalizeResponses(responses: DataQueryResponse[]): DataQueryResponse[] {