From fa67315f5b4769b2d739da23ef253fd504d610d9 Mon Sep 17 00:00:00 2001
From: Geido <60598000+geido@users.noreply.github.com>
Date: Tue, 25 Oct 2022 14:42:11 +0300
Subject: [PATCH] fix: Default temporal column in Datasource (#21857)

---
 .../DatasourceControl.test.tsx                | 123 +++++++++++++++++-
 .../controls/DatasourceControl/index.jsx      |  30 +++--
 2 files changed, 142 insertions(+), 11 deletions(-)

diff --git a/superset-frontend/src/explore/components/controls/DatasourceControl/DatasourceControl.test.tsx b/superset-frontend/src/explore/components/controls/DatasourceControl/DatasourceControl.test.tsx
index 4d01ec7a79fb..8e41812fb844 100644
--- a/superset-frontend/src/explore/components/controls/DatasourceControl/DatasourceControl.test.tsx
+++ b/superset-frontend/src/explore/components/controls/DatasourceControl/DatasourceControl.test.tsx
@@ -18,9 +18,10 @@
  */
 
 import React from 'react';
-import { render, screen, act } from 'spec/helpers/testing-library';
+import { render, screen, act, waitFor } from 'spec/helpers/testing-library';
 import userEvent from '@testing-library/user-event';
 import { SupersetClient, DatasourceType } from '@superset-ui/core';
+import fetchMock from 'fetch-mock';
 import DatasourceControl from '.';
 
 const SupersetClientGet = jest.spyOn(SupersetClient, 'get');
@@ -45,7 +46,10 @@ const createProps = () => ({
   },
   validationErrors: [],
   name: 'datasource',
-  actions: {},
+  actions: {
+    changeDatasource: jest.fn(),
+    setControlValue: jest.fn(),
+  },
   isEditable: true,
   user: {
     createdOn: '2021-04-27T18:12:38.952304',
@@ -62,6 +66,16 @@ const createProps = () => ({
   onDatasourceSave: jest.fn(),
 });
 
+async function openAndSaveChanges(datasource: any) {
+  fetchMock.post('glob:*/datasource/save/', datasource, {
+    overwriteRoutes: true,
+  });
+  userEvent.click(screen.getByTestId('datasource-menu-trigger'));
+  userEvent.click(await screen.findByTestId('edit-dataset'));
+  userEvent.click(await screen.findByTestId('datasource-modal-save'));
+  userEvent.click(await screen.findByText('OK'));
+}
+
 test('Should render', async () => {
   const props = createProps();
   render(<DatasourceControl {...props} />);
@@ -229,3 +243,108 @@ test('Click on Save as dataset', async () => {
   expect(screen.getByRole('button', { name: /close/i })).toBeVisible();
   expect(dropdownField).toBeVisible();
 });
+
+test('should set the default temporal column', async () => {
+  const props = createProps();
+  const overrideProps = {
+    ...props,
+    form_data: {
+      granularity_sqla: 'test-col',
+    },
+    datasource: {
+      ...props.datasource,
+      main_dttm_col: 'test-default',
+      columns: [
+        {
+          column_name: 'test-col',
+          is_dttm: false,
+        },
+        {
+          column_name: 'test-default',
+          is_dttm: true,
+        },
+      ],
+    },
+  };
+  render(<DatasourceControl {...props} {...overrideProps} />, {
+    useRedux: true,
+  });
+
+  await openAndSaveChanges(overrideProps.datasource);
+  await waitFor(() => {
+    expect(props.actions.setControlValue).toHaveBeenCalledWith(
+      'granularity_sqla',
+      'test-default',
+    );
+  });
+});
+
+test('should set the first available temporal column', async () => {
+  const props = createProps();
+  const overrideProps = {
+    ...props,
+    form_data: {
+      granularity_sqla: 'test-col',
+    },
+    datasource: {
+      ...props.datasource,
+      main_dttm_col: null,
+      columns: [
+        {
+          column_name: 'test-col',
+          is_dttm: false,
+        },
+        {
+          column_name: 'test-first',
+          is_dttm: true,
+        },
+      ],
+    },
+  };
+  render(<DatasourceControl {...props} {...overrideProps} />, {
+    useRedux: true,
+  });
+
+  await openAndSaveChanges(overrideProps.datasource);
+  await waitFor(() => {
+    expect(props.actions.setControlValue).toHaveBeenCalledWith(
+      'granularity_sqla',
+      'test-first',
+    );
+  });
+});
+
+test('should not set the temporal column', async () => {
+  const props = createProps();
+  const overrideProps = {
+    ...props,
+    form_data: {
+      granularity_sqla: null,
+    },
+    datasource: {
+      ...props.datasource,
+      main_dttm_col: null,
+      columns: [
+        {
+          column_name: 'test-col',
+          is_dttm: false,
+        },
+        {
+          column_name: 'test-col-2',
+          is_dttm: false,
+        },
+      ],
+    },
+  };
+  render(<DatasourceControl {...props} {...overrideProps} />, {
+    useRedux: true,
+  });
+
+  await openAndSaveChanges(overrideProps.datasource);
+  await waitFor(() => {
+    expect(props.actions.setControlValue).toHaveBeenCalledWith(
+      'granularity_sqla',
+      null,
+    );
+  });
+});
diff --git a/superset-frontend/src/explore/components/controls/DatasourceControl/index.jsx b/superset-frontend/src/explore/components/controls/DatasourceControl/index.jsx
index 9fa4a8a2ba8e..266593adaeed 100644
--- a/superset-frontend/src/explore/components/controls/DatasourceControl/index.jsx
+++ b/superset-frontend/src/explore/components/controls/DatasourceControl/index.jsx
@@ -27,7 +27,7 @@ import {
   t,
   withTheme,
 } from '@superset-ui/core';
-
+import { getTemporalColumns } from '@superset-ui/chart-controls';
 import { getUrlParam } from 'src/utils/urlUtils';
 import { AntdDropdown } from 'src/components';
 import { Menu } from 'src/components/Menu';
@@ -171,19 +171,31 @@ class DatasourceControl extends React.PureComponent {
 
   onDatasourceSave = datasource => {
     this.props.actions.changeDatasource(datasource);
+    const { temporalColumns, defaultTemporalColumn } =
+      getTemporalColumns(datasource);
+    const { columns } = datasource;
+    // the current granularity_sqla might not be a temporal column anymore
     const timeCol = this.props.form_data?.granularity_sqla;
-    const { columns } = this.props.datasource;
-    const firstDttmCol = columns.find(column => column.is_dttm);
-    if (
-      datasource.type === 'table' &&
-      !columns.find(({ column_name }) => column_name === timeCol)?.is_dttm
-    ) {
-      // set `granularity_sqla` to first datatime column name or null
+    const isGranularitySqalTemporal = columns.find(
+      ({ column_name }) => column_name === timeCol,
+    )?.is_dttm;
+    // the current main_dttm_col might not be a temporal column anymore
+    const isDefaultTemporal = columns.find(
+      ({ column_name }) => column_name === defaultTemporalColumn,
+    )?.is_dttm;
+
+    // if the current granularity_sqla is empty or it is not a temporal column anymore
+    // let's update the control value
+    if (datasource.type === 'table' && !isGranularitySqalTemporal) {
+      const temporalColumn = isDefaultTemporal
+        ? defaultTemporalColumn
+        : temporalColumns?.[0];
       this.props.actions.setControlValue(
         'granularity_sqla',
-        firstDttmCol ? firstDttmCol.column_name : null,
+        temporalColumn || null,
       );
     }
+
     if (this.props.onDatasourceSave) {
       this.props.onDatasourceSave(datasource);
     }