add upsertQueryDataThunk

Author: barnabasJ

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

@@ -13,7 +13,7 @@ import { QueryStatus } from './apiState'
 import type { InternalSerializeQueryArgs } from '../defaultSerializeQueryArgs'
 import type { Api, ApiContext } from '../apiTypes'
 import type { ApiEndpointQuery } from './module'
-import type { BaseQueryError } from '../baseQueryTypes'
+import type { BaseQueryError, QueryReturnValue } from '../baseQueryTypes'
 import type { QueryResultSelectorResult } from './buildSelectors'
 
 declare module './module' {
@@ -34,10 +34,11 @@ declare module './module' {
   }
 }
 
-export interface StartQueryActionCreatorOptions {
+export interface StartQueryActionCreatorOptions  {
   subscribe?: boolean
   forceRefetch?: boolean | number
   subscriptionOptions?: SubscriptionOptions
+  forceQueryFn?: () => QueryReturnValue
 }
 
 type StartQueryActionCreator<
@@ -179,6 +180,7 @@ export type MutationActionCreatorResult<
   unsubscribe(): void
 }
 
+
 export function buildInitiate({
   serializeQueryArgs,
   queryThunk,
@@ -259,7 +261,7 @@ Features like automatic cache collection, automatic refetching etc. will not be
     endpointDefinition: QueryDefinition<any, any, any, any>
   ) {
     const queryAction: StartQueryActionCreator<any> =
-      (arg, { subscribe = true, forceRefetch, subscriptionOptions } = {}) =>
+      (arg, { subscribe = true, forceRefetch, subscriptionOptions, forceQueryFn } = {}) =>
       (dispatch, getState) => {
         const queryCacheKey = serializeQueryArgs({
           queryArgs: arg,
@@ -270,6 +272,7 @@ Features like automatic cache collection, automatic refetching etc. will not be
           type: 'query',
           subscribe,
           forceRefetch,
+          forceQueryFn,
           subscriptionOptions,
           endpointName,
           originalArgs: arg,

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

@@ -144,6 +144,7 @@ function defaultTransformResponse(baseQueryReturnValue: unknown) {
 
 export type MaybeDrafted<T> = T | Draft<T>
 export type Recipe<T> = (data: MaybeDrafted<T>) => void | MaybeDrafted<T>
+export type UpsertRecipe<T> = (data: MaybeDrafted<T> | undefined) => void | MaybeDrafted<T>
 
 export type PatchQueryDataThunk<
   Definitions extends EndpointDefinitions,
@@ -163,6 +164,15 @@ export type UpdateQueryDataThunk<
   updateRecipe: Recipe<ResultTypeFrom<Definitions[EndpointName]>>
 ) => ThunkAction<PatchCollection, PartialState, any, AnyAction>
 
+export type UpsertQueryDataThunk<
+  Definitions extends EndpointDefinitions,
+  PartialState
+> = <EndpointName extends QueryKeys<Definitions>>(
+  endpointName: EndpointName,
+  args: QueryArgFrom<Definitions[EndpointName]>,
+  upsertRecipe: UpsertRecipe<ResultTypeFrom<Definitions[EndpointName]>>
+) => ThunkAction<PatchCollection, PartialState, any, AnyAction>
+
 /**
  * An object returned from dispatching a `api.util.updateQueryData` call.
  */
@@ -252,6 +262,53 @@ export function buildThunks<
 
       dispatch(api.util.patchQueryData(endpointName, args, ret.patches))
 
+      return ret
+    }
+
+  const upsertQueryData: UpsertQueryDataThunk<EndpointDefinitions, State> =
+    (endpointName, args, upsertRecipe) => (dispatch, getState) => {
+      const currentState = (
+        api.endpoints[endpointName] as ApiEndpointQuery<any, any>
+      ).select(args)(getState())
+      let ret: PatchCollection = {
+        patches: [],
+        inversePatches: [],
+        undo: () =>
+          dispatch(
+            api.util.patchQueryData(endpointName, args, ret.inversePatches)
+          ),
+      }
+      if (currentState.status === QueryStatus.uninitialized) {
+        return ret
+      }
+      if ('data' in currentState) {
+        if (isDraftable(currentState.data)) {
+          const [, patches, inversePatches] = produceWithPatches(
+            currentState.data,
+            upsertRecipe
+          )
+          ret.patches.push(...patches)
+          ret.inversePatches.push(...inversePatches)
+        } else {
+          const value = upsertRecipe(currentState.data)
+          ret.patches.push({ op: 'replace', path: [], value })
+          ret.inversePatches.push({
+            op: 'replace',
+            path: [],
+            value: currentState.data,
+          })
+        }
+        dispatch(api.util.patchQueryData(endpointName, args, ret.patches))
+      } else {
+        ret.inversePatches.push({
+          op: 'replace',
+          path: [],
+          value: undefined,
+        })
+        dispatch(api.endpoints[endpointName].initiate(args, {subscribe: false, forceRefetch: true, }))
+      }
+
+
       return ret
     }
 
@@ -291,7 +348,10 @@ export function buildThunks<
         forced:
           arg.type === 'query' ? isForcedQuery(arg, getState()) : undefined,
       }
-      if (endpointDefinition.query) {
+      if ('forceQueryFn' in arg && arg.forceQueryFn) {
+        result = arg.forceQueryFn()
+        
+      }else if (endpointDefinition.query) {
         result = await baseQuery(
           endpointDefinition.query(arg.originalArgs),
           baseQueryApi,
@@ -527,6 +587,7 @@ In the case of an unhandled error, no tags will be "provided" or "invalidated".`
     mutationThunk,
     prefetch,
     updateQueryData,
+    upsertQueryData,
     patchQueryData,
     buildMatchThunkActions,
   }