From 314987f32dee789d7aa6af14943727af979ee30b Mon Sep 17 00:00:00 2001
From: Beto Dealmeida <roberto@dealmeida.net>
Date: Tue, 25 Apr 2023 19:44:50 -0700
Subject: [PATCH] fix: pivot v2 charts created before `GENERIC_CHART_AXES` is
 enabled (#23731)

---
 .../src/plugin/buildQuery.ts                  |  6 ++++-
 .../plugin-chart-pivot-table/src/types.ts     |  2 ++
 .../test/plugin/buildQuery.test.ts            | 27 +++++++++++++++++++
 3 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/superset-frontend/plugins/plugin-chart-pivot-table/src/plugin/buildQuery.ts b/superset-frontend/plugins/plugin-chart-pivot-table/src/plugin/buildQuery.ts
index 6e33b41cc7..ac2ddcd332 100644
--- a/superset-frontend/plugins/plugin-chart-pivot-table/src/plugin/buildQuery.ts
+++ b/superset-frontend/plugins/plugin-chart-pivot-table/src/plugin/buildQuery.ts
@@ -42,7 +42,11 @@ export default function buildQuery(formData: PivotTableQueryFormData) {
       isPhysicalColumn(col) &&
       formData.time_grain_sqla &&
       hasGenericChartAxes &&
-      formData?.temporal_columns_lookup?.[col]
+      /* Charts created before `GENERIC_CHART_AXES` is enabled have a different
+       * form data, with `granularity_sqla` set instead.
+       */
+      (formData?.temporal_columns_lookup?.[col] ||
+        formData.granularity_sqla === col)
     ) {
       return {
         timeGrain: formData.time_grain_sqla,
diff --git a/superset-frontend/plugins/plugin-chart-pivot-table/src/types.ts b/superset-frontend/plugins/plugin-chart-pivot-table/src/types.ts
index 8cf9a12ea3..e011f45931 100644
--- a/superset-frontend/plugins/plugin-chart-pivot-table/src/types.ts
+++ b/superset-frontend/plugins/plugin-chart-pivot-table/src/types.ts
@@ -80,6 +80,8 @@ interface PivotTableCustomizeProps {
     filters?: ContextMenuFilters,
   ) => void;
   timeGrainSqla?: TimeGranularity;
+  time_grain_sqla?: TimeGranularity;
+  granularity_sqla?: string;
 }
 
 export type PivotTableQueryFormData = QueryFormData &
diff --git a/superset-frontend/plugins/plugin-chart-pivot-table/test/plugin/buildQuery.test.ts b/superset-frontend/plugins/plugin-chart-pivot-table/test/plugin/buildQuery.test.ts
index a602bd4048..770cb9849b 100644
--- a/superset-frontend/plugins/plugin-chart-pivot-table/test/plugin/buildQuery.test.ts
+++ b/superset-frontend/plugins/plugin-chart-pivot-table/test/plugin/buildQuery.test.ts
@@ -17,6 +17,8 @@
  * under the License.
  */
 
+import { TimeGranularity } from '@superset-ui/core';
+import * as supersetCoreModule from '@superset-ui/core';
 import buildQuery from '../../src/plugin/buildQuery';
 import { PivotTableQueryFormData } from '../../src/types';
 
@@ -55,4 +57,29 @@ describe('PivotTableChart buildQuery', () => {
     const [query] = queryContext.queries;
     expect(query.columns).toEqual(['col1', 'col2', 'row1', 'row2']);
   });
+
+  it('should work with old charts after GENERIC_CHART_AXES is enabled', () => {
+    Object.defineProperty(supersetCoreModule, 'hasGenericChartAxes', {
+      value: true,
+    });
+    const modifiedFormData = {
+      ...formData,
+      time_grain_sqla: TimeGranularity.MONTH,
+      granularity_sqla: 'col1',
+    };
+    const queryContext = buildQuery(modifiedFormData);
+    const [query] = queryContext.queries;
+    expect(query.columns).toEqual([
+      {
+        timeGrain: 'P1M',
+        columnType: 'BASE_AXIS',
+        sqlExpression: 'col1',
+        label: 'col1',
+        expressionType: 'SQL',
+      },
+      'col2',
+      'row1',
+      'row2',
+    ]);
+  });
 });