From 0b3d3dd4caa7f4c31c1ba7229966a40ba0469e85 Mon Sep 17 00:00:00 2001
From: Ville Brofeldt <33317356+villebro@users.noreply.github.com>
Date: Thu, 19 May 2022 13:21:25 +0300
Subject: [PATCH] fix(generic-chart-axes): set x-axis if unset and ff is
 enabled (#20107)

* fix(generic-chart-axes): set x-axis if unset and ff is enabled

* simplify

* simplify

* continue cleanup

* yet more cleanup
---
 .../src/controlPanel.tsx                      | 13 +++++-----
 .../plugin-chart-echarts/src/controls.tsx     | 26 +++++++++++++++++--
 .../explore/controlUtils/getControlState.ts   | 11 +++-----
 superset-frontend/src/explore/store.js        |  9 ++++---
 4 files changed, 40 insertions(+), 19 deletions(-)

diff --git a/superset-frontend/plugins/legacy-plugin-chart-event-flow/src/controlPanel.tsx b/superset-frontend/plugins/legacy-plugin-chart-event-flow/src/controlPanel.tsx
index fd5ab5b5dd..11a0b88987 100644
--- a/superset-frontend/plugins/legacy-plugin-chart-event-flow/src/controlPanel.tsx
+++ b/superset-frontend/plugins/legacy-plugin-chart-event-flow/src/controlPanel.tsx
@@ -19,13 +19,14 @@
 import React from 'react';
 import { t, validateNonEmpty } from '@superset-ui/core';
 import {
-  formatSelectOptionsForRange,
-  ColumnOption,
   columnChoices,
+  ColumnOption,
+  ColumnMeta,
   ControlPanelConfig,
+  ControlState,
+  formatSelectOptionsForRange,
   sections,
   SelectControlConfig,
-  ColumnMeta,
 } from '@superset-ui/chart-controls';
 
 const config: ControlPanelConfig = {
@@ -46,10 +47,8 @@ const config: ControlPanelConfig = {
                 choices: columnChoices(state?.datasource),
               }),
               // choices is from `mapStateToProps`
-              default: (control: { choices?: string[] }) =>
-                control.choices && control.choices.length > 0
-                  ? control.choices[0][0]
-                  : null,
+              default: (control: ControlState) =>
+                control.choices?.[0]?.[0] || null,
               validators: [validateNonEmpty],
             },
           },
diff --git a/superset-frontend/plugins/plugin-chart-echarts/src/controls.tsx b/superset-frontend/plugins/plugin-chart-echarts/src/controls.tsx
index df050e6dbb..eca4723887 100644
--- a/superset-frontend/plugins/plugin-chart-echarts/src/controls.tsx
+++ b/superset-frontend/plugins/plugin-chart-echarts/src/controls.tsx
@@ -17,11 +17,18 @@
  * under the License.
  */
 import React from 'react';
-import { t, validateNonEmpty } from '@superset-ui/core';
+import {
+  FeatureFlag,
+  isFeatureEnabled,
+  t,
+  validateNonEmpty,
+} from '@superset-ui/core';
 import {
   ControlPanelsContainerProps,
+  ControlPanelState,
   ControlSetItem,
   ControlSetRow,
+  ControlState,
   sharedControls,
 } from '@superset-ui/chart-controls';
 import { DEFAULT_LEGEND_FORM_DATA } from './types';
@@ -143,7 +150,22 @@ export const xAxisControl: ControlSetItem = {
   config: {
     ...sharedControls.groupby,
     label: t('X-axis'),
-    default: null,
+    default: (
+      control: ControlState,
+      controlPanel: Partial<ControlPanelState>,
+    ) => {
+      // default to the chosen time column if x-axis is unset and the
+      // GENERIC_CHART_AXES feature flag is enabled
+      const { value } = control;
+      if (value) {
+        return value;
+      }
+      const timeColumn = controlPanel?.form_data?.granularity_sqla;
+      if (isFeatureEnabled(FeatureFlag.GENERIC_CHART_AXES) && timeColumn) {
+        return timeColumn;
+      }
+      return null;
+    },
     multi: false,
     description: t('Dimension to use on x-axis.'),
     validators: [validateNonEmpty],
diff --git a/superset-frontend/src/explore/controlUtils/getControlState.ts b/superset-frontend/src/explore/controlUtils/getControlState.ts
index ae0cbf531d..5014cd3c1a 100644
--- a/superset-frontend/src/explore/controlUtils/getControlState.ts
+++ b/superset-frontend/src/explore/controlUtils/getControlState.ts
@@ -85,7 +85,7 @@ function handleMissingChoice<T = ControlType>(control: ControlState<T>) {
 
 export function applyMapStateToPropsToControl<T = ControlType>(
   controlState: ControlState<T>,
-  controlPanelState: Partial<ControlPanelState>,
+  controlPanelState: Partial<ControlPanelState> | null,
 ) {
   const { mapStateToProps } = controlState;
   let state = { ...controlState };
@@ -120,7 +120,7 @@ export function applyMapStateToPropsToControl<T = ControlType>(
 
 export function getControlStateFromControlConfig<T = ControlType>(
   controlConfig: ControlConfig<T> | null,
-  controlPanelState: Partial<ControlPanelState>,
+  controlPanelState: Partial<ControlPanelState> | null,
   value?: JsonValue,
 ) {
   // skip invalid config values
@@ -130,10 +130,7 @@ export function getControlStateFromControlConfig<T = ControlType>(
   const controlState = { ...controlConfig, value } as ControlState<T>;
   // only apply mapStateToProps when control states have been initialized
   // or when explicitly didn't provide control panel state (mostly for testing)
-  if (
-    (controlPanelState && controlPanelState.controls) ||
-    controlPanelState === null
-  ) {
+  if (controlPanelState?.controls || controlPanelState === null) {
     return applyMapStateToPropsToControl(controlState, controlPanelState);
   }
   return controlState;
@@ -155,7 +152,7 @@ export function getControlState(
 export function getAllControlsState(
   vizType: string,
   datasourceType: DatasourceType,
-  state: ControlPanelState,
+  state: ControlPanelState | null,
   formData: QueryFormData,
 ) {
   const controlsState = {};
diff --git a/superset-frontend/src/explore/store.js b/superset-frontend/src/explore/store.js
index 16ad2324c0..80ad75e3e5 100644
--- a/superset-frontend/src/explore/store.js
+++ b/superset-frontend/src/explore/store.js
@@ -64,9 +64,12 @@ export function getControlsState(state, inputFormData) {
 export function applyDefaultFormData(inputFormData) {
   const datasourceType = inputFormData.datasource.split('__')[1];
   const vizType = inputFormData.viz_type;
-  const controlsState = getAllControlsState(vizType, datasourceType, null, {
-    ...inputFormData,
-  });
+  const controlsState = getAllControlsState(
+    vizType,
+    datasourceType,
+    null,
+    inputFormData,
+  );
   const controlFormData = getFormDataFromControls(controlsState);
 
   const formData = {};