From 79f8079a1021c50d9147bbd036813462f7b0d9c1 Mon Sep 17 00:00:00 2001
From: Nar Saynorath <nar.saynorath@sentry.io>
Date: Fri, 3 Jan 2025 16:56:42 -0500
Subject: [PATCH] fix(widget-builder): Update table sort when field removed or
 modified

---
 .../widgetBuilder/components/visualize.tsx    | 12 +++--
 .../hooks/useWidgetBuilderState.spec.tsx      | 49 +++++++++++++++++++
 .../hooks/useWidgetBuilderState.tsx           | 39 +++++++++++++++
 3 files changed, 95 insertions(+), 5 deletions(-)

diff --git a/static/app/views/dashboards/widgetBuilder/components/visualize.tsx b/static/app/views/dashboards/widgetBuilder/components/visualize.tsx
index dd77bf27e54831..29d6e80bb5ee63 100644
--- a/static/app/views/dashboards/widgetBuilder/components/visualize.tsx
+++ b/static/app/views/dashboards/widgetBuilder/components/visualize.tsx
@@ -260,16 +260,18 @@ function Visualize() {
                             : field.field
                         }
                         onChange={newField => {
+                          const newFields = cloneDeep(fields);
+                          const currentField = newFields[index]!;
                           // Update the current field's aggregate with the new aggregate
-                          if (field.kind === FieldValueKind.FUNCTION) {
-                            field.function[1] = newField.value as string;
+                          if (currentField.kind === FieldValueKind.FUNCTION) {
+                            currentField.function[1] = newField.value as string;
                           }
-                          if (field.kind === FieldValueKind.FIELD) {
-                            field.field = newField.value as string;
+                          if (currentField.kind === FieldValueKind.FIELD) {
+                            currentField.field = newField.value as string;
                           }
                           dispatch({
                             type: updateAction,
-                            payload: fields,
+                            payload: newFields,
                           });
                         }}
                         triggerProps={{
diff --git a/static/app/views/dashboards/widgetBuilder/hooks/useWidgetBuilderState.spec.tsx b/static/app/views/dashboards/widgetBuilder/hooks/useWidgetBuilderState.spec.tsx
index 522596e65243bb..82afad19c711cc 100644
--- a/static/app/views/dashboards/widgetBuilder/hooks/useWidgetBuilderState.spec.tsx
+++ b/static/app/views/dashboards/widgetBuilder/hooks/useWidgetBuilderState.spec.tsx
@@ -827,6 +827,55 @@ describe('useWidgetBuilderState', () => {
         'event.type',
       ]);
     });
+
+    it('resets the sort when the field that is being sorted is removed', () => {
+      mockedUsedLocation.mockReturnValue(
+        LocationFixture({
+          query: {field: ['testField'], sort: ['-testField']},
+        })
+      );
+
+      const {result} = renderHook(() => useWidgetBuilderState(), {
+        wrapper: WidgetBuilderProvider,
+      });
+
+      expect(result.current.state.sort).toEqual([{field: 'testField', kind: 'desc'}]);
+
+      act(() => {
+        result.current.dispatch({
+          type: BuilderStateAction.SET_FIELDS,
+          payload: [{field: 'testField2', kind: FieldValueKind.FIELD}],
+        });
+      });
+
+      expect(result.current.state.sort).toEqual([{field: 'testField2', kind: 'desc'}]);
+    });
+
+    it('modifies the sort when the field that is being sorted is modified', () => {
+      mockedUsedLocation.mockReturnValue(
+        LocationFixture({
+          query: {field: ['testField', 'sortField'], sort: ['-sortField']},
+        })
+      );
+
+      const {result} = renderHook(() => useWidgetBuilderState(), {
+        wrapper: WidgetBuilderProvider,
+      });
+
+      expect(result.current.state.sort).toEqual([{field: 'sortField', kind: 'desc'}]);
+
+      act(() => {
+        result.current.dispatch({
+          type: BuilderStateAction.SET_FIELDS,
+          payload: [
+            {field: 'testField', kind: FieldValueKind.FIELD},
+            {field: 'newSortField', kind: FieldValueKind.FIELD},
+          ],
+        });
+      });
+
+      expect(result.current.state.sort).toEqual([{field: 'newSortField', kind: 'desc'}]);
+    });
   });
 
   describe('yAxis', () => {
diff --git a/static/app/views/dashboards/widgetBuilder/hooks/useWidgetBuilderState.tsx b/static/app/views/dashboards/widgetBuilder/hooks/useWidgetBuilderState.tsx
index 400ab8c71e7e2e..a13c5e9187b0e9 100644
--- a/static/app/views/dashboards/widgetBuilder/hooks/useWidgetBuilderState.tsx
+++ b/static/app/views/dashboards/widgetBuilder/hooks/useWidgetBuilderState.tsx
@@ -233,6 +233,45 @@ function useWidgetBuilderState(): {
           break;
         case BuilderStateAction.SET_FIELDS:
           setFields(action.payload);
+          const isRemoved = action.payload.length < (fields?.length ?? 0);
+          if (
+            displayType === DisplayType.TABLE &&
+            action.payload.length > 0 &&
+            !action.payload.find(
+              field => generateFieldAsString(field) === sort?.[0]?.field
+            )
+          ) {
+            if (isRemoved) {
+              setSort([
+                {
+                  kind: 'desc',
+                  field: generateFieldAsString(action.payload[0] as QueryFieldValue),
+                },
+              ]);
+            } else {
+              // Find the index of the first field that doesn't match the old fields.
+              const changedFieldIndex = action.payload.findIndex(
+                field =>
+                  !fields?.find(
+                    originalField =>
+                      generateFieldAsString(originalField) ===
+                      generateFieldAsString(field)
+                  )
+              );
+              if (changedFieldIndex !== -1) {
+                // At this point, we can assume the fields are the same length so
+                // using the changedFieldIndex in action.payload is safe.
+                setSort([
+                  {
+                    kind: sort?.[0]?.kind ?? 'desc',
+                    field: generateFieldAsString(
+                      action.payload[changedFieldIndex] as QueryFieldValue
+                    ),
+                  },
+                ]);
+              }
+            }
+          }
           break;
         case BuilderStateAction.SET_Y_AXIS:
           setYAxis(action.payload);