Fix/4407 set query data use inifinite query

.changeset/new-toys-occur.md

@@ -0,0 +1,5 @@
+---
+"@blitzjs/rpc": patch
+---
+
+fix(4407): setQueryData with useInfiniteQuery

apps/web/src/pages/page-with-inf-mutate.tsx

@@ -0,0 +1,90 @@
+import {useSuspenseInfiniteQuery} from "@blitzjs/rpc"
+import getInfiniteUsers from "src/queries/getInfiniteUsers"
+import {useActionState} from "react"
+
+function PageWithInfiniteQueryMutate(props) {
+  const [usersPages, extraInfo] = useSuspenseInfiniteQuery(
+    getInfiniteUsers,
+    (page = {take: 3, skip: 0}) => page,
+    {
+      getNextPageParam: (lastPage) => lastPage.nextPage,
+      initialPageParam: {take: 3, skip: 0},
+    },
+  )
+  const {isFetchingNextPage, fetchNextPage, hasNextPage, setQueryData} = extraInfo
+
+  const onOnContactSave = async (previousState, formData: FormData) => {
+    const name = formData.get("name") as string | null
+
+    await setQueryData(
+      (oldData) => {
+        if (!oldData) {
+          return {
+            pages: [],
+            pageParams: [],
+          }
+        }
+
+        return {
+          ...oldData,
+          pages: oldData.pages.map((page, index) => {
+            if (index === 0) {
+              return {
+                ...page,
+                users: [
+                  {
+                    id: Math.random(),
+                    name,
+                    role: "user",
+                    email: `${name}@yopmail.com`,
+                    createdAt: new Date(),
+                    updatedAt: new Date(),
+                    hashedPassword: "alsdklaskdoaskdokdo",
+                  },
+                  ...page.users,
+                ],
+              }
+            }
+            return page
+          }),
+        }
+      },
+      {refetch: false},
+    )
+  }
+
+  const [, formAction] = useActionState(onOnContactSave, {name: ""})
+
+  return (
+    <div>
+      <form action={formAction}>
+        <input type="text" name="name" placeholder="User name" />
+        <button type="submit">Add user</button>
+      </form>
+      {usersPages.map((usersPage) => (
+        <>
+          {usersPage?.users.map((u) => (
+            <div key={u.name}>
+              <p>name: {u.name}</p>
+              <p>role: {u.role}</p>
+              <p>email: {u.email}</p>
+              <hr />
+            </div>
+          ))}
+
+          {usersPage.hasMore && (
+            <button onClick={() => fetchNextPage()} disabled={!hasNextPage || !!isFetchingNextPage}>
+              {isFetchingNextPage
+                ? "Loading more..."
+                : hasNextPage
+                ? "Load More"
+                : "Nothing more to load"}
+            </button>
+          )}
+        </>
+      ))}
+    </div>
+  )
+}
+
+export default PageWithInfiniteQueryMutate

packages/blitz-rpc/src/query/react-query/react-query.ts

@@ -39,6 +39,7 @@ import {
   sanitizeQuery,
   sanitizeMutation,
   getInfiniteQueryKey,
+  QueryType,
 } from "../utils"
 import {useRouter} from "next/compat/router"
 
@@ -282,7 +283,7 @@ export function usePaginatedQuery<
 // -------------------------
 export interface RestInfiniteResult<TResult, TError>
   extends Omit<UseInfiniteQueryResult<TResult, TError>, "data">,
-    QueryCacheFunctions<TResult> {
+    QueryCacheFunctions<InfiniteData<TResult>> {
   pageParams: any
 }
 
@@ -355,7 +356,7 @@ export function useInfiniteQuery<
 
   const rest = {
     ...queryRest,
-    ...getQueryCacheFunctions<FirstParam<T>, TResult, T>(queryFn, getQueryParams),
+    ...getQueryCacheFunctions<FirstParam<T>, InfiniteData<TResult>, T>(queryFn, getQueryParams),
     pageParams: infiniteQueryData?.pageParams,
   }
 
@@ -367,7 +368,7 @@ export function useInfiniteQuery<
 // -------------------------
 export interface RestInfiniteResult<TResult, TError>
   extends Omit<UseInfiniteQueryResult<TResult, TError>, "data">,
-    QueryCacheFunctions<TResult> {
+    QueryCacheFunctions<InfiniteData<TResult>> {
   pageParams: any
 }
 
@@ -449,7 +450,11 @@ export function useSuspenseInfiniteQuery<
 
   const rest = {
     ...queryRest,
-    ...getQueryCacheFunctions<FirstParam<T>, TResult, T>(queryFn, getQueryParams),
+    ...getQueryCacheFunctions<FirstParam<T>, InfiniteData<TResult>, T>(
+      queryFn,
+      getQueryParams,
+      QueryType.INFINITE,
+    ),
     pageParams: infiniteQueryData?.pageParams,
   }
 

packages/blitz-rpc/src/query/utils.ts

@@ -65,9 +65,10 @@ export interface QueryCacheFunctions<T> {
 export const getQueryCacheFunctions = <TInput, TResult, T extends AsyncFunc>(
   resolver: T | Resolver<TInput, TResult> | RpcClient<TInput, TResult>,
   params: TInput,
+  queryType: QueryType = QueryType.STANDARD,
 ): QueryCacheFunctions<TResult> => ({
   setQueryData: (newData, opts = {refetch: true}) => {
-    return setQueryData(resolver, params, newData, opts)
+    return setQueryData(resolver, params, newData, opts, queryType)
   },
 })
 
@@ -170,16 +171,22 @@ export const invalidateQuery: InvalidateQuery = (resolver = undefined, ...params
   })
 }
 
+export enum QueryType {
+  STANDARD = "STANDARD",
+  INFINITE = "INFINITE",
+}
 export function setQueryData<TInput, TResult, T extends AsyncFunc>(
   resolver: T | Resolver<TInput, TResult> | RpcClient<TInput, TResult>,
   params: TInput,
   newData: TResult | ((oldData: TResult | undefined) => TResult | undefined),
   opts: MutateOptions = {refetch: true},
+  queryType: QueryType = QueryType.STANDARD,
 ): Promise<void | ReturnType<ReturnType<typeof getQueryClient>["invalidateQueries"]>> {
   if (typeof resolver === "undefined") {
     throw new Error("setQueryData is missing the first argument - it must be a resolver function")
   }
-  const queryKey = getQueryKey(resolver, params)
+  const getQueryKeyFn = queryType === QueryType.STANDARD ? getQueryKey : getInfiniteQueryKey
+  const queryKey = getQueryKeyFn(resolver, params)
 
   return new Promise((res) => {
     getQueryClient().setQueryData(queryKey, newData)