Add push behavior to useQueryPersistedState

js_modules/dagster-ui/packages/ui-core/src/asset-selection/useAssetSelectionState.oss.tsx

@@ -4,5 +4,6 @@ export function useAssetSelectionState() {
   return useQueryPersistedState<string>({
     queryKey: 'asset-selection',
     defaults: {['asset-selection']: ''},
+    behavior: 'push',
   });
 }

js_modules/dagster-ui/packages/ui-core/src/hooks/__tests__/useQueryPersistedState.test.tsx

@@ -1,7 +1,7 @@
-import {render, screen, waitFor} from '@testing-library/react';
+import {act, render, screen, waitFor} from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
 import {useCallback, useMemo} from 'react';
-import {MemoryRouter} from 'react-router-dom';
+import {MemoryRouter, useHistory} from 'react-router-dom';
 
 import {Route} from '../../app/Route';
 import {useQueryPersistedState} from '../useQueryPersistedState';
@@ -348,4 +348,112 @@ describe('useQueryPersistedState', () => {
       );
     });
   });
+
+  it('supports push behavior', async () => {
+    let querySearch: string | undefined;
+
+    let goback: () => void;
+
+    const Test = ({options}: {options: Parameters<typeof useQueryPersistedState>[0]}) => {
+      const [_, setQuery] = useQueryPersistedState(options);
+      const history = useHistory();
+      goback = () => history.goBack();
+      return (
+        <>
+          <div onClick={() => setQuery('one')}>one</div>
+          <div onClick={() => setQuery('two')}>two</div>
+          <div onClick={() => setQuery('three')}>three</div>
+        </>
+      );
+    };
+
+    render(
+      <MemoryRouter initialEntries={['/page?q=B']}>
+        <Test options={{queryKey: 'q', behavior: 'push'}} />
+        <Route path="*" render={({location}) => (querySearch = location.search) && <span />} />
+      </MemoryRouter>,
+    );
+
+    await userEvent.click(screen.getByText(`one`));
+    await waitFor(() => {
+      expect(querySearch).toEqual('?q=one');
+    });
+
+    await userEvent.click(screen.getByText(`two`));
+    await waitFor(() => {
+      expect(querySearch).toEqual('?q=two');
+    });
+
+    await userEvent.click(screen.getByText(`three`));
+    await waitFor(() => {
+      expect(querySearch).toEqual('?q=three');
+    });
+
+    await act(() => {
+      goback();
+    });
+
+    await waitFor(() => {
+      expect(querySearch).toEqual('?q=two');
+    });
+
+    await act(() => {
+      goback();
+    });
+
+    await waitFor(() => {
+      expect(querySearch).toEqual('?q=one');
+    });
+  });
+
+  it('supports replace behavior', async () => {
+    let querySearch: string | undefined;
+
+    let goback: () => void;
+    let push: (...args: Parameters<ReturnType<typeof useHistory>['push']>) => void;
+
+    const Test = ({options}: {options: Parameters<typeof useQueryPersistedState>[0]}) => {
+      const [_, setQuery] = useQueryPersistedState(options);
+      const history = useHistory();
+      goback = () => history.goBack();
+      push = history.push.bind(history);
+      return (
+        <>
+          <div onClick={() => setQuery('one')}>one</div>
+          <div onClick={() => setQuery('two')}>two</div>
+          <div onClick={() => setQuery('three')}>three</div>
+        </>
+      );
+    };
+
+    render(
+      <MemoryRouter initialEntries={['/page?q=B']}>
+        <Test options={{queryKey: 'q', behavior: 'replace'}} />
+        <Route path="*" render={({location}) => (querySearch = location.search) && <span />} />
+      </MemoryRouter>,
+    );
+
+    push!('/page?');
+
+    await userEvent.click(screen.getByText(`one`));
+    await waitFor(() => {
+      expect(querySearch).toEqual('?q=one');
+    });
+
+    await userEvent.click(screen.getByText(`two`));
+    await waitFor(() => {
+      expect(querySearch).toEqual('?q=two');
+    });
+
+    await userEvent.click(screen.getByText(`three`));
+    await waitFor(() => {
+      expect(querySearch).toEqual('?q=three');
+    });
+
+    await act(() => {
+      goback();
+    });
+
+    expect(querySearch).toEqual('?q=B'); // end up back on initial route
+  });
 });

js_modules/dagster-ui/packages/ui-core/src/hooks/useQueryPersistedState.tsx

@@ -21,6 +21,7 @@ export type QueryPersistedStateConfig<T extends QueryPersistedDataType> = {
   defaults?: {[key: string]: any};
   decode?: (raw: {[key: string]: any}) => T;
   encode?: (raw: T) => {[key: string]: any};
+  behavior?: 'push' | 'replace';
 };
 
 const defaultEncode = memoize(<T,>(queryKey: string) => (raw: T) => ({[queryKey]: raw}));
@@ -63,7 +64,7 @@ const defaultDecode = memoize(
 export function useQueryPersistedState<T extends QueryPersistedDataType>(
   options: QueryPersistedStateConfig<T>,
 ): [T, React.Dispatch<React.SetStateAction<T>>] {
-  const {queryKey, defaults} = options;
+  const {queryKey, defaults, behavior = 'replace'} = options;
   let {encode, decode} = options;
 
   if (queryKey) {
@@ -113,9 +114,15 @@ export function useQueryPersistedState<T extends QueryPersistedDataType>(
       // the `replace` so that we surface any unwanted loops during development.
       if (process.env.NODE_ENV !== 'production' || !areQueriesEqual(currentQueryString, next)) {
         currentQueryString = next;
-        history.replace(
-          `${history.location.pathname}?${qs.stringify(next, {arrayFormat: 'brackets'})}`,
-        );
+        if (behavior === 'replace') {
+          history.replace(
+            `${history.location.pathname}?${qs.stringify(next, {arrayFormat: 'brackets'})}`,
+          );
+        } else {
+          history.push(
+            `${history.location.pathname}?${qs.stringify(next, {arrayFormat: 'brackets'})}`,
+          );
+        }
       }
     },
     [history, encode, options],