Optimize data step in the notebook editor (#39505)

frontend/src/metabase/common/components/FieldPicker/FieldPicker.tsx

@@ -17,8 +17,15 @@ interface FieldPickerProps {
   stageIndex: number;
   columns: Lib.ColumnMetadata[];
   "data-testid"?: string;
-  isColumnSelected: (column: Lib.ColumnMetadata) => boolean;
-  onToggle: (columnIndex: number, isSelected: boolean) => void;
+  isColumnSelected: (
+    column: Lib.ColumnMetadata,
+    columnInfo: Lib.ColumnDisplayInfo,
+  ) => boolean;
+  onToggle: (
+    column: Lib.ColumnMetadata,
+    isSelected: boolean,
+    columnIndex: number,
+  ) => void;
   onSelectAll: () => void;
   onSelectNone: () => void;
 }
@@ -35,27 +42,21 @@ export const FieldPicker = ({
 }: FieldPickerProps) => {
   const items = useMemo(
     () =>
-      columns.map(column => ({
-        ...Lib.displayInfo(query, stageIndex, column),
-        column,
-      })),
-    [query, stageIndex, columns],
+      columns.map(column => {
+        const columnInfo = Lib.displayInfo(query, stageIndex, column);
+        return {
+          column,
+          columnInfo,
+          isSelected: isColumnSelected(column, columnInfo),
+        };
+      }),
+    [query, stageIndex, columns, isColumnSelected],
   );
 
-  const isAll = useMemo(
-    () => columns.every(isColumnSelected),
-    [columns, isColumnSelected],
-  );
-
-  const isNone = useMemo(
-    () => columns.every(column => !isColumnSelected(column)),
-    [columns, isColumnSelected],
-  );
-
-  const isDisabledDeselection = useMemo(
-    () => columns.filter(isColumnSelected).length <= 1,
-    [columns, isColumnSelected],
-  );
+  const isAll = items.every(item => item.isSelected);
+  const isNone = items.every(item => !item.isSelected);
+  const isDisabledDeselection =
+    items.filter(item => item.isSelected).length <= 1;
 
   const handleLabelToggle = () => {
     if (isAll) {
@@ -79,16 +80,18 @@ export const FieldPicker = ({
         </label>
       </ToggleItem>
       {items.map((item, index) => (
-        <ColumnItem key={item.longDisplayName}>
+        <ColumnItem key={index}>
           <label>
             <Checkbox
-              checked={isColumnSelected(item.column)}
-              disabled={isColumnSelected(item.column) && isDisabledDeselection}
-              onChange={event => onToggle(index, event.target.checked)}
+              checked={item.isSelected}
+              disabled={item.isSelected && isDisabledDeselection}
+              onChange={event =>
+                onToggle(item.column, event.target.checked, index)
+              }
             />
 
             <ItemIcon name={getColumnIcon(item.column)} size={18} />
-            <ItemTitle>{item.displayName}</ItemTitle>
+            <ItemTitle>{item.columnInfo.displayName}</ItemTitle>
           </label>
         </ColumnItem>
       ))}

frontend/src/metabase/query_builder/components/notebook/steps/DataStep/DataStep.styled.tsx

 import styled from "@emotion/styled";
 
+import IconButtonWrapper from "metabase/components/IconButtonWrapper";
+import { color } from "metabase/lib/colors";
+
 import { NotebookCell } from "../../NotebookCell";
 
 export const DataStepCell = styled.div`
   padding: ${NotebookCell.CONTAINER_PADDING};
 `;
+
+export const DataStepIconButton = styled(IconButtonWrapper)`
+  color: ${color("white")};
+  padding: ${NotebookCell.CONTAINER_PADDING};
+  opacity: 0.5;
+`;

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

@@ -2,16 +2,15 @@ import { useMemo } from "react";
 import { t } from "ttag";
 
 import { FieldPicker } from "metabase/common/components/FieldPicker";
-import PopoverWithTrigger from "metabase/components/PopoverWithTrigger";
 import { DataSourceSelector } from "metabase/query_builder/components/DataSelector";
+import { Icon, Popover, Tooltip } from "metabase/ui";
 import * as Lib from "metabase-lib";
 import type { DatabaseId, TableId } from "metabase-types/api";
 
-import { FieldsPickerIcon, FIELDS_PICKER_STYLES } from "../../FieldsPickerIcon";
 import { NotebookCell, NotebookCellItem } from "../../NotebookCell";
 import type { NotebookStepUiComponentProps } from "../../types";
 
-import { DataStepCell } from "./DataStep.styled";
+import { DataStepCell, DataStepIconButton } from "./DataStep.styled";
 
 export const DataStep = ({
   query,
@@ -55,15 +54,15 @@ export const DataStep = ({
         inactive={!table}
         right={
           canSelectTableColumns && (
-            <DataFieldsPicker
+            <DataFieldPopover
               query={query}
               stageIndex={stageIndex}
               updateQuery={updateQuery}
             />
           )
         }
-        containerStyle={FIELDS_PICKER_STYLES.notebookItemContainer}
-        rightContainerStyle={FIELDS_PICKER_STYLES.notebookRightItemContainer}
+        containerStyle={{ padding: 0 }}
+        rightContainerStyle={{ width: 37, height: 37, padding: 0 }}
         data-testid="data-step-cell"
       >
         <DataSourceSelector
@@ -81,36 +80,67 @@ export const DataStep = ({
   );
 };
 
-interface DataFieldsPickerProps {
+interface DataFieldPopoverProps {
   query: Lib.Query;
   stageIndex: number;
   updateQuery: (query: Lib.Query) => Promise<void>;
 }
 
-export const DataFieldsPicker = ({
+function DataFieldPopover({
   query,
   stageIndex,
   updateQuery,
-}: DataFieldsPickerProps) => {
+}: DataFieldPopoverProps) {
+  return (
+    <Popover position="bottom-start">
+      <Popover.Target>
+        <Tooltip label={t`Pick columns`}>
+          <DataStepIconButton
+            aria-label={t`Pick columns`}
+            data-testid="fields-picker"
+          >
+            <Icon name="chevrondown" />
+          </DataStepIconButton>
+        </Tooltip>
+      </Popover.Target>
+      <Popover.Dropdown>
+        <DataFieldPicker
+          query={query}
+          stageIndex={stageIndex}
+          updateQuery={updateQuery}
+        />
+      </Popover.Dropdown>
+    </Popover>
+  );
+}
+
+interface DataFieldPickerProps {
+  query: Lib.Query;
+  stageIndex: number;
+  updateQuery: (query: Lib.Query) => Promise<void>;
+}
+
+function DataFieldPicker({
+  query,
+  stageIndex,
+  updateQuery,
+}: DataFieldPickerProps) {
   const columns = useMemo(
     () => Lib.fieldableColumns(query, stageIndex),
     [query, stageIndex],
   );
 
-  const handleToggle = (changedIndex: number, isSelected: boolean) => {
-    const nextColumns = columns.filter((_, currentIndex) => {
-      if (currentIndex === changedIndex) {
-        return isSelected;
-      }
-      const column = columns[currentIndex];
-      return Lib.displayInfo(query, stageIndex, column).selected;
-    });
-    const nextQuery = Lib.withFields(query, stageIndex, nextColumns);
+  const handleToggle = (column: Lib.ColumnMetadata, isSelected: boolean) => {
+    const nextQuery = isSelected
+      ? Lib.addField(query, stageIndex, column)
+      : Lib.removeField(query, stageIndex, column);
     updateQuery(nextQuery);
   };
 
-  const checkColumnSelected = (column: Lib.ColumnMetadata) =>
-    !!Lib.displayInfo(query, stageIndex, column).selected;
+  const isColumnSelected = (column: Lib.ColumnMetadata) => {
+    const columnInfo = Lib.displayInfo(query, stageIndex, column);
+    return Boolean(columnInfo.selected);
+  };
 
   const handleSelectAll = () => {
     const nextQuery = Lib.withFields(query, stageIndex, []);
@@ -123,19 +153,14 @@ export const DataFieldsPicker = ({
   };
 
   return (
-    <PopoverWithTrigger
-      triggerStyle={FIELDS_PICKER_STYLES.trigger}
-      triggerElement={FieldsPickerIcon}
-    >
-      <FieldPicker
-        query={query}
-        stageIndex={stageIndex}
-        columns={columns}
-        isColumnSelected={checkColumnSelected}
-        onToggle={handleToggle}
-        onSelectAll={handleSelectAll}
-        onSelectNone={handleSelectNone}
-      />
-    </PopoverWithTrigger>
+    <FieldPicker
+      query={query}
+      stageIndex={stageIndex}
+      columns={columns}
+      isColumnSelected={isColumnSelected}
+      onToggle={handleToggle}
+      onSelectAll={handleSelectAll}
+      onSelectNone={handleSelectNone}
+    />
   );
-};
+}

frontend/src/metabase/query_builder/components/notebook/steps/JoinStep/JoinTablePicker/JoinTablePicker.tsx

@@ -131,7 +131,11 @@ function JoinTableColumnsPicker({
   onChange,
 }: JoinTableColumnsPickerProps) {
   const [isOpened, setIsOpened] = useState(false);
-  const handleToggle = (changedIndex: number, isSelected: boolean) => {
+  const handleToggle = (
+    _column: Lib.ColumnMetadata,
+    isSelected: boolean,
+    changedIndex: number,
+  ) => {
     const nextColumns = columns.filter((_, currentIndex) =>
       currentIndex === changedIndex
         ? isSelected