Make prefetching work with usePaginatedQuery and useInfiniteQuery (#3014)

(patch)

.github/workflows/compressed.yml

@@ -15,7 +15,8 @@ jobs:
 
     steps:
       - uses: actions/checkout@v2
-      - uses: actions/setup-node@v2
+      - name: Use Node.js
+        uses: actions/setup-node@v2
         with:
         node-version: "14"
       - name: Count size

nextjs/packages/next/build/babel/plugins/rewrite-imports.ts

@@ -56,6 +56,7 @@ const specialImports: Record<string, string> = {
   useMutation: 'next/data-client',
   queryClient: 'next/data-client',
   getQueryKey: 'next/data-client',
+  getInfiniteQueryKey: 'next/data-client',
   invalidateQuery: 'next/data-client',
   setQueryData: 'next/data-client',
   useQueryErrorResetBoundary: 'next/data-client',

nextjs/packages/next/data-client/index.ts

@@ -11,6 +11,7 @@ export {
 export {
   queryClient,
   getQueryKey,
+  getInfiniteQueryKey,
   invalidateQuery,
   setQueryData,
 } from './react-query-utils'

nextjs/packages/next/data-client/react-query-utils.ts

@@ -125,6 +125,23 @@ export function getQueryKey<TInput, TResult, T extends AsyncFunc>(
   return getQueryKeyFromUrlAndParams(sanitizeQuery(resolver)._routePath, params)
 }
 
+export function getInfiniteQueryKey<TInput, TResult, T extends AsyncFunc>(
+  resolver: T | Resolver<TInput, TResult> | RpcClient<TInput, TResult>,
+  params?: TInput
+) {
+  if (typeof resolver === 'undefined') {
+    throw new Error(
+      'getInfiniteQueryKey is missing the first argument - it must be a resolver function'
+    )
+  }
+
+  const queryKey = getQueryKeyFromUrlAndParams(
+    sanitizeQuery(resolver)._routePath,
+    params
+  )
+  return [...queryKey, 'infinite']
+}
+
 export function invalidateQuery<TInput, TResult, T extends AsyncFunc>(
   resolver: T | Resolver<TInput, TResult> | RpcClient<TInput, TResult>,
   params?: TInput

nextjs/packages/next/data-client/react-query.tsx

@@ -20,6 +20,7 @@ import {
   QueryCacheFunctions,
   sanitizeQuery,
   sanitizeMutation,
+  getInfiniteQueryKey,
 } from './react-query-utils'
 import { useRouter } from '../client/router'
 
@@ -166,16 +167,19 @@ export function usePaginatedQuery<
     )
   }
 
-  const suspense =
-    options?.enabled === false || options?.enabled === null
+  const suspenseEnabled = Boolean(process.env.__BLITZ_SUSPENSE_ENABLED)
+  let enabled =
+    isServer && suspenseEnabled
       ? false
-      : options?.suspense
+      : options?.enabled ?? options?.enabled !== null
+  const suspense = enabled === false ? false : options?.suspense
+
   const session = useSession({ suspense })
   if (session.isLoading) {
-    options.enabled = false
+    enabled = false
   }
 
-  const routerIsReady = useRouter().isReady
+  const routerIsReady = useRouter().isReady || (isServer && suspenseEnabled)
   const enhancedResolverRpcClient = sanitizeQuery(queryFn)
   const queryKey = getQueryKey(queryFn, params)
 
@@ -186,8 +190,20 @@ export function usePaginatedQuery<
       : (emptyQueryFn as any),
     ...options,
     keepPreviousData: true,
+    enabled,
   })
 
+  if (
+    queryRest.isIdle &&
+    isServer &&
+    suspenseEnabled !== false &&
+    !data &&
+    (!options || !('suspense' in options) || options.suspense) &&
+    (!options || !('enabled' in options) || options.enabled)
+  ) {
+    throw new Promise(() => {})
+  }
+
   const rest = {
     ...queryRest,
     ...getQueryCacheFunctions<FirstParam<T>, TResult, T>(queryFn, params),
@@ -250,24 +266,26 @@ export function useInfiniteQuery<
     )
   }
 
-  const suspense =
-    options?.enabled === false || options?.enabled === null
+  const suspenseEnabled = Boolean(process.env.__BLITZ_SUSPENSE_ENABLED)
+  let enabled =
+    isServer && suspenseEnabled
       ? false
-      : options?.suspense
+      : options?.enabled ?? options?.enabled !== null
+  const suspense = enabled === false ? false : options?.suspense
   const session = useSession({ suspense })
   if (session.isLoading) {
-    options.enabled = false
+    enabled = false
   }
 
-  const routerIsReady = useRouter().isReady
+  const routerIsReady = useRouter().isReady || (isServer && suspenseEnabled)
   const enhancedResolverRpcClient = sanitizeQuery(queryFn)
-  const queryKey = getQueryKey(queryFn, getQueryParams)
+  const queryKey = getInfiniteQueryKey(queryFn, getQueryParams)
 
   const { data, ...queryRest } = useInfiniteReactQuery({
     // we need an extra cache key for infinite loading so that the cache for
     // for this query is stored separately since the hook result is an array of results.
     // Without this cache for usePaginatedQuery and this will conflict and break.
-    queryKey: routerIsReady ? [...queryKey, 'infinite'] : ['_routerNotReady_'],
+    queryKey: routerIsReady ? queryKey : ['_routerNotReady_'],
     queryFn: routerIsReady
       ? ({ pageParam }) =>
           enhancedResolverRpcClient(getQueryParams(pageParam), {
@@ -275,8 +293,20 @@ export function useInfiniteQuery<
           })
       : (emptyQueryFn as any),
     ...options,
+    enabled,
   })
 
+  if (
+    queryRest.isIdle &&
+    isServer &&
+    suspenseEnabled !== false &&
+    !data &&
+    (!options || !('suspense' in options) || options.suspense) &&
+    (!options || !('enabled' in options) || options.enabled)
+  ) {
+    throw new Promise(() => {})
+  }
+
   const rest = {
     ...queryRest,
     ...getQueryCacheFunctions<FirstParam<T>, TResult, T>(

test/integration/queries/app/queries/getIncrementedWithPagination.ts

@@ -0,0 +1,33 @@
+import {paginate, resolver} from "blitz"
+
+const dataset = Array.from(Array(100).keys())
+
+type Args = {
+  skip: number
+  take: number
+  where?: {value: {gte: number}}
+}
+
+let counter = 0
+export default resolver.pipe(async ({skip = 0, take = 100, where}: Args) => {
+  counter++
+  const {items, hasMore, nextPage, count} = await paginate({
+    skip,
+    take,
+    count: async () => dataset.length,
+    query: async (paginateArgs) =>
+      dataset
+        .filter((i) => {
+          if (!where) return true
+          return i >= where.value.gte
+        })
+        .slice(paginateArgs.skip, paginateArgs.skip + paginateArgs.take),
+  })
+  return {
+    counter,
+    items,
+    hasMore,
+    nextPage,
+    count,
+  }
+})

test/integration/queries/pages/dehydrated-state-use-infinite-query.tsx

@@ -0,0 +1,47 @@
+import getPaginated from "app/queries/getPaginated"
+import {
+  dehydrate,
+  getInfiniteQueryKey,
+  GetServerSideProps,
+  invokeWithMiddleware,
+  QueryClient,
+  useInfiniteQuery,
+} from "blitz"
+import {useState} from "react"
+
+export const getServerSideProps: GetServerSideProps = async (ctx) => {
+  const queryClient = new QueryClient()
+  const queryKey = getInfiniteQueryKey(getPaginated, {where: {value: {gte: 10}}, take: 5, skip: 0})
+
+  await queryClient.prefetchInfiniteQuery(queryKey, () =>
+    invokeWithMiddleware(getPaginated, {where: {value: {gte: 10}}, take: 5, skip: 0}, ctx),
+  )
+
+  return {
+    props: {
+      dehydratedState: dehydrate(queryClient),
+    },
+  }
+}
+
+function Content() {
+  const [value] = useState(10)
+
+  const [groups] = useInfiniteQuery(
+    getPaginated,
+    (page = {take: 5, skip: 0}) => ({
+      where: {value: {gte: value}},
+      ...page,
+    }),
+    {
+      getNextPageParam: (lastGroup) => lastGroup.nextPage,
+    },
+  )
+  return <div id="content">{JSON.stringify(groups)}</div>
+}
+
+function InfiniteQueryDehydratedState() {
+  return <Content />
+}
+
+export default InfiniteQueryDehydratedState

test/integration/queries/pages/dehydrated-state-use-paginated-query.tsx

@@ -0,0 +1,32 @@
+import getMap from "app/queries/getMap"
+import {
+  dehydrate,
+  getQueryKey,
+  GetServerSideProps,
+  invokeWithMiddleware,
+  QueryClient,
+  usePaginatedQuery,
+} from "blitz"
+
+export const getServerSideProps: GetServerSideProps = async (ctx) => {
+  const queryClient = new QueryClient()
+  const queryKey = getQueryKey(getMap, undefined)
+  await queryClient.prefetchQuery(queryKey, () => invokeWithMiddleware(getMap, undefined, ctx))
+
+  return {
+    props: {
+      dehydratedState: dehydrate(queryClient),
+    },
+  }
+}
+
+function Content() {
+  const [map] = usePaginatedQuery(getMap, undefined)
+  return <p id="content">map is Map: {"" + (map instanceof Map)}</p>
+}
+
+function DehydratedStateWithPagination() {
+  return <Content />
+}
+
+export default DehydratedStateWithPagination

test/integration/queries/pages/dehydrated-state-use-query.tsx

@@ -7,7 +7,6 @@ import {
   QueryClient,
   useQuery,
 } from "blitz"
-import {Suspense} from "react"
 
 export const getServerSideProps: GetServerSideProps = async (ctx) => {
   const queryClient = new QueryClient()
@@ -27,11 +26,7 @@ function Content() {
 }
 
 function DehydratedState() {
-  return (
-    <Suspense fallback="Loading ...">
-      <Content />
-    </Suspense>
-  )
+  return <Content />
 }
 
 export default DehydratedState

test/integration/queries/pages/invalidate-use-infinite-query.tsx

@@ -0,0 +1,34 @@
+import getIncremented from "app/queries/getIncrementedWithPagination"
+import {invalidateQuery, useInfiniteQuery} from "blitz"
+import {Suspense} from "react"
+
+function Content() {
+  const [groups] = useInfiniteQuery(
+    getIncremented,
+    (page = {take: 5, skip: 0}) => ({
+      where: {value: {gte: 10}},
+      ...page,
+    }),
+    {
+      getNextPageParam: (lastGroup) => lastGroup.nextPage,
+    },
+  )
+  return (
+    <>
+      <button onClick={() => invalidateQuery(getIncremented)}>click me</button>
+      <div id="content">{JSON.stringify(groups)}</div>
+    </>
+  )
+}
+
+function InvalidateInfiniteQuery() {
+  return (
+    <div id="page">
+      <Suspense fallback={"Loading..."}>
+        <Content />
+      </Suspense>
+    </div>
+  )
+}
+
+export default InvalidateInfiniteQuery

test/integration/queries/test/index.test.ts

@@ -13,7 +13,7 @@ describe("Queries", () => {
       env: {__NEXT_TEST_WITH_DEVTOOL: 1},
     })
 
-    const prerender = ["/use-query", "/invalidate"]
+    const prerender = ["/use-query", "/invalidate-use-query", "/invalidate-use-infinite-query"]
     await Promise.all(prerender.map((route) => renderViaHTTP(context.appPort, route)))
   })
   afterAll(() => killApp(context.server))
@@ -32,7 +32,7 @@ describe("Queries", () => {
 
   describe("invalidateQuery", () => {
     it("should invalidate the query", async () => {
-      const browser = await webdriver(context.appPort, "/invalidate")
+      const browser = await webdriver(context.appPort, "/invalidate-use-query")
       await browser.waitForElementByCss("#content")
       let text = await browser.elementByCss("#content").text()
       expect(text).toMatch(/0/)
@@ -75,10 +75,66 @@ describe("Queries", () => {
 
   describe("DehydratedState", () => {
     it("should work", async () => {
-      const browser = await webdriver(context.appPort, "/dehydrated-state")
+      const browser = await webdriver(context.appPort, "/dehydrated-state-use-query")
       let text = await browser.elementByCss("#content").text()
       expect(text).toMatch(/map is Map: true/)
       if (browser) await browser.close()
     })
   })
+
+  describe("DehydratedState with usePaginatedQuery", () => {
+    it("should work", async () => {
+      const browser = await webdriver(context.appPort, "/dehydrated-state-use-paginated-query")
+      let text = await browser.elementByCss("#content").text()
+      expect(text).toMatch(/map is Map: true/)
+      if (browser) await browser.close()
+    })
+  })
+
+  describe("DehydratedState with useInfiniteQuery", () => {
+    it("should work", async () => {
+      const browser = await webdriver(context.appPort, "/dehydrated-state-use-infinite-query")
+      await browser.waitForElementByCss("#content")
+      let text = await browser.elementByCss("#content").text()
+      expect(JSON.parse(text)).toEqual([
+        {
+          items: [10, 11, 12, 13, 14],
+          hasMore: true,
+          nextPage: {take: 5, skip: 5},
+          count: 100,
+        },
+      ])
+    })
+  })
+
+  describe("invalidateQuery with useInfiniteQuery", () => {
+    it("should invalidate the query", async () => {
+      const browser = await webdriver(context.appPort, "/invalidate-use-infinite-query")
+      await browser.waitForElementByCss("#content")
+      let text = await browser.elementByCss("#content").text()
+      expect(JSON.parse(text)).toEqual([
+        {
+          counter: 1,
+          items: [10, 11, 12, 13, 14],
+          hasMore: true,
+          nextPage: {take: 5, skip: 5},
+          count: 100,
+        },
+      ])
+      await browser.elementByCss("button").click()
+      waitFor(500)
+      text = await browser.elementByCss("#content").text()
+      expect(JSON.parse(text)).toEqual([
+        {
+          counter: 2,
+          items: [10, 11, 12, 13, 14],
+          hasMore: true,
+          nextPage: {take: 5, skip: 5},
+          count: 100,
+        },
+      ])
+
+      if (browser) await browser.close()
+    })
+  })
 })