From 998dcb93c94cf8cf2cf752cd842ea54f8eae1a3e Mon Sep 17 00:00:00 2001
From: Ashton G <ashton.galloway@hpe.com>
Date: Tue, 29 Oct 2024 13:22:50 -0400
Subject: [PATCH] fix: improve sorting for multi-type columns (#10123)

---
 webui/react/src/components/MultiSortMenu.tsx | 63 ++++++++++++++++++--
 webui/react/src/pages/FlatRuns/FlatRuns.tsx  | 34 ++++++++---
 2 files changed, 84 insertions(+), 13 deletions(-)

diff --git a/webui/react/src/components/MultiSortMenu.tsx b/webui/react/src/components/MultiSortMenu.tsx
index 68916cb3fd6..1b677894876 100644
--- a/webui/react/src/components/MultiSortMenu.tsx
+++ b/webui/react/src/components/MultiSortMenu.tsx
@@ -1,12 +1,17 @@
+import Badge from 'hew/Badge';
 import Button from 'hew/Button';
 import { DirectionType, Sort, validSort } from 'hew/DataGrid/DataGrid';
 import Dropdown, { MenuItem } from 'hew/Dropdown';
 import Icon from 'hew/Icon';
 import Select, { Option } from 'hew/Select';
 import { Loadable } from 'hew/utils/loadable';
+import { groupBy, mapValues } from 'lodash';
+import { Fragment, useMemo } from 'react';
 
+import { runColumns } from 'pages/FlatRuns/columns';
 import { V1ColumnType } from 'services/api-ts-sdk';
 import { ProjectColumn } from 'types';
+import { removeColumnTypePrefix } from 'utils/flatRun';
 
 import css from './MultiSortMenu.module.scss';
 
@@ -25,6 +30,7 @@ interface MultiSortRowProps {
   onChange?: (sort: Sort) => void;
   onRemove?: () => void;
   bannedSortColumns: Set<string>;
+  typeMap: Record<string, V1ColumnType[]>;
 }
 interface DirectionOptionsProps {
   onChange?: (direction: DirectionType) => void;
@@ -36,6 +42,7 @@ interface ColumnOptionsProps {
   onChange?: (column: string) => void;
   value?: string;
   bannedSortColumns: Set<string>;
+  typeMap: Record<string, V1ColumnType[]>;
 }
 
 export const optionsByColumnType = {
@@ -153,6 +160,7 @@ const ColumnOptions: React.FC<ColumnOptionsProps> = ({
   columns,
   value,
   bannedSortColumns,
+  typeMap,
 }) => (
   <Select
     autoFocus
@@ -161,10 +169,31 @@ const ColumnOptions: React.FC<ColumnOptionsProps> = ({
     loading={Loadable.isNotLoaded(columns)}
     options={Loadable.getOrElse([], columns)
       .filter((c) => !bannedSortColumns.has(c.column))
-      .map((c) => ({
-        label: c.displayName || c.column,
-        value: c.column,
-      }))}
+      .map((c) => {
+        const badges = typeMap[c.column].map((type) => {
+          const copy =
+            (runColumns as readonly string[]).includes(c.column) &&
+            type === V1ColumnType.UNSPECIFIED
+              ? 'BOOLEAN'
+              : removeColumnTypePrefix(type);
+          return (
+            <Fragment key={type}>
+              <Badge text={copy} />{' '}
+            </Fragment>
+          );
+        });
+        const title = c.displayName || c.column;
+        const label = (
+          <>
+            {title} {badges}
+          </>
+        );
+        return {
+          label,
+          title,
+          value: c.column,
+        };
+      })}
     placeholder="Select column"
     value={value}
     width="100%"
@@ -178,6 +207,7 @@ const MultiSortRow: React.FC<MultiSortRowProps> = ({
   onChange,
   onRemove,
   bannedSortColumns,
+  typeMap,
 }) => {
   const valueType =
     Loadable.getOrElse([], columns).find((c) => c.column === sort.column)?.type ||
@@ -188,6 +218,7 @@ const MultiSortRow: React.FC<MultiSortRowProps> = ({
         <ColumnOptions
           bannedSortColumns={bannedSortColumns}
           columns={columns}
+          typeMap={typeMap}
           value={sort.column}
           onChange={(column) => onChange?.({ ...sort, column })}
         />
@@ -231,6 +262,27 @@ const MultiSort: React.FC<MultiSortProps> = ({ sorts, columns, onChange, bannedS
     window.document.body.dispatchEvent(new Event('mousedown', { bubbles: true }));
     onChange?.([EMPTY_SORT]);
   };
+  // replace duplicate columns with a single unspecified column for copy
+  // reasons. maintain types so we can display in the select
+  const [mergedColumns, typeMap] = useMemo(() => {
+    const loadableTuple = columns.map((c) => {
+      const columnGroups = groupBy(c, 'column');
+      const fixedColumns = Object.values(columnGroups).flatMap((group) => {
+        if (group.length > 1) {
+          return [
+            {
+              ...group[0],
+              type: V1ColumnType.UNSPECIFIED,
+            },
+          ];
+        }
+        return group;
+      });
+      const typeMap = mapValues(columnGroups, (group) => group.map((column) => column.type));
+      return [fixedColumns, typeMap] as const;
+    });
+    return [loadableTuple.map((l) => l[0]), loadableTuple.map((l) => l[1]).getOrElse({})];
+  }, [columns]);
 
   return (
     <div className={css.base} data-test-component="multiSort">
@@ -238,7 +290,7 @@ const MultiSort: React.FC<MultiSortProps> = ({ sorts, columns, onChange, bannedS
       <div className={css.rows} data-test="rows">
         {sorts.map((sort, idx) => {
           const seenColumns = sorts.slice(0, idx).map((s) => s.column);
-          const columnOptions = Loadable.map(columns, (cols) =>
+          const columnOptions = mergedColumns.map((cols) =>
             cols.filter((c) => !seenColumns.includes(c.column)),
           );
           return (
@@ -247,6 +299,7 @@ const MultiSort: React.FC<MultiSortProps> = ({ sorts, columns, onChange, bannedS
               columns={columnOptions}
               key={sort.column || idx}
               sort={sort}
+              typeMap={typeMap}
               onChange={makeOnRowChange(idx)}
               onRemove={makeOnRowRemove(idx)}
             />
diff --git a/webui/react/src/pages/FlatRuns/FlatRuns.tsx b/webui/react/src/pages/FlatRuns/FlatRuns.tsx
index 6e6d3d5fb90..af7a790a47f 100644
--- a/webui/react/src/pages/FlatRuns/FlatRuns.tsx
+++ b/webui/react/src/pages/FlatRuns/FlatRuns.tsx
@@ -22,6 +22,7 @@ import Pagination from 'hew/Pagination';
 import Row from 'hew/Row';
 import { useToast } from 'hew/Toast';
 import { Loadable, Loaded, NotLoaded } from 'hew/utils/loadable';
+import { groupBy, keyBy, mapValues } from 'lodash';
 import { useObservable } from 'micro-observables';
 import React, { useCallback, useEffect, useLayoutEffect, useMemo, useRef, useState } from 'react';
 import { v4 as uuidv4 } from 'uuid';
@@ -178,6 +179,7 @@ const FlatRuns: React.FC<Props> = ({ projectId, workspaceId, searchId }) => {
   const [runs, setRuns] = useState<Loadable<FlatRun>[]>(INITIAL_LOADING_RUNS);
 
   const [sorts, setSorts] = useState<Sort[]>([EMPTY_SORT]);
+
   const sortString = useMemo(() => makeSortString(sorts.filter(validSort.is)), [sorts]);
   const loadableFormset = useObservable(formStore.formset);
   const rootFilterChildren: Array<FormGroup | FormField> = Loadable.match(loadableFormset, {
@@ -237,6 +239,24 @@ const FlatRuns: React.FC<Props> = ({ projectId, workspaceId, searchId }) => {
     }
   }, [projectId]);
 
+  // expand sorts to include types from metadata columns so the column ids match
+  const expandedSorts = useMemo(() => {
+    return projectColumns.match({
+      _: () => sorts,
+      Loaded(pc) {
+        const groupedColumns = groupBy(pc, (c) => c.column);
+        const columnToExpandedColumns = mapValues(groupedColumns, (cols) =>
+          cols.map((c) => formatColumnKey(c)),
+        );
+        return sorts
+          .filter(validSort.is)
+          .flatMap((sort) =>
+            columnToExpandedColumns[sort.column].map((ec) => ({ ...sort, column: ec })),
+          );
+      },
+    });
+  }, [sorts, projectColumns]);
+
   const arrayTypeColumns = useMemo(() => {
     const arrayTypeColumns = projectColumns
       .getOrElse([])
@@ -292,13 +312,11 @@ const FlatRuns: React.FC<Props> = ({ projectId, workspaceId, searchId }) => {
     return [...STATIC_COLUMNS, ...settings.columns.slice(0, settings.pinnedColumnsCount)];
   }, [settings.columns, settings.pinnedColumnsCount]);
 
+  const projectColumnsMap = useMemo(() => {
+    return projectColumns.map((columns) => keyBy(columns, formatColumnKey));
+  }, [projectColumns]);
+
   const columns: ColumnDef<FlatRun>[] = useMemo(() => {
-    const projectColumnsMap: Loadable<Record<string, ProjectColumn>> = Loadable.map(
-      projectColumns,
-      (columns) => {
-        return columns.reduce((acc, col) => ({ ...acc, [formatColumnKey(col)]: col }), {});
-      },
-    );
     const columnDefs = getColumnDefs({
       appTheme,
       columnWidths: settings.columnWidths,
@@ -455,7 +473,7 @@ const FlatRuns: React.FC<Props> = ({ projectId, workspaceId, searchId }) => {
     appTheme,
     columnsIfLoaded,
     isDarkMode,
-    projectColumns,
+    projectColumnsMap,
     projectHeatmap,
     dataGridSelection.rows,
     settings.columnWidths,
@@ -956,7 +974,7 @@ const FlatRuns: React.FC<Props> = ({ projectId, workspaceId, searchId }) => {
     [
       bannedFilterColumns,
       bannedSortColumns,
-      projectColumns,
+      projectColumnsMap,
       settings.pinnedColumnsCount,
       settings.heatmapOn,
       settings.heatmapSkipped,