feat(widget-builder): Cache builder state between dataset changes

static/app/views/dashboards/widgetBuilder/components/datasetSelector.tsx

@@ -11,15 +11,16 @@ import useOrganization from 'sentry/utils/useOrganization';
 import {WidgetType} from 'sentry/views/dashboards/types';
 import {SectionHeader} from 'sentry/views/dashboards/widgetBuilder/components/common/sectionHeader';
 import {useWidgetBuilderContext} from 'sentry/views/dashboards/widgetBuilder/contexts/widgetBuilderContext';
+import {useCacheBuilderState} from 'sentry/views/dashboards/widgetBuilder/hooks/useCacheBuilderState';
 import useDashboardWidgetSource from 'sentry/views/dashboards/widgetBuilder/hooks/useDashboardWidgetSource';
 import useIsEditingWidget from 'sentry/views/dashboards/widgetBuilder/hooks/useIsEditingWidget';
-import {BuilderStateAction} from 'sentry/views/dashboards/widgetBuilder/hooks/useWidgetBuilderState';
 
 function WidgetBuilderDatasetSelector() {
   const organization = useOrganization();
-  const {state, dispatch} = useWidgetBuilderContext();
+  const {state} = useWidgetBuilderContext();
   const source = useDashboardWidgetSource();
   const isEditing = useIsEditingWidget();
+  const {cacheBuilderState, restoreOrSetBuilderState} = useCacheBuilderState();
 
   const datasetChoices: Array<RadioOption<WidgetType>> = [];
   datasetChoices.push([WidgetType.ERRORS, t('Errors')]);
@@ -48,17 +49,21 @@ function WidgetBuilderDatasetSelector() {
         label={t('Dataset')}
         value={state.dataset ?? WidgetType.ERRORS}
         choices={datasetChoices}
-        onChange={(newValue: WidgetType) => {
-          dispatch({
-            type: BuilderStateAction.SET_DATASET,
-            payload: newValue,
-          });
+        onChange={(newDataset: WidgetType) => {
+          // Set the current dataset state in local storage for recovery
+          // when the user navigates back to this dataset
+          cacheBuilderState(state.dataset ?? WidgetType.ERRORS);
+
+          // Restore the builder state for the new dataset
+          // or set the dataset if there is no cached state
+          restoreOrSetBuilderState(newDataset);
+
           trackAnalytics('dashboards_views.widget_builder.change', {
             from: source,
             widget_type: state.dataset ?? '',
             builder_version: WidgetBuilderVersion.SLIDEOUT,
             field: 'dataSet',
-            value: newValue,
+            value: newDataset,
             new_widget: !isEditing,
             organization,
           });

static/app/views/dashboards/widgetBuilder/hooks/useCacheBuilderState.spec.tsx

@@ -0,0 +1,204 @@
+import {LocationFixture} from 'sentry-fixture/locationFixture';
+
+import {renderHook} from 'sentry-test/reactTestingLibrary';
+
+import {useLocation} from 'sentry/utils/useLocation';
+import {DisplayType, WidgetType} from 'sentry/views/dashboards/types';
+import {
+  useWidgetBuilderContext,
+  WidgetBuilderProvider,
+} from 'sentry/views/dashboards/widgetBuilder/contexts/widgetBuilderContext';
+import {
+  BuilderStateAction,
+  type WidgetBuilderState,
+} from 'sentry/views/dashboards/widgetBuilder/hooks/useWidgetBuilderState';
+import {convertBuilderStateToWidget} from 'sentry/views/dashboards/widgetBuilder/utils/convertBuilderStateToWidget';
+
+import {useCacheBuilderState} from './useCacheBuilderState';
+
+jest.mock('sentry/utils/useNavigate', () => ({
+  useNavigate: jest.fn(),
+}));
+jest.mock('sentry/utils/useLocation');
+
+jest.mock('sentry/views/dashboards/widgetBuilder/contexts/widgetBuilderContext', () => ({
+  useWidgetBuilderContext: jest.fn(),
+  WidgetBuilderProvider: jest.requireActual(
+    'sentry/views/dashboards/widgetBuilder/contexts/widgetBuilderContext'
+  ).WidgetBuilderProvider,
+}));
+
+const mockUseWidgetBuilderContext = jest.mocked(useWidgetBuilderContext);
+const mockUseLocation = jest.mocked(useLocation);
+
+function Wrapper({children}: {children: React.ReactNode}) {
+  return <WidgetBuilderProvider>{children}</WidgetBuilderProvider>;
+}
+
+describe('useCacheBuilderState', () => {
+  let mockLocalStorage: Record<string, string>;
+
+  beforeEach(() => {
+    mockLocalStorage = {};
+    mockUseWidgetBuilderContext.mockReturnValue({
+      state: {},
+      dispatch: jest.fn(),
+    });
+    mockUseLocation.mockReturnValue(LocationFixture());
+
+    Storage.prototype.getItem = jest.fn(key => mockLocalStorage[key] ?? null);
+    Storage.prototype.setItem = jest.fn((key, value) => {
+      mockLocalStorage[key] = value;
+    });
+    Storage.prototype.removeItem = jest.fn(key => {
+      delete mockLocalStorage[key];
+    });
+  });
+
+  afterEach(() => {
+    jest.restoreAllMocks();
+  });
+
+  it('caches builder state to localStorage', () => {
+    const cachedWidget: WidgetBuilderState = {
+      dataset: WidgetType.ERRORS,
+      displayType: DisplayType.LINE,
+      yAxis: [
+        {
+          function: ['count', '', undefined, undefined],
+          kind: 'function',
+        },
+      ],
+      query: ['this is a test query'],
+    };
+    mockUseWidgetBuilderContext.mockReturnValue({
+      state: cachedWidget,
+      dispatch: jest.fn(),
+    });
+
+    const {result} = renderHook(() => useCacheBuilderState(), {
+      wrapper: Wrapper,
+    });
+
+    result.current.cacheBuilderState(WidgetType.ERRORS);
+
+    // Verify state was saved to localStorage
+    expect(localStorage.setItem).toHaveBeenCalledWith(
+      'dashboards:widget-builder:dataset:error-events',
+      JSON.stringify(convertBuilderStateToWidget(cachedWidget))
+    );
+
+    result.current.restoreOrSetBuilderState(WidgetType.ERRORS);
+
+    expect(localStorage.getItem).toHaveBeenCalledWith(
+      'dashboards:widget-builder:dataset:error-events'
+    );
+  });
+
+  it('restores builder state from localStorage when available', () => {
+    const cachedWidget: WidgetBuilderState = {
+      title: 'error widget title',
+      description: 'error widget description',
+      dataset: WidgetType.ERRORS,
+      displayType: DisplayType.LINE,
+      yAxis: [
+        {
+          function: ['count', '', undefined, undefined],
+          kind: 'function',
+        },
+      ],
+      query: ['this is a test query'],
+    };
+    const currentWidget: WidgetBuilderState = {
+      title: 'issue widget title',
+      description: 'issue widget description',
+      dataset: WidgetType.ISSUE,
+      displayType: DisplayType.TABLE,
+      query: ['issue.id:123'],
+      fields: [
+        {
+          field: 'issue',
+          kind: 'field',
+        },
+      ],
+    };
+    const mockDispatch = jest.fn();
+    mockUseWidgetBuilderContext.mockReturnValue({
+      state: currentWidget,
+      dispatch: mockDispatch,
+    });
+    // Add cached widget to the localStorage
+    localStorage.setItem(
+      'dashboards:widget-builder:dataset:error-events',
+      JSON.stringify(convertBuilderStateToWidget(cachedWidget))
+    );
+
+    const {result} = renderHook(() => useCacheBuilderState(), {
+      wrapper: Wrapper,
+    });
+
+    // Call the restore helper on the cached dataset
+    result.current.restoreOrSetBuilderState(WidgetType.ERRORS);
+
+    expect(mockDispatch).toHaveBeenCalledWith({
+      type: BuilderStateAction.SET_STATE,
+
+      // the yAxis gets converted to a string when used with this payload
+      payload: expect.objectContaining({
+        ...cachedWidget,
+        yAxis: ['count()'],
+        title: 'issue widget title', // The title was not overridden
+        description: 'issue widget description', // The description was not overridden
+      }),
+    });
+  });
+
+  it('plainly sets the new dataset when no cached state exists', () => {
+    const cachedWidget: WidgetBuilderState = {
+      dataset: WidgetType.ERRORS,
+      displayType: DisplayType.LINE,
+      yAxis: [
+        {
+          function: ['count', '', undefined, undefined],
+          kind: 'function',
+        },
+      ],
+      query: ['this is a test query'],
+    };
+    const currentWidget: WidgetBuilderState = {
+      dataset: WidgetType.ISSUE,
+      displayType: DisplayType.TABLE,
+      query: ['issue.id:123'],
+      fields: [
+        {
+          field: 'issue',
+          kind: 'field',
+        },
+      ],
+    };
+    const mockDispatch = jest.fn();
+    mockUseWidgetBuilderContext.mockReturnValue({
+      state: currentWidget,
+      dispatch: mockDispatch,
+    });
+    // Add cached widget to the localStorage, this will not be the one
+    // used in the test to test that a cache miss falls back to the plain
+    // dataset change
+    localStorage.setItem(
+      'dashboards:widget-builder:dataset:error-events',
+      JSON.stringify(convertBuilderStateToWidget(cachedWidget))
+    );
+
+    const {result} = renderHook(() => useCacheBuilderState(), {
+      wrapper: Wrapper,
+    });
+
+    // Call the restore helper on the cached dataset
+    result.current.restoreOrSetBuilderState(WidgetType.TRANSACTIONS);
+
+    expect(mockDispatch).toHaveBeenCalledWith({
+      type: BuilderStateAction.SET_DATASET,
+      payload: WidgetType.TRANSACTIONS,
+    });
+  });
+});

static/app/views/dashboards/widgetBuilder/hooks/useCacheBuilderState.tsx

@@ -0,0 +1,70 @@
+import {useCallback, useEffect} from 'react';
+
+import type {WidgetType} from 'sentry/views/dashboards/types';
+import {useWidgetBuilderContext} from 'sentry/views/dashboards/widgetBuilder/contexts/widgetBuilderContext';
+import {BuilderStateAction} from 'sentry/views/dashboards/widgetBuilder/hooks/useWidgetBuilderState';
+import {convertBuilderStateToWidget} from 'sentry/views/dashboards/widgetBuilder/utils/convertBuilderStateToWidget';
+import {convertWidgetToBuilderStateParams} from 'sentry/views/dashboards/widgetBuilder/utils/convertWidgetToBuilderStateParams';
+
+const WIDGET_BUILDER_DATASET_STATE_KEY = 'dashboards:widget-builder:dataset';
+
+function cleanUpDatasetState() {
+  for (let i = 0; i < localStorage.length; i++) {
+    const key = localStorage.key(i);
+    if (key?.startsWith(WIDGET_BUILDER_DATASET_STATE_KEY)) {
+      localStorage.removeItem(key);
+    }
+  }
+}
+
+/**
+ * This hook is used to cache the builder state for the given dataset
+ * and restore it when the user navigates back to the same dataset.
+ */
+export function useCacheBuilderState() {
+  const {state, dispatch} = useWidgetBuilderContext();
+
+  useEffect(() => {
+    return cleanUpDatasetState;
+  }, []);
+
+  const cacheBuilderState = useCallback(
+    (dataset: WidgetType) => {
+      localStorage.setItem(
+        `${WIDGET_BUILDER_DATASET_STATE_KEY}:${dataset}`,
+        JSON.stringify(convertBuilderStateToWidget(state))
+      );
+    },
+    [state]
+  );
+
+  // Checks if there is a cached builder state for the given dataset
+  // and restores it if it exists. Otherwise, it sets the dataset.
+  const restoreOrSetBuilderState = useCallback(
+    (nextDataset: WidgetType) => {
+      const previousDatasetState = localStorage.getItem(
+        `${WIDGET_BUILDER_DATASET_STATE_KEY}:${nextDataset}`
+      );
+      if (previousDatasetState) {
+        const builderState = convertWidgetToBuilderStateParams(
+          JSON.parse(previousDatasetState)
+        );
+        dispatch({
+          type: BuilderStateAction.SET_STATE,
+          payload: {...builderState, title: state.title, description: state.description},
+        });
+      } else {
+        dispatch({
+          type: BuilderStateAction.SET_DATASET,
+          payload: nextDataset,
+        });
+      }
+    },
+    [dispatch, state.title, state.description]
+  );
+
+  return {
+    cacheBuilderState,
+    restoreOrSetBuilderState,
+  };
+}