Add multiple branches support inside the experiments table

extension/src/experiments/webview/contract.ts

@@ -43,6 +43,7 @@ export type Experiment = {
   starred?: boolean
   status?: ExperimentStatus
   timestamp?: string | null
+  branch?: string
 }
 
 export const isRunning = (status: ExperimentStatus | undefined): boolean =>

extension/src/test/e2e/extension.test.ts

@@ -52,7 +52,9 @@ describe('Experiments Table Webview', function () {
   const workspaceRow = 1
   const commitRows = 3
   const previousCommitRow = 1
-  const initialRows = headerRows + workspaceRow + commitRows + previousCommitRow
+  const actionsRow = 1
+  const initialRows =
+    headerRows + workspaceRow + commitRows + previousCommitRow + actionsRow
 
   it('should load as an editor', async function () {
     const workbench = await browser.getWorkbench()

extension/src/test/fixtures/expShow/base/rows.ts

@@ -15,8 +15,9 @@ const valueWithNoChanges = (str: string) => ({
 
 const colorsList = copyOriginalColors()
 
-const data: Commit[] = [
+export const rowsFixtureWithBranches: Commit[] = [
   {
+    branch: 'current',
     deps: {
       [join('data', 'data.xml')]: valueWithNoChanges(
         '22a1a2931c8370d3aeedd7183606fd7f'
@@ -72,6 +73,7 @@ const data: Commit[] = [
     starred: false
   },
   {
+    branch: 'current',
     deps: {
       [join('data', 'data.xml')]: valueWithNoChanges(
         '22a1a2931c8370d3aeedd7183606fd7f'
@@ -126,6 +128,7 @@ const data: Commit[] = [
     starred: false,
     subRows: [
       {
+        branch: 'current',
         deps: {
           [join('data', 'data.xml')]: valueWithNoChanges(
             '22a1a2931c8370d3aeedd7183606fd7f'
@@ -184,6 +187,7 @@ const data: Commit[] = [
         Created: '2020-12-29T15:31:52'
       },
       {
+        branch: 'current',
         deps: {
           [join('data', 'data.xml')]: valueWithNoChanges(
             '22a1a2931c8370d3aeedd7183606fd7f'
@@ -240,6 +244,7 @@ const data: Commit[] = [
         Created: '2020-12-29T15:28:59'
       },
       {
+        branch: 'current',
         deps: {
           [join('data', 'data.xml')]: valueWithNoChanges(
             '22a1a2931c8370d3aeedd7183606fd7f'
@@ -297,6 +302,7 @@ const data: Commit[] = [
         Created: '2020-12-29T15:27:02'
       },
       {
+        branch: 'current',
         displayColor: undefined,
         id: '489fd8b',
         sha: '489fd8bdaa709f7330aac342e051a9431c625481',
@@ -308,6 +314,7 @@ const data: Commit[] = [
         status: ExperimentStatus.FAILED
       },
       {
+        branch: 'current',
         deps: {
           [join('data', 'data.xml')]: valueWithNoChanges(
             '22a1a2931c8370d3aeedd7183606fd7f'
@@ -359,6 +366,7 @@ const data: Commit[] = [
         Created: '2020-12-29T15:26:36'
       },
       {
+        branch: 'current',
         displayColor: undefined,
         deps: {
           [join('data', 'data.xml')]: valueWithNoChanges(
@@ -407,6 +415,7 @@ const data: Commit[] = [
         Created: '2020-12-29T15:25:27'
       },
       {
+        branch: 'current',
         displayColor: undefined,
         deps: {
           [join('data', 'data.xml')]: valueWithNoChanges(
@@ -461,4 +470,15 @@ const data: Commit[] = [
   }
 ]
 
-export default data
+const rowsFixtureWithoutBranches = [
+  ...rowsFixtureWithBranches.map(({ branch, ...row }) =>
+    row.subRows
+      ? {
+          ...row,
+          subRows: row.subRows?.map(({ branch, ...subRow }) => subRow)
+        }
+      : row
+  )
+]
+
+export default rowsFixtureWithoutBranches

extension/src/test/fixtures/expShow/base/tableData.ts

@@ -1,5 +1,5 @@
 import { TableData } from '../../../../experiments/webview/contract'
-import rowsFixture from './rows'
+import { rowsFixtureWithBranches } from './rows'
 import columnsFixture from './columns'
 
 const tableDataFixture: TableData = {
@@ -18,7 +18,7 @@ const tableDataFixture: TableData = {
   hasValidDvcYaml: true,
   isShowingMoreCommits: true,
   isBranchesView: false,
-  rows: rowsFixture,
+  rows: rowsFixtureWithBranches,
   selectedForPlotsCount: 2,
   sorts: []
 }

webview/src/experiments/components/App.tsx

@@ -91,7 +91,19 @@ export const App: React.FC<Record<string, unknown>> = () => {
                 )
                 continue
               case 'rows':
-                dispatch(updateRows(data.data.rows))
+                dispatch(
+                  updateRows(
+                    // Setting any branch for now just so that it isn't undefined. It does not matter the label for now as it is not shown
+                    data.data.rows.map(row => ({
+                      ...row,
+                      branch: 'current',
+                      subRows: row.subRows?.map(subRow => ({
+                        ...subRow,
+                        branch: 'current'
+                      }))
+                    }))
+                  )
+                )
                 continue
               case 'selectedForPlotsCount':
                 dispatch(

webview/src/experiments/components/table/Indicators.tsx

@@ -5,16 +5,19 @@ import { CellHintTooltip } from './body/CellHintTooltip'
 import {
   focusFiltersTree,
   focusSortsTree,
-  openPlotsWebview
+  openPlotsWebview,
+  selectBranches
 } from '../../util/messages'
 import { Icon } from '../../../shared/components/Icon'
 import {
   Filter,
+  GitMerge,
   GraphScatter,
   SortPrecedence
 } from '../../../shared/components/icons'
 import { pluralize } from '../../../util/strings'
 import { ExperimentsState } from '../../store'
+import { featureFlag } from '../../../util/flags'
 
 export type CounterBadgeProps = {
   count?: number
@@ -81,6 +84,10 @@ export const Indicators = () => {
   const selectedForPlotsCount = useSelector(
     (state: ExperimentsState) => state.tableData.selectedForPlotsCount
   )
+  const branchesSelected = useSelector(
+    (state: ExperimentsState) =>
+      Math.max(state.tableData.branches.length - 1, 0) // We always have one branch by default (the current one which is not selected)
+  )
 
   const sortsCount = sorts?.length
   const filtersCount = filters?.length
@@ -124,6 +131,20 @@ export const Indicators = () => {
       >
         <Icon width={16} height={16} icon={Filter} />
       </Indicator>
+      {featureFlag.ADD_REMOVE_BRANCHES && (
+        <Indicator
+          count={branchesSelected}
+          aria-label="branches"
+          onClick={selectBranches}
+          tooltipContent={formatCountMessage(
+            'Branches',
+            branchesSelected,
+            'Selected'
+          )}
+        >
+          <Icon width={16} height={16} icon={GitMerge} />
+        </Indicator>
+      )}
     </div>
   )
 }

webview/src/experiments/components/table/Table.test.tsx

@@ -27,6 +27,7 @@ import { dragAndDrop, dragEnter, dragLeave } from '../../../test/dragDrop'
 import { DragEnterDirection } from '../../../shared/components/dragDrop/util'
 import { experimentsReducers } from '../../store'
 import { customQueries } from '../../../test/queries'
+import { TableDataState } from '../../state/tableDataSlice'
 
 jest.mock('../../../shared/api')
 
@@ -39,7 +40,8 @@ describe('Table', () => {
   ) => {
     const tableData = {
       ...sortingTableDataFixture,
-      ...partialTableData
+      ...partialTableData,
+      branches: ['current']
     }
     return render(
       <Provider
@@ -257,8 +259,9 @@ describe('Table', () => {
         id: 333
       }
 
-      const tableDataWithColumnSetting: TableData = {
+      const tableDataWithColumnSetting: TableDataState = {
         ...sortingTableDataFixture,
+        branches: ['current'],
         columnWidths
       }
       render(
@@ -297,8 +300,9 @@ describe('Table', () => {
         id: 333
       }
 
-      const tableDataWithColumnSetting: TableData = {
+      const tableDataWithColumnSetting: TableDataState = {
         ...sortingTableDataFixture,
+        branches: ['current'],
         columnWidths
       }
       render(

webview/src/experiments/components/table/Table.tsx

@@ -5,7 +5,6 @@ import styles from './styles.module.scss'
 import { TableHead } from './header/TableHead'
 import { RowSelectionContext } from './RowSelectionContext'
 import { Indicators } from './Indicators'
-import { CommitsAndBranchesNavigation } from './commitsAndBranches/CommitsAndBranchesNavigation'
 import { TableContent } from './body/TableContent'
 import { InstanceProp } from '../../util/interfaces'
 
@@ -54,7 +53,6 @@ export const Table: React.FC<TableProps> = ({
           tableHeadHeight={tableHeadHeight}
         />
       </table>
-      <CommitsAndBranchesNavigation />
       <Indicators />
     </div>
   )

webview/src/experiments/components/table/body/PreviousCommitsRow.tsx

@@ -0,0 +1,22 @@
+import cx from 'classnames'
+import React from 'react'
+import styles from '../styles.module.scss'
+
+interface PreviousCommitsRowProps {
+  isBranchesView?: boolean
+  nbColumns: number
+}
+
+export const PreviousCommitsRow: React.FC<PreviousCommitsRowProps> = ({
+  isBranchesView,
+  nbColumns
+}) => (
+  <thead>
+    <tr className={cx(styles.previousCommitsRow)}>
+      <th className={cx(styles.previousCommitsText, styles.experimentsTd)}>
+        {isBranchesView ? 'Other Branches' : 'Previous Commits'}
+      </th>
+      <th colSpan={nbColumns - 1}></th>
+    </tr>
+  </thead>
+)

webview/src/experiments/components/table/body/TableBody.tsx

@@ -5,26 +5,29 @@ import { EXPERIMENT_WORKSPACE_ID } from 'dvc/src/cli/dvc/contract'
 import { ExperimentGroup } from './ExperimentGroup'
 import { BatchSelectionProp, RowContent } from './Row'
 import { WorkspaceRowGroup } from './WorkspaceRowGroup'
+import { PreviousCommitsRow } from './PreviousCommitsRow'
 import styles from '../styles.module.scss'
 import { InstanceProp, RowProp } from '../../../util/interfaces'
 import { ExperimentsState } from '../../../store'
 
-export const TableBody: React.FC<
-  RowProp &
-    InstanceProp &
-    BatchSelectionProp & {
-      root: HTMLElement | null
-      tableHeaderHeight: number
-    }
-> = ({
+interface TableBodyProps extends RowProp, InstanceProp, BatchSelectionProp {
+  root: HTMLElement | null
+  tableHeaderHeight: number
+  showPreviousRow?: boolean
+  isLast?: boolean
+}
+
+export const TableBody: React.FC<TableBodyProps> = ({
   row,
   instance,
   contextMenuDisabled,
   projectHasCheckpoints,
   hasRunningExperiment,
   batchRowSelection,
   root,
-  tableHeaderHeight
+  tableHeaderHeight,
+  showPreviousRow,
+  isLast
 }) => {
   const contentProps = {
     batchRowSelection,
@@ -54,22 +57,17 @@ export const TableBody: React.FC<
     </WorkspaceRowGroup>
   ) : (
     <>
-      {row.index === 2 && row.depth === 0 && (
-        <tbody>
-          <tr className={cx(styles.previousCommitsRow, styles.experimentsTr)}>
-            <td
-              className={cx(styles.previousCommitsText, styles.experimentsTd)}
-            >
-              {isBranchesView ? 'Other Branches' : 'Previous Commits'}
-            </td>
-            <td colSpan={row.getAllCells().length - 1}></td>
-          </tr>
-        </tbody>
+      {showPreviousRow && row.depth === 0 && (
+        <PreviousCommitsRow
+          isBranchesView={isBranchesView}
+          nbColumns={row.getAllCells().length}
+        />
       )}
       <tbody
         className={cx(styles.rowGroup, {
           [styles.experimentGroup]: row.depth > 0,
-          [styles.expandedGroup]: row.getIsExpanded() && row.subRows.length > 0
+          [styles.expandedGroup]: row.getIsExpanded() && row.subRows.length > 0,
+          [styles.lastRowGroup]: isLast
         })}
       >
         {content}