fix deselecting columns hides custom fields (#20239)

frontend/src/metabase/query_builder/components/notebook/steps/DataStep.jsx

@@ -5,6 +5,7 @@ import { t } from "ttag";
 
 import { CollectionDatasetOrDataSourceSelector } from "metabase/query_builder/components/DataSelector";
 import { getDatabasesList } from "metabase/query_builder/selectors";
+import { isLocalField } from "metabase/lib/query/field_ref";
 
 import { NotebookCell, NotebookCellItem } from "../NotebookCell";
 import {
@@ -76,30 +77,35 @@ export default connect(state => ({ databases: getDatabasesList(state) }))(
 
 const DataFieldsPicker = ({ query, updateQuery, ...props }) => {
   const dimensions = query.tableDimensions();
+  const expressionDimensions = query.expressionDimensions();
   const selectedDimensions = query.columnDimensions();
   const selected = new Set(selectedDimensions.map(d => d.key()));
   const fields = query.fields();
 
   const handleSelectNone = () => {
-    query.setFields([dimensions[0].mbql()]).update(updateQuery);
-  };
-
-  const handleToggleDimension = dimension =>
     query
-      .setFields(
-        dimensions
-          .filter(d => {
-            if (d === dimension) {
-              return !selected.has(d.key());
-            } else {
-              return selected.has(d.key());
-            }
-          })
-          .map(d => d.mbql()),
-      )
+      .setFields([
+        dimensions[0].mbql(),
+        ...expressionDimensions.map(d => d.mbql()),
+      ])
       .update(updateQuery);
+  };
+
+  const handleToggleDimension = dimension => {
+    const newFields = [...dimensions, ...expressionDimensions]
+      .filter(d => {
+        if (d === dimension) {
+          return !selected.has(d.key());
+        } else {
+          return selected.has(d.key());
+        }
+      })
+      .map(d => d.mbql());
+
+    query.setFields(newFields).update(updateQuery);
+  };
 
-  const hasOneColumnSelected = fields.length === 1;
+  const hasOneColumnSelected = fields.filter(isLocalField).length === 1;
 
   return (
     <FieldsPicker

frontend/test/metabase/scenarios/custom-column/reproductions/20229-cc-missing-if-all-columns-not-selected.cy.spec.js

@@ -19,7 +19,7 @@ const questionDetails = {
   },
 };
 
-describe.skip("issue 20229", () => {
+describe("issue 20229", () => {
   beforeEach(() => {
     restore();
     cy.signInAsAdmin();