fix: keep calculated columns when datasource is updated (#32523)

superset-frontend/src/components/Datasource/utils.js

@@ -56,8 +56,11 @@ export function updateColumns(prevCols, newCols, addSuccessToast) {
     added: [],
     modified: [],
     removed: prevCols
-      .map(col => col.column_name)
-      .filter(col => !databaseColumnNames.includes(col)),
+      .filter(
+        col =>
+          !(col.expression || databaseColumnNames.includes(col.column_name)),
+      )
+      .map(col => col.column_name),
     finalColumns: [],
   };
   newCols.forEach(col => {
@@ -89,6 +92,13 @@ export function updateColumns(prevCols, newCols, addSuccessToast) {
       columnChanges.finalColumns.push(currentCol);
     }
   });
+  // push all calculated columns
+  prevCols
+    .filter(col => col.expression)
+    .forEach(col => {
+      columnChanges.finalColumns.push(col);
+    });
+
   if (columnChanges.modified.length) {
     addSuccessToast(
       tn(

superset-frontend/src/components/Datasource/utils.test.tsx

@@ -0,0 +1,202 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import { tn } from '@superset-ui/core';
+import { updateColumns } from './utils';
+
+describe('updateColumns', () => {
+  let addSuccessToast: jest.Mock;
+
+  beforeEach(() => {
+    addSuccessToast = jest.fn();
+  });
+
+  it('adds new columns when prevCols is empty', () => {
+    interface Column {
+      column_name: string;
+      type: string;
+      is_dttm: boolean;
+    }
+
+    const prevCols: Column[] = [];
+    const newCols = [
+      { column_name: 'col1', type: 'string', is_dttm: false },
+      { column_name: 'col2', type: 'number', is_dttm: true },
+    ];
+    const result = updateColumns(prevCols, newCols, addSuccessToast);
+    expect(result.added.sort()).toEqual(['col1', 'col2'].sort());
+    expect(result.modified).toEqual([]);
+    expect(result.removed).toEqual([]);
+    expect(result.finalColumns).toHaveLength(2);
+    // Only the added toast should be fired
+    expect(addSuccessToast).toHaveBeenCalledTimes(1);
+    expect(addSuccessToast).toHaveBeenCalledWith(
+      tn(
+        'Added 1 new column to the virtual dataset',
+        'Added %s new columns to the virtual dataset',
+        2,
+      ),
+    );
+  });
+
+  it('modifies columns when type or is_dttm changes', () => {
+    const prevCols = [
+      { column_name: 'col1', type: 'string', is_dttm: false },
+      { column_name: 'col2', type: 'number', is_dttm: false },
+    ];
+    const newCols = [
+      // col1 unchanged
+      { column_name: 'col1', type: 'string', is_dttm: false },
+      // col2 modified (type changed)
+      { column_name: 'col2', type: 'float', is_dttm: false },
+    ];
+    const result = updateColumns(prevCols, newCols, addSuccessToast);
+    expect(result.added).toEqual([]);
+    expect(result.modified).toEqual(['col2']);
+    // No columns removed
+    expect(result.removed).toEqual([]);
+    // Final columns: first is unchanged, second is updated
+    expect(result.finalColumns).toHaveLength(2);
+    const updatedCol2 = (
+      result.finalColumns as {
+        column_name: string;
+        type: string;
+        is_dttm: boolean;
+      }[]
+    ).find(c => c.column_name === 'col2');
+    expect(updatedCol2?.type).toBe('float');
+    // Modified toast should be fired
+    expect(addSuccessToast).toHaveBeenCalledTimes(1);
+    expect(addSuccessToast).toHaveBeenCalledWith(
+      tn(
+        'Modified 1 column in the virtual dataset',
+        'Modified %s columns in the virtual dataset',
+        1,
+      ),
+    );
+  });
+
+  it('removes columns not present in newCols', () => {
+    const prevCols = [
+      { column_name: 'col1', type: 'string', is_dttm: false },
+      { column_name: 'col2', type: 'number', is_dttm: true },
+    ];
+    const newCols = [
+      // Only col2 is present
+      { column_name: 'col2', type: 'number', is_dttm: true },
+    ];
+    const result = updateColumns(prevCols, newCols, addSuccessToast);
+    // col1 should be marked as removed
+    expect(result.removed).toEqual(['col1']);
+    expect(result.added).toEqual([]);
+    expect(result.modified).toEqual([]);
+    expect(result.finalColumns).toHaveLength(1);
+    // Removed toast should be fired
+    expect(addSuccessToast).toHaveBeenCalledTimes(1);
+    expect(addSuccessToast).toHaveBeenCalledWith(
+      tn(
+        'Removed 1 column from the virtual dataset',
+        'Removed %s columns from the virtual dataset',
+        1,
+      ),
+    );
+  });
+
+  it('handles combined additions, modifications, and removals', () => {
+    const prevCols = [
+      { column_name: 'col1', type: 'string', is_dttm: false },
+      { column_name: 'col2', type: 'number', is_dttm: false },
+      { column_name: 'col3', type: 'number', is_dttm: false },
+    ];
+    const newCols = [
+      // col1 modified
+      { column_name: 'col1', type: 'string', is_dttm: true },
+      // col2 unchanged
+      { column_name: 'col2', type: 'number', is_dttm: false },
+      // col4 is a new column
+      { column_name: 'col4', type: 'boolean', is_dttm: false },
+    ];
+    const result = updateColumns(prevCols, newCols, addSuccessToast);
+    expect(result.added).toEqual(['col4']);
+    expect(result.modified).toEqual(['col1']);
+    // col3 is removed since it is missing in newCols and has no expression
+    expect(result.removed).toEqual(['col3']);
+    expect(result.finalColumns).toHaveLength(3);
+    // Three types of changes should fire three separate toasts
+    expect(addSuccessToast).toHaveBeenCalledTimes(3);
+    expect(addSuccessToast.mock.calls).toEqual([
+      [
+        tn(
+          'Modified 1 column in the virtual dataset',
+          'Modified %s columns in the virtual dataset',
+          1,
+        ),
+      ],
+      [
+        tn(
+          'Removed 1 column from the virtual dataset',
+          'Removed %s columns from the virtual dataset',
+          1,
+        ),
+      ],
+      [
+        tn(
+          'Added 1 new column to the virtual dataset',
+          'Added %s new columns to the virtual dataset',
+          1,
+        ),
+      ],
+    ]);
+  });
+  it('should not remove columns with expressions', () => {
+    const prevCols = [
+      { column_name: 'col1', type: 'string', is_dttm: false },
+      { column_name: 'col2', type: 'number', is_dttm: false },
+      {
+        column_name: 'col3',
+        type: 'number',
+        is_dttm: false,
+        expression: 'expr',
+      },
+    ];
+    const newCols = [
+      // col1 modified
+      { column_name: 'col1', type: 'string', is_dttm: true },
+      // col2 unchanged
+      { column_name: 'col2', type: 'number', is_dttm: false },
+    ];
+    const result = updateColumns(prevCols, newCols, addSuccessToast);
+    expect(result.added).toEqual([]);
+    expect(result.modified).toEqual(['col1']);
+    // col3 is not removed since it has an expression
+    expect(result.removed).toEqual([]);
+    expect(result.finalColumns).toHaveLength(3);
+    // Two types of changes should fire two separate toasts
+    expect(addSuccessToast).toHaveBeenCalledTimes(1);
+    expect(addSuccessToast.mock.calls).toEqual([
+      [
+        tn(
+          'Modified 1 column in the virtual dataset',
+          'Modified %s columns in the virtual dataset',
+          1,
+        ),
+      ],
+    ]);
+  });
+});