From f1a2a768971e64f98697c38da79bf1810ec2ac84 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ida=20=C5=A0tambuk?= <ida.stambuk@grafana.com>
Date: Mon, 30 Jan 2023 16:45:03 +0100
Subject: [PATCH] Datasources: Use getDefaultQuery in annotations editors
 (#61870)

+ Add Cloudwatch default annotation
---
 .../grafana-data/src/types/annotations.ts     |  5 ++
 .../StandardAnnotationQueryEditor.test.tsx    | 70 +++++++++++++++++++
 .../StandardAnnotationQueryEditor.tsx         |  8 ++-
 .../annotations/executeAnnotationQuery.ts     |  6 +-
 public/app/features/query/state/runRequest.ts | 10 +--
 .../cloudwatch/annotationSupport.ts           |  8 ++-
 .../datasource/cloudwatch/datasource.test.ts  |  2 +-
 .../datasource/cloudwatch/defaultQueries.ts   |  8 +++
 8 files changed, 106 insertions(+), 11 deletions(-)
 create mode 100644 public/app/features/annotations/components/StandardAnnotationQueryEditor.test.tsx

diff --git a/packages/grafana-data/src/types/annotations.ts b/packages/grafana-data/src/types/annotations.ts
index 099ad4fe034..da077daa76e 100644
--- a/packages/grafana-data/src/types/annotations.ts
+++ b/packages/grafana-data/src/types/annotations.ts
@@ -113,4 +113,9 @@ export interface AnnotationSupport<TQuery extends DataQuery = DataQuery, TAnno =
    * Specify a custom QueryEditor for the annotation page.  If not specified, the standard one will be used
    */
   QueryEditor?: ComponentType<AnnotationQueryEditorProps<TQuery>>;
+
+  /**
+   * Define this method if you want to pre-populate the editor with a default query
+   */
+  getDefaultQuery?(): Partial<TQuery>;
 }
diff --git a/public/app/features/annotations/components/StandardAnnotationQueryEditor.test.tsx b/public/app/features/annotations/components/StandardAnnotationQueryEditor.test.tsx
new file mode 100644
index 00000000000..a840ae1641c
--- /dev/null
+++ b/public/app/features/annotations/components/StandardAnnotationQueryEditor.test.tsx
@@ -0,0 +1,70 @@
+import { render } from '@testing-library/react';
+import React from 'react';
+
+import { AnnotationQuery, DataSourceApi, DataSourceInstanceSettings } from '@grafana/data/src';
+
+import StandardAnnotationQueryEditor, { Props as EditorProps } from './StandardAnnotationQueryEditor';
+
+const setup = (customProps: Partial<EditorProps>) => {
+  const props: EditorProps = {
+    datasource: {} as unknown as DataSourceApi,
+    datasourceInstanceSettings: {} as DataSourceInstanceSettings,
+    annotation: {} as AnnotationQuery,
+    onChange: jest.fn(),
+    ...customProps,
+  };
+  const { rerender } = render(<StandardAnnotationQueryEditor {...props} />);
+  return { rerender, props };
+};
+
+jest.mock('app/features/dashboard/services/DashboardSrv', () => ({
+  getDashboardSrv: jest.fn().mockReturnValue({
+    getCurrent: jest.fn().mockReturnValue(null),
+  }),
+}));
+
+jest.mock('app/features/dashboard/services/TimeSrv', () => ({
+  getTimeSrv: jest.fn().mockReturnValue({
+    timeRange: jest.fn().mockReturnValue({}),
+  }),
+}));
+
+describe('StandardAnnotationQueryEditor', () => {
+  it('should fill out a default query if it is defined and pass it to the Query Editor', () => {
+    const { props } = setup({
+      annotation: { name: 'initialAnn', target: { refId: 'initialAnnotationRef' } } as AnnotationQuery,
+
+      datasource: {
+        annotations: {
+          QueryEditor: jest.fn(() => <div>Editor</div>),
+          getDefaultQuery: jest.fn().mockImplementation(() => ({ queryType: 'defaultAnnotationsQuery' })),
+          prepareAnnotation: (annotation: AnnotationQuery) => annotation,
+        },
+      } as unknown as DataSourceApi,
+    });
+    expect(props.datasource?.annotations?.getDefaultQuery).toBeDefined();
+    expect(props.datasource?.annotations?.QueryEditor).toHaveBeenCalledWith(
+      expect.objectContaining({
+        query: expect.objectContaining({ queryType: 'defaultAnnotationsQuery', refId: 'initialAnnotationRef' }),
+      }),
+      expect.anything()
+    );
+  });
+  it('should keep and pass the initial query if the defaultQuery is not defined', () => {
+    const { props } = setup({
+      annotation: { name: 'initialAnn', target: { refId: 'initialAnnotationRef' } } as AnnotationQuery,
+      datasource: {
+        annotations: {
+          QueryEditor: jest.fn(() => <div>Editor</div>),
+          prepareAnnotation: (annotation: AnnotationQuery) => annotation,
+        },
+      } as unknown as DataSourceApi,
+    });
+    expect(props.datasource?.annotations?.QueryEditor).toHaveBeenCalledWith(
+      expect.objectContaining({
+        query: expect.objectContaining({ refId: 'initialAnnotationRef' }),
+      }),
+      expect.anything()
+    );
+  });
+});
diff --git a/public/app/features/annotations/components/StandardAnnotationQueryEditor.tsx b/public/app/features/annotations/components/StandardAnnotationQueryEditor.tsx
index 61d335c0642..e6df175dc2e 100644
--- a/public/app/features/annotations/components/StandardAnnotationQueryEditor.tsx
+++ b/public/app/features/annotations/components/StandardAnnotationQueryEditor.tsx
@@ -22,7 +22,7 @@ import { AnnotationQueryResponse } from '../types';
 
 import { AnnotationFieldMapper } from './AnnotationResultMapper';
 
-interface Props {
+export interface Props {
   datasource: DataSourceApi;
   datasourceInstanceSettings: DataSourceInstanceSettings;
   annotation: AnnotationQuery<DataQuery>;
@@ -186,7 +186,11 @@ export default class StandardAnnotationQueryEditor extends PureComponent<Props,
       return <div>Annotations are not supported. This datasource needs to export a QueryEditor</div>;
     }
 
-    const query = annotation.target ?? { refId: 'Anno' };
+    const query = {
+      ...datasource.annotations?.getDefaultQuery?.(),
+      ...(annotation.target ?? { refId: 'Anno' }),
+    };
+
     return (
       <>
         <DataSourcePluginContextProvider instanceSettings={datasourceInstanceSettings}>
diff --git a/public/app/features/annotations/executeAnnotationQuery.ts b/public/app/features/annotations/executeAnnotationQuery.ts
index df9a7e0365c..e3870a75437 100644
--- a/public/app/features/annotations/executeAnnotationQuery.ts
+++ b/public/app/features/annotations/executeAnnotationQuery.ts
@@ -23,7 +23,11 @@ export function executeAnnotationQuery(
     ...datasource.annotations,
   };
 
-  const annotation = processor.prepareAnnotation!(savedJsonAnno);
+  const annotationWithDefaults = {
+    ...processor.getDefaultQuery?.(),
+    ...savedJsonAnno,
+  };
+  const annotation = processor.prepareAnnotation!(annotationWithDefaults);
   if (!annotation) {
     return of({});
   }
diff --git a/public/app/features/query/state/runRequest.ts b/public/app/features/query/state/runRequest.ts
index 440bf4f4e6c..fe0a7d120cf 100644
--- a/public/app/features/query/state/runRequest.ts
+++ b/public/app/features/query/state/runRequest.ts
@@ -24,7 +24,6 @@ import {
 import { toDataQueryError } from '@grafana/runtime';
 import { isExpressionReference } from '@grafana/runtime/src/utils/DataSourceWithBackend';
 import { backendSrv } from 'app/core/services/backend_srv';
-import { queryIsEmpty } from 'app/core/utils/query';
 import { dataSource as expressionDatasource } from 'app/features/expressions/ExpressionDatasource';
 import { ExpressionQuery } from 'app/features/expressions/types';
 
@@ -176,10 +175,11 @@ export function callQueryMethod(
   request: DataQueryRequest,
   queryFunction?: typeof datasource.query
 ) {
-  // If the datasource has defined a default query, make sure it's applied if the query is empty
-  request.targets = request.targets.map((t) =>
-    queryIsEmpty(t) ? { ...datasource?.getDefaultQuery?.(CoreApp.PanelEditor), ...t } : t
-  );
+  // If the datasource has defined a default query, make sure it's applied
+  request.targets = request.targets.map((t) => ({
+    ...datasource?.getDefaultQuery?.(CoreApp.PanelEditor),
+    ...t,
+  }));
 
   // If its a public datasource, just return the result. Expressions will be handled on the backend.
   if (datasource.type === 'public-ds') {
diff --git a/public/app/plugins/datasource/cloudwatch/annotationSupport.ts b/public/app/plugins/datasource/cloudwatch/annotationSupport.ts
index 88a6c3fe9c7..08d4659a707 100644
--- a/public/app/plugins/datasource/cloudwatch/annotationSupport.ts
+++ b/public/app/plugins/datasource/cloudwatch/annotationSupport.ts
@@ -1,6 +1,7 @@
 import { AnnotationQuery } from '@grafana/data';
 
 import { AnnotationQueryEditor } from './components/AnnotationQueryEditor';
+import { DEFAULT_ANNOTATIONS_QUERY } from './defaultQueries';
 import { isCloudWatchAnnotation } from './guards';
 import { CloudWatchAnnotationQuery, CloudWatchQuery, LegacyAnnotationQuery } from './types';
 
@@ -24,8 +25,8 @@ export const CloudWatchAnnotationSupport = {
       target: {
         ...query.target,
         ...query,
-        statistic: query.statistic || 'Average',
-        region: query.region || 'default',
+        statistic: query.statistic || DEFAULT_ANNOTATIONS_QUERY.statistic,
+        region: query.region || DEFAULT_ANNOTATIONS_QUERY.region,
         queryMode: 'Annotations',
         refId: query.refId || 'annotationQuery',
       },
@@ -56,5 +57,8 @@ export const CloudWatchAnnotationSupport = {
 
     return undefined;
   },
+  getDefaultQuery() {
+    return DEFAULT_ANNOTATIONS_QUERY;
+  },
   QueryEditor: AnnotationQueryEditor,
 };
diff --git a/public/app/plugins/datasource/cloudwatch/datasource.test.ts b/public/app/plugins/datasource/cloudwatch/datasource.test.ts
index 92bf5f4f1e8..258023ac047 100644
--- a/public/app/plugins/datasource/cloudwatch/datasource.test.ts
+++ b/public/app/plugins/datasource/cloudwatch/datasource.test.ts
@@ -14,10 +14,10 @@ import { setupForLogs } from './__mocks__/logsTestContext';
 import { validLogsQuery, validMetricSearchBuilderQuery } from './__mocks__/queries';
 import { TimeRangeMock } from './__mocks__/timeRange';
 import {
+  CloudWatchDefaultQuery,
   CloudWatchLogsQuery,
   CloudWatchMetricsQuery,
   CloudWatchQuery,
-  CloudWatchDefaultQuery,
   MetricEditorMode,
   MetricQueryType,
 } from './types';
diff --git a/public/app/plugins/datasource/cloudwatch/defaultQueries.ts b/public/app/plugins/datasource/cloudwatch/defaultQueries.ts
index a3e804861ff..62b72054e99 100644
--- a/public/app/plugins/datasource/cloudwatch/defaultQueries.ts
+++ b/public/app/plugins/datasource/cloudwatch/defaultQueries.ts
@@ -1,4 +1,5 @@
 import {
+  CloudWatchAnnotationQuery,
   CloudWatchLogsQuery,
   CloudWatchMetricsQuery,
   LogGroup,
@@ -24,6 +25,13 @@ export const DEFAULT_METRICS_QUERY: Omit<CloudWatchMetricsQuery, 'refId'> = {
   matchExact: true,
 };
 
+export const DEFAULT_ANNOTATIONS_QUERY: Omit<CloudWatchAnnotationQuery, 'refId'> = {
+  queryMode: 'Annotations',
+  namespace: '',
+  region: 'default',
+  statistic: 'Average',
+};
+
 export const DEFAULT_LOGS_QUERY_STRING = 'fields @timestamp, @message |\n sort @timestamp desc |\n limit 20';
 
 export const getDefaultLogsQuery = (