SQR: Fix datasource-template runtime routing for multi-select fan-out

packages/scenes/src/querying/SceneQueryRunner.test.ts

@@ -47,62 +47,68 @@ import { SafeSerializableSceneObject } from '../utils/SafeSerializableSceneObjec
 import { SceneQueryStateControllerState } from '../behaviors/types';
 import { config } from '@grafana/runtime';
 import { ScopesVariable } from '../variables/variants/ScopesVariable';
+import { ConstantVariable } from '../variables/variants/ConstantVariable';
+
+const getDataSourceMock = jest.fn().mockImplementation((datasource) => {
+  const isMixed = datasource?.uid === '-- Mixed --';
+
+  return {
+    uid: isMixed ? '-- Mixed --' : 'test-uid',
+    meta: isMixed ? { mixed: true } : undefined,
+    getRef: () => (isMixed ? { type: 'mixed', uid: '-- Mixed --' } : { uid: 'test-uid' }),
+    query: (request: DataQueryRequest) => {
+      if (request.targets.find((t) => t.refId === 'withAnnotations')) {
+        return of({
+          data: [
+            toDataFrame({
+              refId: 'withAnnotations',
+              datapoints: [
+                [100, 1],
+                [400, 2],
+                [500, 3],
+              ],
+            }),
+            toDataFrame({
+              name: 'exemplar',
+              refId: 'withAnnotations',
+              meta: {
+                typeVersion: [0, 0],
+                custom: {
+                  resultType: 'exemplar',
+                },
+                dataTopic: 'annotations',
+              },
+              fields: [
+                {
+                  name: 'foo',
+                  type: 'string',
+                  values: ['foo1', 'foo2', 'foo3'],
+                },
+                {
+                  name: 'bar',
+                  type: 'string',
+                  values: ['bar1', 'bar2', 'bar3'],
+                },
+              ],
+            }),
+          ],
+        });
+      }
 
-const getDataSourceMock = jest.fn().mockReturnValue({
-  uid: 'test-uid',
-  getRef: () => ({ uid: 'test-uid' }),
-  query: (request: DataQueryRequest) => {
-    if (request.targets.find((t) => t.refId === 'withAnnotations')) {
       return of({
         data: [
           toDataFrame({
-            refId: 'withAnnotations',
+            refId: 'A',
             datapoints: [
               [100, 1],
-              [400, 2],
-              [500, 3],
-            ],
-          }),
-          toDataFrame({
-            name: 'exemplar',
-            refId: 'withAnnotations',
-            meta: {
-              typeVersion: [0, 0],
-              custom: {
-                resultType: 'exemplar',
-              },
-              dataTopic: 'annotations',
-            },
-            fields: [
-              {
-                name: 'foo',
-                type: 'string',
-                values: ['foo1', 'foo2', 'foo3'],
-              },
-              {
-                name: 'bar',
-                type: 'string',
-                values: ['bar1', 'bar2', 'bar3'],
-              },
+              [200, 2],
+              [300, 3],
             ],
           }),
         ],
       });
-    }
-
-    return of({
-      data: [
-        toDataFrame({
-          refId: 'A',
-          datapoints: [
-            [100, 1],
-            [200, 2],
-            [300, 3],
-          ],
-        }),
-      ],
-    });
-  },
+    },
+  };
 });
 
 const runRequestMock = jest.fn().mockImplementation((ds: DataSourceApi, request: DataQueryRequest) => {
@@ -1461,6 +1467,51 @@ describe.each(['11.1.2', '11.1.1'])('SceneQueryRunner', (v) => {
       expect(getDataSourceCall[0]).toEqual({ uid: 'Muuu' });
     });
 
+    it('should keep non-Mixed runtime datasource for templated datasource variable with single selected value', async () => {
+      const queryRunner = new SceneQueryRunner({
+        queries: [{ refId: 'A', datasource: { uid: '${ds}', type: 'prometheus' } }],
+      });
+
+      const scene = new SceneFlexLayout({
+        $variables: new SceneVariableSet({
+          variables: [new ConstantVariable({ name: 'ds', value: 'uid-1' })],
+        }),
+        $timeRange: new SceneTimeRange(),
+        $data: queryRunner,
+        children: [],
+      });
+
+      scene.activate();
+      await new Promise((r) => setTimeout(r, 1));
+
+      expect(getDataSourceMock.mock.calls[0][0]).toEqual({ uid: '${ds}', type: 'prometheus' });
+      expect(sentRequest?.targets[0].datasource).toEqual({ uid: 'test-uid' });
+      expect(queryRunner.state.datasource).toBeUndefined();
+    });
+
+    it('should resolve runtime datasource to Mixed when a datasource variable has multiple selected values', async () => {
+      const queryRunner = new SceneQueryRunner({
+        datasource: { uid: '${ds}', type: 'prometheus' },
+        queries: [{ refId: 'A' }],
+      });
+
+      const scene = new SceneFlexLayout({
+        $variables: new SceneVariableSet({
+          variables: [new ConstantVariable({ name: 'ds', value: ['uid-1', 'uid-2'] })],
+        }),
+        $timeRange: new SceneTimeRange(),
+        $data: queryRunner,
+        children: [],
+      });
+
+      scene.activate();
+      await new Promise((r) => setTimeout(r, 1));
+
+      expect(getDataSourceMock.mock.calls[0][0]).toEqual({ type: 'mixed', uid: '-- Mixed --' });
+      expect(sentRequest?.targets[0].datasource).toEqual({ uid: '${ds}', type: 'prometheus' });
+      expect(queryRunner.state.datasource).toEqual({ uid: '${ds}', type: 'prometheus' });
+    });
+
     it('Should interpolate a variable when used in query options', async () => {
       const variable = new TestVariable({ name: 'A', value: '1h' });
       const queryRunner = new SceneQueryRunner({

packages/scenes/src/querying/SceneQueryRunner.ts

@@ -51,6 +51,8 @@ import { wrapInSafeSerializableSceneObject } from '../utils/wrapInSafeSerializab
 import { DrilldownDependenciesManager } from '../variables/DrilldownDependenciesManager';
 
 let counter = 100;
+const MIXED_DATASOURCE_UID = '-- Mixed --';
+const MIXED_DATASOURCE_REF: DataSourceRef = { type: 'mixed', uid: MIXED_DATASOURCE_UID };
 
 export function getNextRequestId(prefix = 'SQR') {
   return prefix + counter++;
@@ -82,6 +84,83 @@ export interface QueryRunnerState extends SceneObjectState {
   _hasFetchedData?: boolean;
 }
 
+function isTemplatedDatasourceUid(uid: string | undefined): boolean {
+  if (!uid) {
+    return false;
+  }
+
+  // Detect any variable template usage in a datasource UID so we can evaluate runtime fan-out.
+  return /^(?:\$[A-Za-z0-9_]+|\$\{[^}]+\})/.test(uid);
+}
+
+function isTemplatedDatasourceRef(datasource: DataSourceRef | null | undefined): boolean {
+  return isTemplatedDatasourceUid(datasource?.uid);
+}
+
+function getSimpleVariableNameFromDatasourceUid(uid: string | undefined): string | undefined {
+  if (!uid) {
+    return undefined;
+  }
+  const shortMatch = uid.match(/^\$([A-Za-z0-9_]+)$/);
+  if (shortMatch) {
+    return shortMatch[1];
+  }
+
+  const bracketMatch = uid.match(/^\$\{([A-Za-z0-9_]+)\}$/);
+  if (bracketMatch) {
+    return bracketMatch[1];
+  }
+
+  return undefined;
+}
+
+function getDatasourceVariableSelectionCount(sceneObject: SceneQueryRunner, variableName: string): number | undefined {
+  const variable = sceneGraph.lookupVariable(variableName, sceneObject);
+  if (!variable) {
+    return undefined;
+  }
+
+  const value = variable.getValue();
+  if (Array.isArray(value)) {
+    return value.filter((item) => item !== 'default').length;
+  }
+
+  if (value === null || value === undefined) {
+    return undefined;
+  }
+
+  return value === 'default' ? 0 : 1;
+}
+
+function shouldUseMixedRuntimeDatasource(
+  sceneObject: SceneQueryRunner,
+  datasource: DataSourceRef | null | undefined,
+  queries: SceneDataQuery[]
+) {
+  // Check both panel datasource and per-query datasources.
+  const refs: Array<DataSourceRef | null | undefined> = [datasource, ...queries.map((query) => query.datasource)];
+
+  for (const ref of refs) {
+    const uid = ref?.uid;
+    if (!uid) {
+      continue;
+    }
+
+    const variableName = getSimpleVariableNameFromDatasourceUid(uid);
+    if (variableName) {
+      const selectionCount = getDatasourceVariableSelectionCount(sceneObject, variableName);
+
+      if (selectionCount !== undefined && selectionCount > 1) {
+        return true;
+      }
+
+      continue;
+    }
+  }
+
+  return false;
+}
+
 export interface DataQueryExtended extends DataQuery {
   [key: string]: any;
 
@@ -465,7 +544,10 @@ export class SceneQueryRunner extends SceneObjectBase<QueryRunnerState> implemen
 
     try {
       const datasource = this.state.datasource ?? findFirstDatasource(queries);
-      const ds = await getDataSource(datasource, this._scopedVars);
+      const runtimeDatasource = shouldUseMixedRuntimeDatasource(this, datasource, queries)
+        ? MIXED_DATASOURCE_REF
+        : datasource;
+      const ds = await getDataSource(runtimeDatasource, this._scopedVars);
 
       this._drilldownDependenciesManager.findAndSubscribeToDrilldowns(ds.uid, this);
 
@@ -577,7 +659,11 @@ export class SceneQueryRunner extends SceneObjectBase<QueryRunnerState> implemen
           isExpressionReference /* TODO: Remove this check when isExpressionReference is properly exported from grafan runtime */ &&
           !isExpressionReference(query.datasource))
       ) {
-        query.datasource = ds.getRef();
+        if (!query.datasource && ds.meta?.mixed && isTemplatedDatasourceRef(this.state.datasource)) {
+          query.datasource = this.state.datasource;
+        } else {
+          query.datasource = ds.getRef();
+        }
       }
       return query;
     });

packages/scenes/src/querying/__snapshots__/SceneQueryRunner.test.ts.snap

@@ -57,7 +57,7 @@ exports[`SceneQueryRunner when running query should build DataQueryRequest objec
     "from": "now-6h",
     "to": "now",
   },
-  "requestId": "SQR192",
+  "requestId": "SQR194",
   "scopes": undefined,
   "startTime": 1689063488000,
   "targets": [