From cff473f825825a419eb544d56960ce3a8a541592 Mon Sep 17 00:00:00 2001
From: Beto Dealmeida <roberto@dealmeida.net>
Date: Thu, 30 Nov 2023 11:13:12 -0500
Subject: [PATCH] fix(annotations): time grain column (#26140)

---
 .../src/components/Chart/chartAction.js       |  9 ++-
 .../src/components/Chart/chartActions.test.js | 67 +++++++++++++++++++
 2 files changed, 73 insertions(+), 3 deletions(-)

diff --git a/superset-frontend/src/components/Chart/chartAction.js b/superset-frontend/src/components/Chart/chartAction.js
index 42aa3fc5b..8cd3785ae 100644
--- a/superset-frontend/src/components/Chart/chartAction.js
+++ b/superset-frontend/src/components/Chart/chartAction.js
@@ -269,9 +269,12 @@ export function runAnnotationQuery({
       return Promise.resolve();
     }
 
-    const granularity = fd.time_grain_sqla || fd.granularity;
-    fd.time_grain_sqla = granularity;
-    fd.granularity = granularity;
+    // In the original formData the `granularity` attribute represents the time grain (eg
+    // `P1D`), but in the request payload it corresponds to the name of the column where
+    // the time grain should be applied (eg, `Date`), so we need to move things around.
+    fd.time_grain_sqla = fd.time_grain_sqla || fd.granularity;
+    fd.granularity = fd.granularity_sqla;
+
     const overridesKeys = Object.keys(annotation.overrides);
     if (overridesKeys.includes('since') || overridesKeys.includes('until')) {
       annotation.overrides = {
diff --git a/superset-frontend/src/components/Chart/chartActions.test.js b/superset-frontend/src/components/Chart/chartActions.test.js
index b44ca7c8d..b3a6fed9f 100644
--- a/superset-frontend/src/components/Chart/chartActions.test.js
+++ b/superset-frontend/src/components/Chart/chartActions.test.js
@@ -21,6 +21,7 @@ import fetchMock from 'fetch-mock';
 import sinon from 'sinon';
 
 import * as chartlib from '@superset-ui/core';
+import { SupersetClient } from '@superset-ui/core';
 import { LOG_EVENT } from 'src/logger/actions';
 import * as exploreUtils from 'src/explore/exploreUtils';
 import * as actions from 'src/components/Chart/chartAction';
@@ -233,4 +234,70 @@ describe('chart actions', () => {
       expect(json.result[0].value.toString()).toEqual(expectedBigNumber);
     });
   });
+
+  describe('runAnnotationQuery', () => {
+    const mockDispatch = jest.fn();
+    const mockGetState = () => ({
+      charts: {
+        chartKey: {
+          latestQueryFormData: {
+            time_grain_sqla: 'P1D',
+            granularity_sqla: 'Date',
+          },
+        },
+      },
+    });
+
+    beforeEach(() => {
+      jest.clearAllMocks();
+    });
+
+    it('should dispatch annotationQueryStarted and annotationQuerySuccess on successful query', async () => {
+      const annotation = {
+        name: 'Holidays',
+        annotationType: 'EVENT',
+        sourceType: 'NATIVE',
+        color: null,
+        opacity: '',
+        style: 'solid',
+        width: 1,
+        showMarkers: false,
+        hideLine: false,
+        value: 1,
+        overrides: {
+          time_range: null,
+        },
+        show: true,
+        showLabel: false,
+        titleColumn: '',
+        descriptionColumns: [],
+        timeColumn: '',
+        intervalEndColumn: '',
+      };
+      const key = undefined;
+
+      const postSpy = jest.spyOn(SupersetClient, 'post');
+      postSpy.mockImplementation(() =>
+        Promise.resolve({ json: { result: [] } }),
+      );
+      const buildV1ChartDataPayloadSpy = jest.spyOn(
+        exploreUtils,
+        'buildV1ChartDataPayload',
+      );
+
+      const queryFunc = actions.runAnnotationQuery({ annotation, key });
+      await queryFunc(mockDispatch, mockGetState);
+
+      expect(buildV1ChartDataPayloadSpy).toHaveBeenCalledWith({
+        formData: {
+          granularity: 'Date',
+          granularity_sqla: 'Date',
+          time_grain_sqla: 'P1D',
+        },
+        force: false,
+        resultFormat: 'json',
+        resultType: 'full',
+      });
+    });
+  });
 });