New columns picker for notebook editor (#17511)

frontend/src/metabase/query_builder/components/notebook/FieldsPickerIcon.jsx

+import React from "react";
+import styled from "styled-components";
+import { t } from "ttag";
+
+import Icon from "metabase/components/Icon";
+import Tooltip from "metabase/components/Tooltip";
+
+const StyledIcon = styled(Icon)`
+  opacity: 0.5;
+`;
+
+export function FieldsPickerIcon() {
+  return (
+    <Tooltip tooltip={<span>{t`Pick columns`}</span>}>
+      <StyledIcon name="table" size={14} />
+    </Tooltip>
+  );
+}
+
+export const FIELDS_PICKER_STYLES = {
+  notebookItemContainer: {
+    width: 37,
+    height: 37,
+  },
+  trigger: {
+    marginTop: 1,
+  },
+};

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

@@ -4,45 +4,50 @@ import { connect } from "react-redux";
 import { t } from "ttag";
 
 import { DatabaseSchemaAndTableDataSelector } from "metabase/query_builder/components/DataSelector";
-import { NotebookCell, NotebookCellItem } from "../NotebookCell";
-
 import { getDatabasesList } from "metabase/query_builder/selectors";
 
-function DataStep({ color, query, databases, updateQuery }) {
+import { NotebookCell, NotebookCellItem } from "../NotebookCell";
+import { FieldsPickerIcon, FIELDS_PICKER_STYLES } from "../FieldsPickerIcon";
+import FieldsPicker from "./FieldsPicker";
+
+function DataStep({ color, query, updateQuery }) {
   const table = query.table();
+  const canSelectTableColumns = table && query.isRaw();
   return (
     <NotebookCell color={color}>
-      <DatabaseSchemaAndTableDataSelector
-        hasTableSearch
-        databaseQuery={{ saved: true }}
-        selectedDatabaseId={query.databaseId()}
-        selectedTableId={query.tableId()}
-        setSourceTableFn={tableId =>
-          query
-            .setTableId(tableId)
-            .setDefaultQuery()
-            .update(updateQuery)
-        }
-        isInitiallyOpen={!query.tableId()}
-        triggerElement={
-          !query.tableId() ? (
-            <NotebookCellItem color={color} inactive>
-              {t`Pick your starting data`}
-            </NotebookCellItem>
-          ) : (
-            <NotebookCellItem color={color} data-testid="data-step-cell">
-              {table && table.displayName()}
-            </NotebookCellItem>
+      <NotebookCellItem
+        color={color}
+        inactive={!table}
+        right={
+          canSelectTableColumns && (
+            <DataFieldsPicker
+              query={query}
+              updateQuery={updateQuery}
+              triggerStyle={FIELDS_PICKER_STYLES.trigger}
+              triggerElement={<FieldsPickerIcon />}
+            />
           )
         }
-      />
-      {table && query.isRaw() && (
-        <DataFieldsPicker
-          className="ml-auto text-bold"
-          query={query}
-          updateQuery={updateQuery}
+        rightContainerStyle={FIELDS_PICKER_STYLES.notebookItemContainer}
+        data-testid="data-step-cell"
+      >
+        <DatabaseSchemaAndTableDataSelector
+          hasTableSearch
+          databaseQuery={{ saved: true }}
+          selectedDatabaseId={query.databaseId()}
+          selectedTableId={query.tableId()}
+          setSourceTableFn={tableId =>
+            query
+              .setTableId(tableId)
+              .setDefaultQuery()
+              .update(updateQuery)
+          }
+          isInitiallyOpen={!query.tableId()}
+          triggerElement={
+            table ? table.displayName() : t`Pick your starting data`
+          }
         />
-      )}
+      </NotebookCellItem>
     </NotebookCell>
   );
 }
@@ -51,16 +56,14 @@ export default connect(state => ({ databases: getDatabasesList(state) }))(
   DataStep,
 );
 
-import FieldsPicker from "./FieldsPicker";
-
-const DataFieldsPicker = ({ className, query, updateQuery }) => {
+const DataFieldsPicker = ({ query, updateQuery, ...props }) => {
   const dimensions = query.tableDimensions();
   const selectedDimensions = query.columnDimensions();
   const selected = new Set(selectedDimensions.map(d => d.key()));
   const fields = query.fields();
   return (
     <FieldsPicker
-      className={className}
+      {...props}
       dimensions={dimensions}
       selectedDimensions={selectedDimensions}
       isAll={!fields || fields.length === 0}

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

@@ -16,13 +16,16 @@ export default function FieldsPicker({
   onSelectAll,
   onSelectNone,
   onToggleDimension,
+  triggerElement = t`Columns`,
+  ...props
 }) {
   const selected = new Set(selectedDimensions.map(d => d.key()));
   return (
     <PopoverWithTrigger
-      triggerElement={t`Columns`}
+      triggerElement={triggerElement}
       triggerClasses={className}
       sizeToFit
+      {...props}
     >
       <ul className="pt1">
         {(onSelectAll || onSelectNone) && (

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

@@ -14,6 +14,7 @@ import {
   NotebookCellItem,
   NotebookCellAdd,
 } from "../NotebookCell";
+import { FieldsPickerIcon, FIELDS_PICKER_STYLES } from "../FieldsPickerIcon";
 import FieldsPicker from "./FieldsPicker";
 import {
   JoinClausesContainer,
@@ -211,14 +212,6 @@ function JoinClause({ color, join, updateQuery, showRemove }) {
         </JoinedTableControlRoot>
       )}
 
-      {join.isValid() && (
-        <JoinFieldsPicker
-          className="ml-auto text-bold"
-          join={join}
-          updateQuery={updateQuery}
-        />
-      )}
-
       {showRemove && <RemoveJoinIcon onClick={removeJoin} />}
     </JoinClauseRoot>
   );
@@ -258,7 +251,21 @@ function JoinTablePicker({
   }
 
   return (
-    <NotebookCellItem color={color} inactive={!joinedTable}>
+    <NotebookCellItem
+      color={color}
+      inactive={!joinedTable}
+      right={
+        joinedTable && (
+          <JoinFieldsPicker
+            join={join}
+            updateQuery={updateQuery}
+            triggerElement={<FieldsPickerIcon />}
+            triggerStyle={FIELDS_PICKER_STYLES.trigger}
+          />
+        )
+      }
+      rightContainerStyle={FIELDS_PICKER_STYLES.notebookItemContainer}
+    >
       <DatabaseSchemaAndTableDataSelector
         hasTableSearch
         canChangeDatabase={false}
@@ -429,10 +436,9 @@ JoinDimensionPicker.propTypes = joinDimensionPickerPropTypes;
 const joinFieldsPickerPropTypes = {
   join: PropTypes.object.isRequired,
   updateQuery: PropTypes.func.isRequired,
-  className: PropTypes.string,
 };
 
-const JoinFieldsPicker = ({ className, join, updateQuery }) => {
+const JoinFieldsPicker = ({ join, updateQuery, ...props }) => {
   const dimensions = join.joinedDimensions();
   const selectedDimensions = join.fieldsDimensions();
   const selected = new Set(selectedDimensions.map(d => d.key()));
@@ -470,7 +476,7 @@ const JoinFieldsPicker = ({ className, join, updateQuery }) => {
 
   return (
     <FieldsPicker
-      className={className}
+      {...props}
       dimensions={dimensions}
       selectedDimensions={selectedDimensions}
       isAll={join.fields === "all"}

frontend/src/metabase/query_builder/components/notebook/steps/JoinStep.unit.spec.js

@@ -271,7 +271,7 @@ describe("Notebook Editor > Join Step", () => {
     const { onQueryChange } = setup();
     await selectTable(/Products/i);
 
-    fireEvent.click(screen.getByText("Columns"));
+    fireEvent.click(screen.getByLabelText("table icon"));
     fireEvent.click(screen.getByText("Select None"));
     fireEvent.click(screen.getByText("Category"));