Unsubscribe existing query when arg changes on a lazy query

Author: msutkowski

package.json

@@ -110,7 +110,8 @@
   ],
   "dependencies": {
     "@babel/runtime": "^7.12.5",
-    "immer": ">=8.0.0"
+    "immer": ">=8.0.0",
+    "react-fast-compare": "^3.2.0"
   },
   "peerDependencies": {
     "@reduxjs/toolkit": "^1.5.0",

src/react-hooks/buildHooks.ts

@@ -24,6 +24,7 @@ import { Id, NoInfer, Override } from '../tsHelpers';
 import { ApiEndpointMutation, ApiEndpointQuery, CoreModule, PrefetchOptions } from '../core/module';
 import { ReactHooksModuleOptions } from './module';
 import { useShallowStableValue } from './useShallowStableValue';
+import isDeepEqual from 'react-fast-compare';
 
 export interface QueryHooks<Definition extends QueryDefinition<any, any, any, any, any>> {
   useQuery: UseQuery<Definition>;
@@ -253,6 +254,7 @@ export function buildHooks<Definitions extends EndpointDefinitions>({
       const lastPromise = promiseRef.current;
 
       const optionsRef = useRef<SubscriptionOptions>();
+      const argsRef = useRef<any>();
 
       useEffect(() => {
         const options = {
@@ -274,23 +276,36 @@ export function buildHooks<Definitions extends EndpointDefinitions>({
           promiseRef.current?.unsubscribe();
           promiseRef.current = undefined;
           optionsRef.current = undefined;
+          argsRef.current = undefined;
         };
       }, []);
 
       const trigger = useCallback(
         function (args: any) {
+          if (!argsRef.current) {
+            argsRef.current = args;
+          } else if (argsRef.current) {
+            // args have changed, we need to unsubscribe before creating the new subscription ref.
+            if (!isDeepEqual(argsRef.current, args)) {
+              // eslint-disable-next-line @typescript-eslint/no-unused-expressions
+              lastPromise?.unsubscribe();
+              argsRef.current = args;
+            }
+          }
+
+          // Set the subscription options on the initial query
+          if (!optionsRef.current) {
+            optionsRef.current = { pollingInterval, refetchOnReconnect, refetchOnFocus };
+          }
+
           promiseRef.current = dispatch(
             initiate(args, {
               subscriptionOptions: { pollingInterval, refetchOnReconnect, refetchOnFocus },
               forceRefetch: true,
             })
           );
-          // Set the subscription options on the initial query
-          if (!optionsRef.current) {
-            optionsRef.current = { pollingInterval, refetchOnReconnect, refetchOnFocus };
-          }
         },
-        [dispatch, initiate, pollingInterval, refetchOnFocus, refetchOnReconnect]
+        [dispatch, initiate, lastPromise, pollingInterval, refetchOnFocus, refetchOnReconnect]
       );
 
       return [trigger, promiseRef];

test/buildHooks.test.tsx

@@ -466,6 +466,65 @@ describe('hooks tests', () => {
         storeRef.store.getState().actions.filter(api.internalActions.updateSubscriptionOptions.match)
       ).toHaveLength(1);
     });
+
+    test('useLazyQuery accepts updated args and unsubscribes the original query', async () => {
+      function User() {
+        const [fetchUser, { data: hookData, isFetching, isUninitialized }] = api.endpoints.getUser.useLazyQuery();
+
+        data = hookData;
+
+        return (
+          <div>
+            <div data-testid="isUninitialized">{String(isUninitialized)}</div>
+            <div data-testid="isFetching">{String(isFetching)}</div>
+
+            <button data-testid="fetchUser1" onClick={() => fetchUser(1)}>
+              fetchUser1
+            </button>
+            <button data-testid="fetchUser2" onClick={() => fetchUser(2)}>
+              fetchUser2
+            </button>
+          </div>
+        );
+      }
+
+      render(<User />, { wrapper: storeRef.wrapper });
+
+      await waitFor(() => expect(screen.getByTestId('isUninitialized').textContent).toBe('true'));
+      await waitFor(() => expect(data).toBeUndefined());
+
+      fireEvent.click(screen.getByTestId('fetchUser1'));
+
+      await waitFor(() => expect(screen.getByTestId('isFetching').textContent).toBe('true'));
+      await waitFor(() => expect(screen.getByTestId('isFetching').textContent).toBe('false'));
+
+      // Being that there is only the initial query, no unsubscribe should be dispatched
+      expect(storeRef.store.getState().actions.filter(api.internalActions.unsubscribeQueryResult.match)).toHaveLength(
+        0
+      );
+
+      fireEvent.click(screen.getByTestId('fetchUser2'));
+
+      await waitFor(() => expect(screen.getByTestId('isFetching').textContent).toBe('true'));
+      await waitFor(() => expect(screen.getByTestId('isFetching').textContent).toBe('false'));
+
+      // `args` changed from 1 -> 2, should unsubscribe the original for 1.
+      expect(storeRef.store.getState().actions.filter(api.internalActions.unsubscribeQueryResult.match)).toHaveLength(
+        1
+      );
+
+      fireEvent.click(screen.getByTestId('fetchUser1'));
+
+      expect(storeRef.store.getState().actions.filter(api.internalActions.unsubscribeQueryResult.match)).toHaveLength(
+        2
+      );
+
+      // no `arg` change, no unsubscribe happens, just another pending request/fulfilled
+      fireEvent.click(screen.getByTestId('fetchUser1'));
+      expect(storeRef.store.getState().actions.filter(api.internalActions.unsubscribeQueryResult.match)).toHaveLength(
+        2
+      );
+    });
   });
 
   describe('useMutation', () => {

test/createApi.test.ts

@@ -439,6 +439,10 @@ describe('additional transformResponse behaviors', () => {
       echo: build.mutation({
         query: () => ({ method: 'PUT', url: '/echo' }),
       }),
+      mutation: build.mutation({
+        query: () => ({ url: '/echo', method: 'POST', body: { nested: { banana: 'bread' } } }),
+        transformResponse: (response: { body: { nested: EchoResponseData } }) => response.body.nested,
+      }),
       query: build.query<SuccessResponse & EchoResponseData, void>({
         query: () => '/success',
         transformResponse: async (response: SuccessResponse) => {
@@ -455,11 +459,17 @@ describe('additional transformResponse behaviors', () => {
 
   const storeRef = setupApiStore(api);
 
-  test('transformResponse handles an async transformation and returns the merged data', async () => {
+  test('transformResponse handles an async transformation and returns the merged data (query)', async () => {
     const result = await storeRef.store.dispatch(api.endpoints.query.initiate());
 
     expect(result.data).toEqual({ value: 'success', banana: 'bread' });
   });
+
+  test('transformResponse transforms a response from a mutation', async () => {
+    const result = await storeRef.store.dispatch(api.endpoints.mutation.initiate({}));
+
+    expect(result.data).toEqual({ banana: 'bread' });
+  });
 });
 
 describe('query endpoint lifecycles - onStart, onSuccess, onError', () => {

yarn.lock

@@ -7429,6 +7429,11 @@ [email protected]:
     object-assign "^4.1.1"
     scheduler "^0.20.0"
 
+react-fast-compare@^3.2.0:
+  version "3.2.0"
+  resolved "https://registry.yarnpkg.com/react-fast-compare/-/react-fast-compare-3.2.0.tgz#641a9da81b6a6320f270e89724fb45a0b39e43bb"
+  integrity sha512-rtGImPZ0YyLrscKI9xTpV8psd6I8VAtjKCzQDlzyDvqJA8XOW78TXYQwNRNd8g8JZnDu8q9Fu/1v4HPAVwVdHA==
+
 "react-is@^16.12.0 || ^17.0.0", react-is@^17.0.1:
   version "17.0.1"
   resolved "https://registry.yarnpkg.com/react-is/-/react-is-17.0.1.tgz#5b3531bd76a645a4c9fb6e693ed36419e3301339"