fix: Default temporal column in Datasource

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,17 @@ const createProps = () => ({
   onDatasourceSave: jest.fn(),
 });
 
+fetchMock.post('glob:*/datasource/save/', {
+  ...createProps().datasource,
+});
+
+async function openAndSaveChanges() {
+  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 +244,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();
+  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();
+  await waitFor(() => {
+    expect(props.actions.setControlValue).toHaveBeenCalledWith(
+      'granularity_sqla',
+      'test-first',
+    );
+  });
+});
+
+test('should set the temporal column at null', 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();
+  await waitFor(() => {
+    expect(props.actions.setControlValue).toHaveBeenCalledWith(
+      'granularity_sqla',
+      null,
+    );
+  });
+});

superset-frontend/src/explore/components/controls/DatasourceControl/index.jsx

@@ -171,19 +171,26 @@ class DatasourceControl extends React.PureComponent {
 
   onDatasourceSave = datasource => {
     this.props.actions.changeDatasource(datasource);
-    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
-      this.props.actions.setControlValue(
-        'granularity_sqla',
-        firstDttmCol ? firstDttmCol.column_name : null,
-      );
+    const timeCol = this.props.form_data?.granularity_sqla;
+    const timeColConf = columns.find(
+      ({ column_name }) => column_name === timeCol,
+    );
+    const firstDttmCol =
+      columns.find(column => column.is_dttm)?.column_name || null;
+    const currentMainDttmCol = columns.find(
+      c => c.column_name === datasource.main_dttm_col,
+    );
+    const setDttmCol = currentMainDttmCol?.is_dttm
+      ? currentMainDttmCol.column_name
+      : firstDttmCol;
+    if (datasource.type === 'table') {
+      // resets new default time col or picks the first available one
+      if (!timeColConf?.is_dttm) {
+        this.props.actions.setControlValue('granularity_sqla', setDttmCol);
+      }
     }
+
     if (this.props.onDatasourceSave) {
       this.props.onDatasourceSave(datasource);
     }