Add possibility to manually unsubscribe from useLazyQuery #2055

packages/toolkit/src/query/react/buildHooks.ts

@@ -1,6 +1,6 @@
 import type { AnyAction, ThunkAction, ThunkDispatch } from '@reduxjs/toolkit'
 import { createSelector } from '@reduxjs/toolkit'
-import type { Selector } from '@reduxjs/toolkit'
+import type { Selector, Unsubscribe } from '@reduxjs/toolkit'
 import type { DependencyList } from 'react'
 import {
   useCallback,
@@ -164,10 +164,11 @@ export type UseQuerySubscription<
   options?: UseQuerySubscriptionOptions
 ) => Pick<QueryActionCreatorResult<D>, 'refetch'>
 
-export type UseLazyQueryLastPromiseInfo<
+export type UseLazyQueryExtras<
   D extends QueryDefinition<any, any, any, any>
 > = {
   lastArg: QueryArgFrom<D>
+  unsubscribe: Unsubscribe | undefined
 }
 
 /**
@@ -194,7 +195,7 @@ export type UseLazyQuery<D extends QueryDefinition<any, any, any, any>> = <
 ) => [
   LazyQueryTrigger<D>,
   UseQueryStateResult<D, R>,
-  UseLazyQueryLastPromiseInfo<D>
+  UseLazyQueryExtras<D>
 ]
 
 export type LazyQueryTrigger<D extends QueryDefinition<any, any, any, any>> = {
@@ -239,7 +240,7 @@ export type UseLazyQuerySubscription<
   D extends QueryDefinition<any, any, any, any>
 > = (
   options?: SubscriptionOptions
-) => readonly [LazyQueryTrigger<D>, QueryArgFrom<D> | UninitializedValue]
+) => readonly [LazyQueryTrigger<D>, QueryArgFrom<D> | UninitializedValue, Unsubscribe | undefined]
 
 export type QueryStateSelector<
   R extends Record<string, any>,
@@ -772,7 +773,7 @@ export function buildHooks<Definitions extends EndpointDefinitions>({
         }
       }, [arg, trigger])
 
-      return useMemo(() => [trigger, arg] as const, [trigger, arg])
+      return useMemo(() => [trigger, arg, promiseRef.current?.unsubscribe] as const, [trigger, arg])
     }
 
     const useQueryState: UseQueryState<any> = (
@@ -835,16 +836,17 @@ export function buildHooks<Definitions extends EndpointDefinitions>({
       useQuerySubscription,
       useLazyQuerySubscription,
       useLazyQuery(options) {
-        const [trigger, arg] = useLazyQuerySubscription(options)
+        const [trigger, arg, unsubscribe] = useLazyQuerySubscription(options)
         const queryStateResults = useQueryState(arg, {
           ...options,
           skip: arg === UNINITIALIZED_VALUE,
         })
 
-        const info = useMemo(() => ({ lastArg: arg }), [arg])
+        const extras = useMemo(() => ({ lastArg: arg, unsubscribe }), [arg, unsubscribe])
+
         return useMemo(
-          () => [trigger, queryStateResults, info],
-          [trigger, queryStateResults, info]
+          () => [trigger, queryStateResults, extras],
+          [trigger, queryStateResults, extras]
         )
       },
       useQuery(arg, options) {

packages/toolkit/src/query/tests/buildHooks.test.tsx

@@ -1027,6 +1027,73 @@ describe('hooks tests', () => {
 
       expect(screen.getByTestId('error').textContent).toBe('')
     })
+
+    test('useLazyQuery can be manually unsubscribed', async () => {
+      function User() {
+        const [getUser, { data, error }, { unsubscribe }] =
+          api.endpoints.getUser.useLazyQuery()
+
+        const [unwrappedResult, setUnwrappedResult] = React.useState<
+          undefined | { name: string }
+        >()
+
+        const handleFetchClick = async () => {
+          const res = getUser(1)
+
+          const result = await res.unwrap()
+          setUnwrappedResult(result)
+        }
+
+        return (
+          <div>
+            <button onClick={handleFetchClick}>Fetch User</button>
+            <button onClick={unsubscribe}>Unsubscribe</button>
+            <div data-testid="result">{JSON.stringify(data)}</div>
+            <div data-testid="error">{JSON.stringify(error)}</div>
+            <div data-testid="unwrappedResult">
+              {JSON.stringify(unwrappedResult)}
+            </div>
+          </div>
+        )
+      }
+
+      render(<User />, { wrapper: storeRef.wrapper })
+
+      const fetchButton = screen.getByRole('button', { name: 'Fetch User' })
+      fireEvent.click(fetchButton)
+
+      await waitFor(() => {
+        const dataResult = screen.getByTestId('result')?.textContent
+        const unwrappedDataResult =
+          screen.getByTestId('unwrappedResult')?.textContent
+
+        if (dataResult && unwrappedDataResult) {
+          expect(JSON.parse(dataResult)).toMatchObject({
+            name: 'Timmy',
+          })
+          expect(JSON.parse(unwrappedDataResult)).toMatchObject(
+            JSON.parse(dataResult)
+          )
+        }
+      })
+
+      expect(screen.getByTestId('error').textContent).toBe('')
+
+      const initialSubscriptionsCount = Object.keys(
+        storeRef.store.getState().api.subscriptions['getUser(1)']!
+      ).length
+
+      expect(initialSubscriptionsCount).toEqual(1)
+
+      const unsubscribeButton = screen.getByRole('button', { name: 'Unsubscribe' })
+      fireEvent.click(unsubscribeButton)
+
+      const subscriptionsCountAfterUnsubscribing = Object.keys(
+        storeRef.store.getState().api.subscriptions['getUser(1)']!
+      ).length
+
+      expect(subscriptionsCountAfterUnsubscribing).toEqual(0)
+    })
   })
 
   describe('useMutation', () => {