allow for "shared component results" using the useMutation hook and fixedCacheKey (#1477)

Author: phryneas

packages/toolkit/src/query/core/apiState.ts

@@ -12,7 +12,15 @@ import type { Id, WithRequiredProp } from '../tsHelpers'
 
 export type QueryCacheKey = string & { _type: 'queryCacheKey' }
 export type QuerySubstateIdentifier = { queryCacheKey: QueryCacheKey }
-export type MutationSubstateIdentifier = { requestId: string }
+export type MutationSubstateIdentifier =
+  | {
+      requestId: string
+      fixedCacheKey?: string
+    }
+  | {
+      requestId?: string
+      fixedCacheKey: string
+    }
 
 export type RefetchConfigOptions = {
   refetchOnMountOrArgChange: boolean | number
@@ -175,6 +183,7 @@ export type QuerySubState<D extends BaseEndpointDefinition<any, any, any>> = Id<
 >
 
 type BaseMutationSubState<D extends BaseEndpointDefinition<any, any, any>> = {
+  requestId: string
   data?: ResultTypeFrom<D>
   error?:
     | SerializedError
@@ -200,6 +209,7 @@ export type MutationSubState<D extends BaseEndpointDefinition<any, any, any>> =
       status: QueryStatus.rejected
     } & WithRequiredProp<BaseMutationSubState<D>, 'error'>)
   | {
+      requestId?: undefined
       status: QueryStatus.uninitialized
       data?: undefined
       error?: undefined

packages/toolkit/src/query/core/buildInitiate.ts

@@ -79,6 +79,7 @@ type StartMutationActionCreator<
      * (defaults to `true`)
      */
     track?: boolean
+    fixedCacheKey?: string
   }
 ) => ThunkAction<MutationActionCreatorResult<D>, any, any, AnyAction>
 
@@ -113,6 +114,7 @@ export type MutationActionCreatorResult<
      * Whether the mutation is being tracked in the store.
      */
     track?: boolean
+    fixedCacheKey?: string
   }
   /**
    * A unique string generated for the request sequence
@@ -287,12 +289,13 @@ Features like automatic cache collection, automatic refetching etc. will not be
     endpointName: string,
     definition: MutationDefinition<any, any, any, any>
   ): StartMutationActionCreator<any> {
-    return (arg, { track = true } = {}) =>
+    return (arg, { track = true, fixedCacheKey } = {}) =>
       (dispatch, getState) => {
         const thunk = mutationThunk({
           endpointName,
           originalArgs: arg,
           track,
+          fixedCacheKey,
         })
         const thunkResult = dispatch(thunk)
         middlewareWarning(getState)
@@ -303,7 +306,7 @@ Features like automatic cache collection, automatic refetching etc. will not be
           .catch((error) => ({ error }))
 
         const reset = () => {
-          if (track) dispatch(removeMutationResult({ requestId }))
+          dispatch(removeMutationResult({ requestId, fixedCacheKey }))
         }
 
         return Object.assign(returnValuePromise, {

packages/toolkit/src/query/core/buildMiddleware/cacheLifecycle.ts

@@ -8,6 +8,7 @@ import type {
   MutationResultSelectorResult,
   QueryResultSelectorResult,
 } from '../buildSelectors'
+import { getMutationCacheKey } from '../buildSlice'
 import type { PatchCollection, Recipe } from '../buildThunks'
 import type {
   PromiseWithKnownReason,
@@ -258,7 +259,7 @@ export const build: SubMiddlewareBuilder = ({
       if (api.internalActions.removeQueryResult.match(action))
         return action.payload.queryCacheKey
       if (api.internalActions.removeMutationResult.match(action))
-        return action.payload.requestId
+        return getMutationCacheKey(action.payload)
       return ''
     }
 

packages/toolkit/src/query/core/buildSelectors.ts

@@ -15,6 +15,7 @@ import type {
   ReducerPathFrom,
 } from '../endpointDefinitions'
 import type { InternalSerializeQueryArgs } from '../defaultSerializeQueryArgs'
+import { getMutationCacheKey } from './buildSlice'
 
 export type SkipToken = typeof skipToken
 /**
@@ -88,7 +89,10 @@ type MutationResultSelectorFactory<
   Definition extends MutationDefinition<any, any, any, any>,
   RootState
 > = (
-  requestId: string | SkipToken
+  requestId:
+    | string
+    | { requestId: string | undefined; fixedCacheKey: string | undefined }
+    | SkipToken
 ) => (state: RootState) => MutationResultSelectorResult<Definition>
 
 export type MutationResultSelectorResult<
@@ -172,7 +176,13 @@ export function buildSelectors<
     any,
     RootState
   > {
-    return (mutationId) => {
+    return (id) => {
+      let mutationId: string | typeof skipToken
+      if (typeof id === 'object') {
+        mutationId = getMutationCacheKey(id) ?? skipToken
+      } else {
+        mutationId = id
+      }
       const selectMutationSubstate = createSelector(
         selectInternalState,
         (internalState) =>

packages/toolkit/src/query/core/buildSlice.ts

@@ -28,13 +28,7 @@ import type {
   ConfigState,
 } from './apiState'
 import { QueryStatus } from './apiState'
-import type {
-  MutationThunk,
-  MutationThunkArg,
-  QueryThunk,
-  QueryThunkArg,
-  ThunkResult,
-} from './buildThunks'
+import type { MutationThunk, QueryThunk } from './buildThunks'
 import { calculateProvidedByThunk } from './buildThunks'
 import type {
   AssertTagTypes,
@@ -61,12 +55,33 @@ function updateQuerySubstateIfExists(
   }
 }
 
+export function getMutationCacheKey(
+  id:
+    | MutationSubstateIdentifier
+    | { requestId: string; arg: { fixedCacheKey?: string | undefined } }
+): string
+export function getMutationCacheKey(id: {
+  fixedCacheKey?: string
+  requestId?: string
+}): string | undefined
+
+export function getMutationCacheKey(
+  id:
+    | { fixedCacheKey?: string; requestId?: string }
+    | MutationSubstateIdentifier
+    | { requestId: string; arg: { fixedCacheKey?: string | undefined } }
+): string | undefined {
+  return ('arg' in id ? id.arg.fixedCacheKey : id.fixedCacheKey) ?? id.requestId
+}
+
 function updateMutationSubstateIfExists(
   state: MutationState<any>,
-  { requestId }: MutationSubstateIdentifier,
+  id:
+    | MutationSubstateIdentifier
+    | { requestId: string; arg: { fixedCacheKey?: string | undefined } },
   update: (substate: MutationSubState<any>) => void
 ) {
-  const substate = state[requestId]
+  const substate = state[getMutationCacheKey(id)]
   if (substate) {
     update(substate)
   }
@@ -174,50 +189,50 @@ export function buildSlice({
     reducers: {
       removeMutationResult(
         draft,
-        action: PayloadAction<MutationSubstateIdentifier>
+        { payload }: PayloadAction<MutationSubstateIdentifier>
       ) {
-        if (action.payload.requestId in draft) {
-          delete draft[action.payload.requestId]
+        const cacheKey = getMutationCacheKey(payload)
+        if (cacheKey in draft) {
+          delete draft[cacheKey]
         }
       },
     },
     extraReducers(builder) {
       builder
         .addCase(
           mutationThunk.pending,
-          (draft, { meta: { arg, requestId, startedTimeStamp } }) => {
+          (draft, { meta, meta: { requestId, arg, startedTimeStamp } }) => {
             if (!arg.track) return
 
-            draft[requestId] = {
+            draft[getMutationCacheKey(meta)] = {
+              requestId,
               status: QueryStatus.pending,
               endpointName: arg.endpointName,
               startedTimeStamp,
             }
           }
         )
-        .addCase(
-          mutationThunk.fulfilled,
-          (draft, { payload, meta, meta: { requestId } }) => {
-            if (!meta.arg.track) return
+        .addCase(mutationThunk.fulfilled, (draft, { payload, meta }) => {
+          if (!meta.arg.track) return
 
-            updateMutationSubstateIfExists(draft, { requestId }, (substate) => {
-              substate.status = QueryStatus.fulfilled
-              substate.data = payload
-              substate.fulfilledTimeStamp = meta.fulfilledTimeStamp
-            })
-          }
-        )
-        .addCase(
-          mutationThunk.rejected,
-          (draft, { payload, error, meta: { requestId, arg } }) => {
-            if (!arg.track) return
+          updateMutationSubstateIfExists(draft, meta, (substate) => {
+            if (substate.requestId !== meta.requestId) return
 
-            updateMutationSubstateIfExists(draft, { requestId }, (substate) => {
-              substate.status = QueryStatus.rejected
-              substate.error = (payload ?? error) as any
-            })
-          }
-        )
+            substate.status = QueryStatus.fulfilled
+            substate.data = payload
+            substate.fulfilledTimeStamp = meta.fulfilledTimeStamp
+          })
+        })
+        .addCase(mutationThunk.rejected, (draft, { payload, error, meta }) => {
+          if (!meta.arg.track) return
+
+          updateMutationSubstateIfExists(draft, meta, (substate) => {
+            if (substate.requestId !== meta.requestId) return
+
+            substate.status = QueryStatus.rejected
+            substate.error = (payload ?? error) as any
+          })
+        })
     },
   })
 

packages/toolkit/src/query/core/buildThunks.ts

@@ -111,6 +111,7 @@ export interface MutationThunkArg {
   originalArgs: unknown
   endpointName: string
   track?: boolean
+  fixedCacheKey?: string
 }
 
 export type ThunkResult = unknown

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

@@ -401,6 +401,7 @@ export type UseMutationStateOptions<
   R extends Record<string, any>
 > = {
   selectFromResult?: MutationStateSelector<R, D>
+  fixedCacheKey?: string
 }
 
 export type UseMutationStateResult<
@@ -781,35 +782,60 @@ export function buildHooks<Definitions extends EndpointDefinitions>({
   }
 
   function buildMutationHook(name: string): UseMutation<any> {
-    return ({ selectFromResult = defaultMutationStateSelector } = {}) => {
+    return ({
+      selectFromResult = defaultMutationStateSelector,
+      fixedCacheKey,
+    } = {}) => {
       const { select, initiate } = api.endpoints[name] as ApiEndpointMutation<
         MutationDefinition<any, any, any, any, any>,
         Definitions
       >
       const dispatch = useDispatch<ThunkDispatch<any, any, AnyAction>>()
       const [promise, setPromise] = useState<MutationActionCreatorResult<any>>()
 
-      useEffect(() => () => promise?.reset(), [promise])
+      useEffect(
+        () => () => {
+          if (!promise?.arg.fixedCacheKey) {
+            promise?.reset()
+          }
+        },
+        [promise]
+      )
 
       const triggerMutation = useCallback(
         function (arg) {
-          const promise = dispatch(initiate(arg))
+          const promise = dispatch(initiate(arg, { fixedCacheKey }))
           setPromise(promise)
           return promise
         },
-        [dispatch, initiate]
+        [dispatch, initiate, fixedCacheKey]
       )
 
       const { requestId } = promise || {}
       const mutationSelector = useMemo(
         () =>
-          createSelector([select(requestId || skipToken)], selectFromResult),
-        [select, requestId, selectFromResult]
+          createSelector(
+            [select({ fixedCacheKey, requestId: promise?.requestId })],
+            selectFromResult
+          ),
+        [select, promise, selectFromResult, fixedCacheKey]
       )
 
       const currentState = useSelector(mutationSelector, shallowEqual)
-      const originalArgs = promise?.arg.originalArgs
-      const reset = useCallback(() => setPromise(undefined), [])
+      const originalArgs =
+        fixedCacheKey == null ? promise?.arg.originalArgs : undefined
+      const reset = useCallback(() => {
+        if (promise) {
+          setPromise(undefined)
+        } else if (fixedCacheKey) {
+          dispatch(
+            api.internalActions.removeMutationResult({
+              requestId,
+              fixedCacheKey,
+            })
+          )
+        }
+      }, [dispatch, fixedCacheKey, promise, requestId])
       const finalState = useMemo(
         () => ({ ...currentState, originalArgs, reset }),
         [currentState, originalArgs, reset]