Refactor notebook editor's join step (#17445)

* Fix imports order

* Extract JoinClausesContainer

* Refactor JoinClause container element

* Extract JoinTypePicker component

* Refactor JoinStepPicker

* Extract JoinTablePicker

* Extract JoinedTableControlRoot

* Extract label components

* Extract RemoveJoinIcon

* Add prop types

* Turn JoinClause into func component, fix ref usage

* Update notebook's step prop-types shape

* Extract callbacks

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

-/* eslint-disable react/prop-types */
-import React from "react";
-
-import { Flex } from "grid-styled";
-import cx from "classnames";
+import React, { useRef } from "react";
+import PropTypes from "prop-types";
 import _ from "underscore";
 import { t } from "ttag";
 
+import PopoverWithTrigger from "metabase/components/PopoverWithTrigger";
+
+import { DatabaseSchemaAndTableDataSelector } from "metabase/query_builder/components/DataSelector";
+import FieldList from "metabase/query_builder/components/FieldList";
+import Join from "metabase-lib/lib/queries/structured/Join";
+
 import {
   NotebookCell,
   NotebookCellItem,
   NotebookCellAdd,
 } from "../NotebookCell";
+import FieldsPicker from "./FieldsPicker";
+import {
+  JoinClausesContainer,
+  JoinClauseContainer,
+  JoinClauseRoot,
+  JoinStrategyIcon,
+  JoinTypeSelectRoot,
+  JoinTypeOptionRoot,
+  JoinTypeIcon,
+  JoinedTableControlRoot,
+  JoinWhereConditionLabel,
+  JoinOnConditionLabel,
+  RemoveJoinIcon,
+} from "./JoinStep.styled";
 
-import Icon from "metabase/components/Icon";
-import PopoverWithTrigger from "metabase/components/PopoverWithTrigger";
+const stepShape = {
+  id: PropTypes.string.isRequired,
+  type: PropTypes.string.isRequired,
+  query: PropTypes.object.isRequired,
+  previewQuery: PropTypes.object,
+  valid: PropTypes.bool.isRequired,
+  visible: PropTypes.bool.isRequired,
+  stageIndex: PropTypes.number.isRequired,
+  itemIndex: PropTypes.number.isRequired,
+  update: PropTypes.func.isRequired,
+  revert: PropTypes.func.isRequired,
+  clean: PropTypes.func.isRequired,
+  actions: PropTypes.array.isRequired,
 
-import { DatabaseSchemaAndTableDataSelector } from "metabase/query_builder/components/DataSelector";
-import FieldList from "metabase/query_builder/components/FieldList";
-import Join from "metabase-lib/lib/queries/structured/Join";
+  previous: stepShape,
+  next: stepShape,
+};
+
+const joinStepPropTypes = {
+  query: PropTypes.object.isRequired,
+  step: PropTypes.shape(stepShape).isRequired,
+  color: PropTypes.string,
+  isLastOpened: PropTypes.bool,
+  updateQuery: PropTypes.func.isRequired,
+};
 
 export default function JoinStep({
   color,
@@ -25,7 +61,6 @@ export default function JoinStep({
   step,
   updateQuery,
   isLastOpened,
-  ...props
 }) {
   const isSingleJoinStep = step.itemIndex != null;
   let joins = query.joins();
@@ -37,190 +72,273 @@ export default function JoinStep({
     joins = [new Join({ fields: "all" }, query.joins().length, query)];
   }
   const valid = _.all(joins, join => join.isValid());
+
+  function addNewJoinClause() {
+    query.join(new Join({ fields: "all" })).update(updateQuery);
+  }
+
   return (
     <NotebookCell color={color} flexWrap="nowrap">
-      <Flex flexDirection="column" className="flex-full">
-        {joins.map((join, index) => (
-          <JoinClause
-            mb={index === joins.length - 1 ? 0 : 2}
-            key={index}
-            color={color}
-            join={join}
-            showRemove={joins.length > 1}
-            updateQuery={updateQuery}
-            isLastOpened={isLastOpened && index === join.length - 1}
-          />
-        ))}
-      </Flex>
+      <JoinClausesContainer>
+        {joins.map((join, index) => {
+          const isLast = index === joins.length - 1;
+          return (
+            <JoinClauseContainer key={index} isLast={isLast}>
+              <JoinClause
+                join={join}
+                color={color}
+                showRemove={joins.length > 1}
+                updateQuery={updateQuery}
+                isLastOpened={isLastOpened && isLast}
+              />
+            </JoinClauseContainer>
+          );
+        })}
+      </JoinClausesContainer>
       {!isSingleJoinStep && valid && (
         <NotebookCellAdd
           color={color}
           className="cursor-pointer ml-auto"
-          onClick={() => {
-            query.join(new Join({ fields: "all" })).update(updateQuery);
-          }}
+          onClick={addNewJoinClause}
         />
       )}
     </NotebookCell>
   );
 }
 
-class JoinClause extends React.Component {
-  render() {
-    const { color, join, updateQuery, showRemove, ...props } = this.props;
-    const query = join.query();
-    if (!query) {
-      return null;
+JoinStep.propTypes = joinStepPropTypes;
+
+const joinClausePropTypes = {
+  color: PropTypes.string,
+  join: PropTypes.object,
+  updateQuery: PropTypes.func,
+  showRemove: PropTypes.bool,
+};
+
+function JoinClause({ color, join, updateQuery, showRemove }) {
+  const joinDimensionPickerRef = useRef();
+  const parentDimensionPickerRef = useRef();
+
+  const query = join.query();
+  if (!query) {
+    return null;
+  }
+
+  let lhsTable;
+  if (join.index() === 0) {
+    // first join's lhs is always the parent table
+    lhsTable = join.parentTable();
+  } else if (join.parentDimension()) {
+    // subsequent can be one of the previously joined tables
+    // NOTE: `lhsDimension` would probably be a better name for `parentDimension`
+    lhsTable = join.parentDimension().field().table;
+  }
+
+  const joinedTable = join.joinedTable();
+
+  function onSourceTableSet(newJoin) {
+    if (!newJoin.parentDimension()) {
+      setTimeout(() => {
+        parentDimensionPickerRef.current.open();
+      });
     }
+  }
 
-    let lhsTable;
-    if (join.index() === 0) {
-      // first join's lhs is always the parent table
-      lhsTable = join.parentTable();
-    } else if (join.parentDimension()) {
-      // subsequent can be one of the previously joined tables
-      // NOTE: `lhsDimension` would probably be a better name for `parentDimension`
-      lhsTable = join.parentDimension().field().table;
+  function onParentDimensionChange(fieldRef) {
+    join
+      .setParentDimension(fieldRef)
+      .setDefaultAlias()
+      .parent()
+      .update(updateQuery);
+    if (!join.joinDimension()) {
+      joinDimensionPickerRef.current.open();
     }
+  }
 
-    const joinedTable = join.joinedTable();
-    const strategyOption = join.strategyOption();
-    return (
-      <Flex align="center" flex="1 1 auto" {...props}>
-        <NotebookCellItem color={color} icon="table2">
-          {(lhsTable && lhsTable.displayName()) || `Previous results`}
-        </NotebookCellItem>
-
-        <PopoverWithTrigger
-          triggerElement={
-            strategyOption ? (
-              <Icon
-                tooltip={t`Change join type`}
-                className="text-brand mr1"
-                name={strategyOption.icon}
-                size={32}
-              />
-            ) : (
-              <NotebookCellItem color={color}>
-                {`Choose a join type`}
-              </NotebookCellItem>
-            )
-          }
-        >
-          {({ onClose }) => (
-            <JoinTypeSelect
-              value={strategyOption && strategyOption.value}
-              onChange={strategy => {
-                join
-                  .setStrategy(strategy)
-                  .parent()
-                  .update(updateQuery);
-                onClose();
-              }}
-              options={join.strategyOptions()}
-            />
-          )}
-        </PopoverWithTrigger>
-
-        <DatabaseSchemaAndTableDataSelector
-          hasTableSearch
-          canChangeDatabase={false}
-          databases={[
-            query.database(),
-            query.database().savedQuestionsDatabase(),
-          ].filter(d => d)}
-          tableFilter={table => table.db_id === query.database().id}
-          selectedDatabaseId={query.databaseId()}
-          selectedTableId={join.joinSourceTableId()}
-          setSourceTableFn={tableId => {
-            const newJoin = join
-              .setJoinSourceTableId(tableId)
-              .setDefaultCondition()
-              .setDefaultAlias();
-            newJoin.parent().update(updateQuery);
-            // _parentDimensionPicker won't be rendered until next update
-            if (!newJoin.parentDimension()) {
-              setTimeout(() => {
-                this._parentDimensionPicker.open();
-              });
-            }
-          }}
-          isInitiallyOpen={join.joinSourceTableId() == null}
-          triggerElement={
-            <NotebookCellItem
-              color={color}
-              icon="table2"
-              inactive={!joinedTable}
-            >
-              {joinedTable ? joinedTable.displayName() : t`Pick a table...`}
-            </NotebookCellItem>
-          }
-        />
+  function onJoinDimensionChange(fieldRef) {
+    join
+      .setJoinDimension(fieldRef)
+      .parent()
+      .update(updateQuery);
+  }
 
-        {joinedTable && (
-          <Flex align="center">
-            <span className="text-medium text-bold ml1 mr2">where</span>
-
-            <JoinDimensionPicker
-              color={color}
-              query={query}
-              dimension={join.parentDimension()}
-              options={join.parentDimensionOptions()}
-              onChange={fieldRef => {
-                join
-                  .setParentDimension(fieldRef)
-                  .setDefaultAlias()
-                  .parent()
-                  .update(updateQuery);
-                if (!join.joinDimension()) {
-                  this._joinDimensionPicker.open();
-                }
-              }}
-              ref={ref => (this._parentDimensionPicker = ref)}
-            />
-
-            <span className="text-medium text-bold mr1">=</span>
-
-            <JoinDimensionPicker
-              color={color}
-              query={query}
-              dimension={join.joinDimension()}
-              options={join.joinDimensionOptions()}
-              onChange={fieldRef => {
-                join
-                  .setJoinDimension(fieldRef)
-                  .parent()
-                  .update(updateQuery);
-              }}
-              ref={ref => (this._joinDimensionPicker = ref)}
-            />
-          </Flex>
-        )}
+  function removeJoin() {
+    join.remove().update(updateQuery);
+  }
 
-        {join.isValid() && (
-          <JoinFieldsPicker
-            className="mb1 ml-auto text-bold"
-            join={join}
-            updateQuery={updateQuery}
+  return (
+    <JoinClauseRoot>
+      <NotebookCellItem color={color} icon="table2">
+        {(lhsTable && lhsTable.displayName()) || `Previous results`}
+      </NotebookCellItem>
+
+      <JoinTypePicker join={join} color={color} updateQuery={updateQuery} />
+
+      <JoinTablePicker
+        join={join}
+        query={query}
+        joinedTable={joinedTable}
+        color={color}
+        updateQuery={updateQuery}
+        onSourceTableSet={onSourceTableSet}
+      />
+
+      {joinedTable && (
+        <JoinedTableControlRoot>
+          <JoinWhereConditionLabel />
+
+          <JoinDimensionPicker
+            color={color}
+            query={query}
+            dimension={join.parentDimension()}
+            options={join.parentDimensionOptions()}
+            onChange={onParentDimensionChange}
+            ref={parentDimensionPickerRef}
           />
-        )}
 
-        {showRemove && (
-          <Icon
-            name="close"
-            size={18}
-            className="cursor-pointer text-light text-medium-hover"
-            onClick={() => join.remove().update(updateQuery)}
+          <JoinOnConditionLabel />
+
+          <JoinDimensionPicker
+            color={color}
+            query={query}
+            dimension={join.joinDimension()}
+            options={join.joinDimensionOptions()}
+            onChange={onJoinDimensionChange}
+            ref={joinDimensionPickerRef}
           />
-        )}
-      </Flex>
-    );
+        </JoinedTableControlRoot>
+      )}
+
+      {join.isValid() && (
+        <JoinFieldsPicker
+          className="mb1 ml-auto text-bold"
+          join={join}
+          updateQuery={updateQuery}
+        />
+      )}
+
+      {showRemove && <RemoveJoinIcon onClick={removeJoin} />}
+    </JoinClauseRoot>
+  );
+}
+
+JoinClause.propTypes = joinClausePropTypes;
+
+const joinTablePickerPropTypes = {
+  join: PropTypes.object,
+  query: PropTypes.object,
+  joinedTable: PropTypes.object,
+  color: PropTypes.string,
+  updateQuery: PropTypes.func,
+  onSourceTableSet: PropTypes.func.isRequired,
+};
+
+function JoinTablePicker({
+  join,
+  query,
+  joinedTable,
+  color,
+  updateQuery,
+  onSourceTableSet,
+}) {
+  const databases = [
+    query.database(),
+    query.database().savedQuestionsDatabase(),
+  ].filter(Boolean);
+
+  function onChange(tableId) {
+    const newJoin = join
+      .setJoinSourceTableId(tableId)
+      .setDefaultCondition()
+      .setDefaultAlias();
+    newJoin.parent().update(updateQuery);
+    onSourceTableSet(newJoin);
   }
+
+  return (
+    <NotebookCellItem color={color} icon="table2" inactive={!joinedTable}>
+      <DatabaseSchemaAndTableDataSelector
+        hasTableSearch
+        canChangeDatabase={false}
+        databases={databases}
+        tableFilter={table => table.db_id === query.database().id}
+        selectedDatabaseId={query.databaseId()}
+        selectedTableId={join.joinSourceTableId()}
+        setSourceTableFn={onChange}
+        isInitiallyOpen={join.joinSourceTableId() == null}
+        triggerElement={
+          joinedTable ? joinedTable.displayName() : t`Pick a table...`
+        }
+      />
+    </NotebookCellItem>
+  );
+}
+
+JoinTablePicker.propTypes = joinTablePickerPropTypes;
+
+const joinTypePickerPropTypes = {
+  join: PropTypes.object,
+  color: PropTypes.string,
+  updateQuery: PropTypes.func,
+};
+
+function JoinTypePicker({ join, color, updateQuery }) {
+  const strategyOption = join.strategyOption();
+
+  function onChange(strategy) {
+    join
+      .setStrategy(strategy)
+      .parent()
+      .update(updateQuery);
+  }
+
+  return (
+    <PopoverWithTrigger
+      triggerElement={
+        strategyOption ? (
+          <JoinStrategyIcon
+            tooltip={t`Change join type`}
+            name={strategyOption.icon}
+          />
+        ) : (
+          <NotebookCellItem color={color}>
+            {`Choose a join type`}
+          </NotebookCellItem>
+        )
+      }
+    >
+      {({ onClose }) => (
+        <JoinTypeSelect
+          value={strategyOption && strategyOption.value}
+          onChange={strategy => {
+            onChange(strategy);
+            onClose();
+          }}
+          options={join.strategyOptions()}
+        />
+      )}
+    </PopoverWithTrigger>
+  );
 }
 
+JoinTypePicker.propTypes = joinTypePickerPropTypes;
+
+const joinStrategyOptionShape = {
+  name: PropTypes.string.isRequired,
+  value: PropTypes.string.isRequired,
+  icon: PropTypes.string.isRequired,
+};
+
+const joinTypeSelectPropTypes = {
+  value: PropTypes.string,
+  onChange: PropTypes.func.isRequired,
+  options: PropTypes.arrayOf(PropTypes.shape(joinStrategyOptionShape))
+    .isRequired,
+};
+
 function JoinTypeSelect({ value, onChange, options }) {
   return (
-    <div className="px1 pt1">
+    <JoinTypeSelectRoot>
       {options.map(option => (
         <JoinTypeOption
           {...option}
@@ -229,32 +347,41 @@ function JoinTypeSelect({ value, onChange, options }) {
           onChange={onChange}
         />
       ))}
-    </div>
+    </JoinTypeSelectRoot>
   );
 }
 
+JoinTypeSelect.propTypes = joinTypeSelectPropTypes;
+
+const joinTypeOptionPropTypes = {
+  ...joinStrategyOptionShape,
+  selected: PropTypes.bool.isRequired,
+  onChange: PropTypes.func.isRequired,
+};
+
 function JoinTypeOption({ name, value, icon, selected, onChange }) {
   return (
-    <Flex
-      align="center"
-      className={cx(
-        "p1 mb1 rounded cursor-pointer text-white-hover bg-brand-hover",
-        {
-          "bg-brand text-white": selected,
-        },
-      )}
-      onClick={() => onChange(value)}
-    >
-      <Icon
-        className={cx("mr1", { "text-brand": !selected })}
-        name={icon}
-        size={24}
-      />
+    <JoinTypeOptionRoot isSelected={selected} onClick={() => onChange(value)}>
+      <JoinTypeIcon name={icon} isSelected={selected} />
       {name}
-    </Flex>
+    </JoinTypeOptionRoot>
   );
 }
 
+JoinTypeOption.propTypes = joinTypeOptionPropTypes;
+
+const joinDimensionPickerPropTypes = {
+  dimension: PropTypes.object.isRequired,
+  onChange: PropTypes.func.isRequired,
+  options: PropTypes.shape({
+    count: PropTypes.number.isRequired,
+    fks: PropTypes.array.isRequired,
+    dimensions: PropTypes.arrayOf(PropTypes.object).isRequired,
+  }).isRequired,
+  query: PropTypes.object.isRequired,
+  color: PropTypes.string,
+};
+
 class JoinDimensionPicker extends React.Component {
   open() {
     this._popover.open();
@@ -292,12 +419,50 @@ class JoinDimensionPicker extends React.Component {
   }
 }
 
-import FieldsPicker from "./FieldsPicker";
+JoinDimensionPicker.propTypes = joinDimensionPickerPropTypes;
+
+const joinFieldsPickerPropTypes = {
+  join: PropTypes.object.isRequired,
+  updateQuery: PropTypes.func.isRequired,
+  className: PropTypes.string,
+};
 
 const JoinFieldsPicker = ({ className, join, updateQuery }) => {
   const dimensions = join.joinedDimensions();
   const selectedDimensions = join.fieldsDimensions();
   const selected = new Set(selectedDimensions.map(d => d.key()));
+
+  function onSelectAll() {
+    join
+      .setFields("all")
+      .parent()
+      .update(updateQuery);
+  }
+
+  function onSelectNone() {
+    join
+      .setFields("none")
+      .parent()
+      .update(updateQuery);
+  }
+
+  function onToggleDimension(dimension) {
+    join
+      .setFields(
+        dimensions
+          .filter(d => {
+            if (d === dimension) {
+              return !selected.has(d.key());
+            } else {
+              return selected.has(d.key());
+            }
+          })
+          .map(d => d.mbql()),
+      )
+      .parent()
+      .update(updateQuery);
+  }
+
   return (
     <FieldsPicker
       className={className}
@@ -305,34 +470,11 @@ const JoinFieldsPicker = ({ className, join, updateQuery }) => {
       selectedDimensions={selectedDimensions}
       isAll={join.fields === "all"}
       isNone={join.fields === "none"}
-      onSelectAll={() =>
-        join
-          .setFields("all")
-          .parent()
-          .update(updateQuery)
-      }
-      onSelectNone={() =>
-        join
-          .setFields("none")
-          .parent()
-          .update(updateQuery)
-      }
-      onToggleDimension={(dimension, enable) => {
-        join
-          .setFields(
-            dimensions
-              .filter(d => {
-                if (d === dimension) {
-                  return !selected.has(d.key());
-                } else {
-                  return selected.has(d.key());
-                }
-              })
-              .map(d => d.mbql()),
-          )
-          .parent()
-          .update(updateQuery);
-      }}
+      onSelectAll={onSelectAll}
+      onSelectNone={onSelectNone}
+      onToggleDimension={onToggleDimension}
     />
   );
 };
+
+JoinFieldsPicker.propTypes = joinFieldsPickerPropTypes;

frontend/src/metabase/query_builder/components/notebook/steps/JoinStep.styled.js

+import styled from "styled-components";
+import { color } from "metabase/lib/colors";
+import { space } from "metabase/styled-components/theme";
+import Icon from "metabase/components/Icon";
+
+export const JoinClausesContainer = styled.div`
+  display: flex;
+  flex-direction: column;
+  flex: 1 0 auto;
+`;
+
+export const JoinClauseContainer = styled.div`
+  margin-bottom: ${props => (props.isLast ? 0 : "2px")};
+`;
+
+export const JoinClauseRoot = styled.div`
+  display: flex;
+  align-items: center;
+  flex: 1 1 auto;
+`;
+
+export const JoinStrategyIcon = styled(Icon).attrs({ size: 32 })`
+  color: ${color("brand")};
+  margin-right: ${space(1)};
+`;
+
+export const JoinTypeSelectRoot = styled.div`
+  margin: ${space(1)} ${space(1)} 0 ${space(1)};
+`;
+
+export const JoinTypeOptionRoot = styled.div`
+  display: flex;
+  align-items: center;
+  padding: ${space(1)};
+  margin-bottom: ${space(1)};
+  cursor: pointer;
+  border-radius: ${space(1)};
+
+  color: ${props => props.isSelected && color("text-white")};
+  background-color: ${props => props.isSelected && color("brand")};
+
+  :hover {
+    color: ${color("text-white")};
+    background-color: ${color("brand")};
+
+    .Icon {
+      color: ${color("text-white")};
+    }
+  }
+`;
+
+export const JoinTypeIcon = styled(Icon).attrs({ size: 24 })`
+  margin-right: ${space(1)};
+  color: ${props => (props.isSelected ? color("text-white") : color("brand"))};
+`;
+
+export const JoinedTableControlRoot = styled.div`
+  display: flex;
+  align-items: center;
+`;
+
+export const JoinWhereConditionLabel = styled.span.attrs({ children: "where" })`
+  color: ${color("text-medium")};
+  font-weight: bold;
+  margin-left: ${space(1)};
+  margin-right: ${space(2)};
+`;
+
+export const JoinOnConditionLabel = styled.span.attrs({ children: "=" })`
+  font-weight: bold;
+  color: ${color("text-medium")};
+  margin-right: 8px;
+`;
+
+export const RemoveJoinIcon = styled(Icon).attrs({ name: "close", size: 18 })`
+  cursor: pointer;
+  color: ${color("text-light")};
+
+  :hover {
+    color: ${color("text-medium")};
+  }
+`;