From 541ae1d0fd1a044d8b1d75fe2484e772e18ad1d2 Mon Sep 17 00:00:00 2001
From: "Michael S. Molina" <michael.s.molina@gmail.com>
Date: Wed, 19 Jul 2023 10:05:24 -0300
Subject: [PATCH] fix: Dashboard time grain in Table

---
 .../plugin-chart-table/src/buildQuery.ts      | 13 ++++--
 .../test/buildQuery.test.ts                   | 42 ++++++++++++++++++-
 2 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/superset-frontend/plugins/plugin-chart-table/src/buildQuery.ts b/superset-frontend/plugins/plugin-chart-table/src/buildQuery.ts
index 211070f8bcc4a..13bf2e9c2d4c8 100644
--- a/superset-frontend/plugins/plugin-chart-table/src/buildQuery.ts
+++ b/superset-frontend/plugins/plugin-chart-table/src/buildQuery.ts
@@ -52,10 +52,15 @@ const buildQuery: BuildQuery<TableChartFormData> = (
   formData: TableChartFormData,
   options,
 ) => {
-  const { percent_metrics: percentMetrics, order_desc: orderDesc = false } =
-    formData;
+  const {
+    percent_metrics: percentMetrics,
+    order_desc: orderDesc = false,
+    extra_form_data,
+  } = formData;
   const queryMode = getQueryMode(formData);
   const sortByMetric = ensureIsArray(formData.timeseries_limit_metric)[0];
+  const time_grain_sqla =
+    extra_form_data?.time_grain_sqla || formData.time_grain_sqla;
   let formDataCopy = formData;
   // never include time in raw records mode
   if (queryMode === QueryMode.raw) {
@@ -102,12 +107,12 @@ const buildQuery: BuildQuery<TableChartFormData> = (
       columns = columns.map(col => {
         if (
           isPhysicalColumn(col) &&
-          formData.time_grain_sqla &&
+          time_grain_sqla &&
           hasGenericChartAxes &&
           formData?.temporal_columns_lookup?.[col]
         ) {
           return {
-            timeGrain: formData.time_grain_sqla,
+            timeGrain: time_grain_sqla,
             columnType: 'BASE_AXIS',
             sqlExpression: col,
             label: col,
diff --git a/superset-frontend/plugins/plugin-chart-table/test/buildQuery.test.ts b/superset-frontend/plugins/plugin-chart-table/test/buildQuery.test.ts
index 4578dd139ddc1..164f31aa05492 100644
--- a/superset-frontend/plugins/plugin-chart-table/test/buildQuery.test.ts
+++ b/superset-frontend/plugins/plugin-chart-table/test/buildQuery.test.ts
@@ -16,7 +16,8 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import { QueryMode } from '@superset-ui/core';
+import { QueryMode, TimeGranularity } from '@superset-ui/core';
+import * as supersetCoreModule from '@superset-ui/core';
 import buildQuery from '../src/buildQuery';
 import { TableChartFormData } from '../src/types';
 
@@ -81,5 +82,44 @@ describe('plugin-chart-table', () => {
       expect(query.columns).toEqual(['rawcol']);
       expect(query.post_processing).toEqual([]);
     });
+    it('should prefer extra_form_data.time_grain_sqla over formData.time_grain_sqla', () => {
+      Object.defineProperty(supersetCoreModule, 'hasGenericChartAxes', {
+        value: true,
+      });
+      const query = buildQuery({
+        ...basicFormData,
+        groupby: ['col1'],
+        query_mode: QueryMode.aggregate,
+        time_grain_sqla: TimeGranularity.MONTH,
+        extra_form_data: { time_grain_sqla: TimeGranularity.QUARTER },
+        temporal_columns_lookup: { col1: true },
+      }).queries[0];
+      expect(query.columns?.[0]).toEqual({
+        timeGrain: TimeGranularity.QUARTER,
+        columnType: 'BASE_AXIS',
+        sqlExpression: 'col1',
+        label: 'col1',
+        expressionType: 'SQL',
+      });
+    });
+    it('should fallback to formData.time_grain_sqla if extra_form_data.time_grain_sqla is not set', () => {
+      Object.defineProperty(supersetCoreModule, 'hasGenericChartAxes', {
+        value: true,
+      });
+      const query = buildQuery({
+        ...basicFormData,
+        time_grain_sqla: TimeGranularity.MONTH,
+        groupby: ['col1'],
+        query_mode: QueryMode.aggregate,
+        temporal_columns_lookup: { col1: true },
+      }).queries[0];
+      expect(query.columns?.[0]).toEqual({
+        timeGrain: TimeGranularity.MONTH,
+        columnType: 'BASE_AXIS',
+        sqlExpression: 'col1',
+        label: 'col1',
+        expressionType: 'SQL',
+      });
+    });
   });
 });