Dashboard -Schema V2 Stateless annotations (dsRef independent) (#102949)

* Implement basic behavior of stateless annotation * Fix ds defined getting lost * Fix linter * adjust test to be aligned with datasources not automatically being assigned to the model * Add unit test for annotations --------- Co-authored-by: Haris Rozajac <58232930+harisrozajac@users.noreply.github.com> Co-authored-by: Ivan Ortega Alba <ivanortegaalba@gmail.com>
2025-04-11 11:42:48 +02:00
parent 8ebce76535
commit 00dcf482cf
5 changed files with 203 additions and 36 deletions
@@ -1078,6 +1078,51 @@ describe('DashboardSceneSerializer', () => {
        serializer.initializeDSReferencesMapping({ elements: {} } as DashboardV2Spec);
        expect(serializer.getDSReferencesMapping().panels.size).toBe(0);
      });
+
+      it('should initialize datasource references mapping when annotations dont have datasources', () => {
+        const saveModel: DashboardV2Spec = {
+          ...defaultDashboardV2Spec(),
+          title: 'Dashboard with annotations without datasource',
+          annotations: [
+            {
+              kind: 'AnnotationQuery',
+              spec: {
+                name: 'Annotation 1',
+                query: { kind: 'prometheus', spec: {} },
+                enable: true,
+                hide: false,
+                iconColor: 'red',
+              },
+            },
+          ],
+        };
+
+        serializer.initializeDSReferencesMapping(saveModel);
+
+        const dsReferencesMap = serializer.getDSReferencesMapping();
+
+        // Annotation 1 should have no datasource
+        expect(dsReferencesMap.annotations.has('Annotation 1')).toBe(true);
+      });
+
+      it('should return early if the saveModel is not a V2 dashboard', () => {
+        const v1SaveModel: Dashboard = {
+          title: 'Test Dashboard',
+          uid: 'my-uid',
+          schemaVersion: 30,
+          panels: [
+            { id: 1, title: 'Panel 1', type: 'text' },
+            { id: 2, title: 'Panel 2', type: 'text' },
+          ],
+        };
+        serializer.initializeDSReferencesMapping(v1SaveModel as unknown as DashboardV2Spec);
+        expect(serializer.getDSReferencesMapping()).toEqual({
+          panels: new Map(),
+          variables: new Set(),
+          annotations: new Set(),
+        });
+        expect(serializer.getDSReferencesMapping().panels.size).toBe(0);
+      });
    });

    describe('V1DashboardSerializer', () => {
@@ -2,6 +2,7 @@ import { Dashboard } from '@grafana/schema';
 import { Spec as DashboardV2Spec } from '@grafana/schema/dist/esm/schema/dashboard/v2alpha1/types.spec.gen';
 import { AnnoKeyDashboardSnapshotOriginalUrl } from 'app/features/apiserver/types';
 import { DashboardWithAccessInfo } from 'app/features/dashboard/api/types';
+import { isDashboardV2Spec } from 'app/features/dashboard/api/utils';
 import { SaveDashboardAsOptions } from 'app/features/dashboard/components/SaveDashboard/types';
 import { DASHBOARD_SCHEMA_VERSION } from 'app/features/dashboard/state/DashboardMigrator';
 import {
@@ -235,6 +236,12 @@ export class V2DashboardSerializer
  }

  initializeDSReferencesMapping(saveModel: DashboardV2Spec | undefined) {
+    // The saveModel could be undefined or not a DashboardV2Spec
+    // when dashboardsNewLayout is enabled, saveModel could be v1
+    // in those cases, only when saving we will convert to v2
+    if (saveModel === undefined || (saveModel && !isDashboardV2Spec(saveModel))) {
+      return;
+    }
    // initialize the object
    this.defaultDsReferencesMap = {
      panels: new Map<string, Set<string>>(),
@@ -274,6 +281,15 @@ export class V2DashboardSerializer
        }
      }
    }
+
+    // initialize annotations ds references map
+    if (saveModel?.annotations) {
+      for (const annotation of saveModel.annotations) {
+        if (!annotation.spec.datasource) {
+          this.defaultDsReferencesMap.annotations.add(annotation.spec.name);
+        }
+      }
+    }
  }

  getDSReferencesMapping() {
@@ -35,10 +35,6 @@ exports[`transformSceneToSaveModelSchemaV2 should transform scene to save model
      "kind": "AnnotationQuery",
      "spec": {
        "builtIn": false,
-        "datasource": {
-          "type": "loki",
-          "uid": "Loki",
-        },
        "enable": true,
        "hide": true,
        "iconColor": "green",
@@ -20,6 +20,7 @@ import {
  SceneDataQuery,
  SceneQueryRunner,
  sceneUtils,
+  dataLayers,
 } from '@grafana/scenes';
 import {
  DashboardCursorSync as DashboardCursorSyncV1,
@@ -57,6 +58,7 @@ import {
  transformSceneToSaveModelSchemaV2,
  validateDashboardSchemaV2,
  getDataQueryKind,
+  getAutoAssignedDSRef,
 } from './transformSceneToSaveModelSchemaV2';

 // Mock dependencies
@@ -421,8 +423,8 @@ describe('transformSceneToSaveModelSchemaV2', () => {

    // Check that the annotation layers are correctly transformed
    expect(result.annotations).toHaveLength(3);
-    // check annotation layer 3 with no datasource has the default datasource defined as type
-    expect(result.annotations?.[2].spec.datasource?.type).toBe('loki');
+    // Check annotation layer 3 without initial data source isn't updated with runtime default
+    expect(result.annotations?.[2].spec.datasource?.type).toBe(undefined);
  });

  it('should transform the minimum scene to save model schema v2', () => {
@@ -748,6 +750,58 @@ describe('getElementDatasource', () => {
    expect(result).toBeUndefined();
  });

+  it('should handle annotation datasources correctly', () => {
+    // Use the dataLayers.AnnotationsDataLayer directly
+    const annotationLayer = new dataLayers.AnnotationsDataLayer({
+      key: 'annotation-1',
+      name: 'Test Annotation',
+      isEnabled: true,
+      isHidden: false,
+      query: {
+        name: 'Test Annotation',
+        enable: true,
+        hide: false,
+        iconColor: 'red',
+        datasource: { uid: 'prometheus', type: 'prometheus' },
+      },
+    });
+
+    // Create an annotation query without datasource
+    const annotationWithoutDS = {
+      name: 'No DS Annotation',
+      enable: true,
+      hide: false,
+      iconColor: 'blue',
+    };
+
+    // Mock dsReferencesMapping
+    const dsReferencesMapping = {
+      panels: new Map([['panel-1', new Set(['A'])]]),
+      variables: new Set<string>(),
+      annotations: new Set<string>(['No DS Annotation']),
+    };
+
+    // Test with annotation that has datasource defined
+    const resultWithDS = getElementDatasource(
+      annotationLayer,
+      annotationLayer.state.query,
+      'annotation',
+      undefined,
+      dsReferencesMapping
+    );
+    expect(resultWithDS).toEqual({ uid: 'prometheus', type: 'prometheus' });
+
+    // Test with annotation that has no datasource defined
+    const resultWithoutDS = getElementDatasource(
+      annotationLayer,
+      annotationWithoutDS,
+      'annotation',
+      undefined,
+      dsReferencesMapping
+    );
+    expect(resultWithoutDS).toBeUndefined();
+  });
+
  it('should handle invalid input combinations', () => {
    const vizPanel = new VizPanel({
      key: 'panel-1',
@@ -772,6 +826,24 @@ describe('getElementDatasource', () => {
    // Variable set with query
    expect(getElementDatasource(variableSet, query, 'variable')).toBeUndefined();
  });
+
+  it('should throw error when invalid type is passed to getAutoAssignedDSRef', () => {
+    const vizPanel = new VizPanel({
+      key: 'panel-1',
+      pluginId: 'timeseries',
+    });
+
+    const dsReferencesMapping = {
+      panels: new Map([['panel-1', new Set(['A'])]]),
+      variables: new Set<string>(),
+      annotations: new Set<string>(),
+    };
+
+    expect(() => {
+      // @ts-expect-error - intentionally passing invalid type to test error handling
+      getAutoAssignedDSRef(vizPanel, 'invalid-type', dsReferencesMapping);
+    }).toThrow('Invalid type invalid-type for getAutoAssignedDSRef');
+  });
 });

 function getMinimalSceneState(body: DashboardLayoutManager): Partial<DashboardSceneState> {
@@ -12,7 +12,6 @@ import {
  SceneVariables,
  SceneVariableSet,
  VizPanel,
-  sceneUtils,
 } from '@grafana/scenes';
 import { DataSourceRef } from '@grafana/schema';
 import { sortedDeepCloneWithoutNulls } from 'app/core/utils/object';
@@ -107,7 +106,7 @@ export function transformSceneToSaveModelSchemaV2(scene: DashboardScene, isSnaps
    // EOF elements

    // annotations
-    annotations: getAnnotations(sceneDash),
+    annotations: getAnnotations(sceneDash, dsReferencesMapping),
    // EOF annotations

    // layout
@@ -392,7 +391,7 @@ function getVariables(oldDash: DashboardSceneState, dsReferencesMapping?: DSRefe
  return variables;
 }

-function getAnnotations(state: DashboardSceneState): AnnotationQueryKind[] {
+function getAnnotations(state: DashboardSceneState, dsReferencesMapping?: DSReferencesMapping): AnnotationQueryKind[] {
  const data = state.$data;
  if (!(data instanceof DashboardDataLayerSet)) {
    return [];
@@ -407,7 +406,7 @@ function getAnnotations(state: DashboardSceneState): AnnotationQueryKind[] {
      spec: {
        builtIn: Boolean(layer.state.query.builtIn),
        name: layer.state.query.name,
-        datasource: layer.state.query.datasource || getDefaultDataSourceRef(),
+        datasource: getElementDatasource(layer, layer.state.query, 'annotation', undefined, dsReferencesMapping),
        enable: Boolean(layer.state.isEnabled),
        hide: Boolean(layer.state.isHidden),
        iconColor: layer.state.query.iconColor,
@@ -657,9 +656,9 @@ function validateRowsLayout(layout: unknown) {
  }
 }

-function getAutoAssignedDSRef(
-  element: VizPanel | SceneVariables,
-  type: 'panels' | 'variables',
+export function getAutoAssignedDSRef(
+  element: VizPanel | SceneVariables | dataLayers.AnnotationsDataLayer,
+  type: 'panels' | 'variables' | 'annotations',
  elementMapReferences?: DSReferencesMapping
 ): Set<string> {
  if (!elementMapReferences) {
@@ -670,16 +669,25 @@ function getAutoAssignedDSRef(
    return elementMapReferences.panels.get(elementKey) || new Set();
  }

-  return elementMapReferences.variables;
+  if (type === 'variables') {
+    return elementMapReferences.variables;
+  }
+
+  if (type === 'annotations') {
+    return elementMapReferences.annotations;
+  }
+
+  // if type is not panels, annotations, or variables, throw error
+  throw new Error(`Invalid type ${type} for getAutoAssignedDSRef`);
 }

 /**
 * Determines if a data source reference should be persisted for a query or variable
 */
-export function getPersistedDSFor<T extends SceneDataQuery | QueryVariable>(
+export function getPersistedDSFor<T extends SceneDataQuery | QueryVariable | AnnotationQuery>(
  element: T,
  autoAssignedDsRef: Set<string>,
-  type: 'query' | 'variable',
+  type: 'query' | 'variable' | 'annotation',
  context?: SceneQueryRunner
 ): DataSourceRef | undefined {
  // Get the element identifier - refId for queries, name for variables
@@ -705,6 +713,10 @@ export function getPersistedDSFor<T extends SceneDataQuery | QueryVariable>(
    return element.state.datasource || {};
  }

+  if (type === 'annotation' && 'datasource' in element) {
+    return element.datasource || {};
+  }
+
  return undefined;
 }

@@ -713,32 +725,52 @@ export function getPersistedDSFor<T extends SceneDataQuery | QueryVariable>(
 * @returns refId for queries, name for variables
 * TODO: we will add annotations in the future
 */
-function getElementIdentifier<T extends SceneDataQuery | QueryVariable>(
+function getElementIdentifier<T extends SceneDataQuery | QueryVariable | AnnotationQuery>(
  element: T,
-  type: 'query' | 'variable'
+  type: 'query' | 'variable' | 'annotation'
 ): string {
  // when is type query look for refId
  if (type === 'query') {
    return 'refId' in element ? element.refId : '';
  }
-  // when is type variable look for the name of the variable
-  return 'state' in element && 'name' in element.state ? element.state.name : '';
+
+  if (type === 'variable') {
+    // when is type variable look for the name of the variable
+    return 'state' in element && 'name' in element.state ? element.state.name : '';
+  }
+
+  // when is type annotation look for annotation name
+  if (type === 'annotation') {
+    return 'name' in element ? element.name : '';
+  }
+
+  throw new Error(`Invalid type ${type} for getElementIdentifier`);
 }

-function isVizPanel(element: VizPanel | SceneVariables): element is VizPanel {
+function isVizPanel(element: VizPanel | SceneVariables | dataLayers.AnnotationsDataLayer): element is VizPanel {
  // FIXME: is there another way to do this?
  return 'pluginId' in element.state;
 }

-function isSceneVariables(element: VizPanel | SceneVariables): element is SceneVariables {
+function isSceneVariables(
+  element: VizPanel | SceneVariables | dataLayers.AnnotationsDataLayer
+): element is SceneVariables {
  // Check for properties unique to SceneVariables but not in VizPanel
  return !('pluginId' in element.state) && ('variables' in element.state || 'getValue' in element);
 }

-function isSceneDataQuery(query: SceneDataQuery | QueryVariable): query is SceneDataQuery {
+function isSceneDataQuery(query: SceneDataQuery | QueryVariable | AnnotationQuery): query is SceneDataQuery {
  return 'refId' in query && !('state' in query);
 }

+function isAnnotationQuery(query: SceneDataQuery | QueryVariable | AnnotationQuery): query is AnnotationQuery {
+  return 'datasource' in query && 'name' in query;
+}
+
+function isQueryVariable(query: SceneDataQuery | QueryVariable | AnnotationQuery): query is QueryVariable {
+  return 'state' in query && 'name' in query.state;
+}
+
 /**
 * Get the persisted datasource for a query or variable
 * When a query or variable is created it could not have a datasource set
@@ -747,35 +779,41 @@ function isSceneDataQuery(query: SceneDataQuery | QueryVariable): query is Scene
 *
 */
 export function getElementDatasource(
-  element: VizPanel | SceneVariables,
-  queryElement: SceneDataQuery | QueryVariable,
-  type: 'panel' | 'variable',
+  element: VizPanel | SceneVariables | dataLayers.AnnotationsDataLayer,
+  queryElement: SceneDataQuery | QueryVariable | AnnotationQuery,
+  type: 'panel' | 'variable' | 'annotation',
  queryRunner?: SceneQueryRunner,
  dsReferencesMapping?: DSReferencesMapping
 ): DataSourceRef | undefined {
+  let result: DataSourceRef | undefined;
  if (type === 'panel') {
    if (!queryRunner || !isVizPanel(element) || !isSceneDataQuery(queryElement)) {
      return undefined;
    }
    // Get datasource for panel query
    const autoAssignedRefs = getAutoAssignedDSRef(element, 'panels', dsReferencesMapping);
-    return getPersistedDSFor(queryElement, autoAssignedRefs, 'query', queryRunner);
+    result = getPersistedDSFor(queryElement, autoAssignedRefs, 'query', queryRunner);
  }

  if (type === 'variable') {
-    if (!isSceneVariables(element) || isSceneDataQuery(queryElement)) {
+    if (!isSceneVariables(element) || !isQueryVariable(queryElement)) {
      return undefined;
    }
    // Get datasource for variable
-    if (!sceneUtils.isQueryVariable(queryElement)) {
-      return undefined;
-    }
    const autoAssignedRefs = getAutoAssignedDSRef(element, 'variables', dsReferencesMapping);
-    // Important: Only return the datasource if it's not in auto-assigned refs
-    // and if the result would not be an empty object
-    const result = getPersistedDSFor(queryElement, autoAssignedRefs, 'variable');
-    return Object.keys(result || {}).length > 0 ? result : undefined;
+
+    result = getPersistedDSFor(queryElement, autoAssignedRefs, 'variable');
  }

-  return undefined;
+  if (type === 'annotation') {
+    if (!isAnnotationQuery(queryElement)) {
+      return undefined;
+    }
+    // Get datasource for annotation
+    const autoAssignedRefs = getAutoAssignedDSRef(element, 'annotations', dsReferencesMapping);
+    result = getPersistedDSFor(queryElement, autoAssignedRefs, 'annotation');
+  }
+  // Important: Only return the datasource if it's not in auto-assigned refs
+  // and if the result would not be an empty object
+  return Object.keys(result || {}).length > 0 ? result : undefined;
 }