feat: pass a function to the revalidate option in mutate (#2862)

koba04 · web-flow · commit 3668c90af08d · 2024-02-16T00:07:27.000+01:00
feat: pass a function to revalidate option in mutate
diff --git a/src/_internal/types.ts b/src/_internal/types.ts
@@ -315,7 +315,7 @@ export type MutatorCallback<Data = any> = (
  * @typeParam MutationData - The type of the data returned by the mutator
  */
 export type MutatorOptions<Data = any, MutationData = Data> = {
-  revalidate?: boolean
+  revalidate?: boolean | ((data: Data, key: Arguments) => boolean)
   populateCache?:
     | boolean
     | ((result: MutationData, currentData: Data | undefined) => Data)
diff --git a/src/_internal/utils/mutate.ts b/src/_internal/utils/mutate.ts
@@ -60,7 +60,6 @@ export async function internalMutate<Data>(
   const rollbackOnErrorOption = options.rollbackOnError
   let optimisticData = options.optimisticData
 
-  const revalidate = options.revalidate !== false
   const rollbackOnError = (error: unknown): boolean => {
     return typeof rollbackOnErrorOption === 'function'
       ? rollbackOnErrorOption(error)
@@ -99,6 +98,9 @@ export async function internalMutate<Data>(
 
     const startRevalidate = () => {
       const revalidators = EVENT_REVALIDATORS[key]
+      const revalidate = isFunction(options.revalidate)
+        ? options.revalidate(get().data, _k)
+        : options.revalidate !== false
       if (revalidate) {
         // Invalidate the key by deleting the concurrent request markers so new
         // requests will not be deduped.
diff --git a/src/infinite/index.ts b/src/infinite/index.ts
@@ -21,7 +21,6 @@ import type {
   SWRHook,
   MutatorCallback,
   Middleware,
-  MutatorOptions,
   GlobalState
 } from '../_internal'
 import type {
@@ -31,7 +30,8 @@ import type {
   SWRInfiniteKeyLoader,
   SWRInfiniteFetcher,
   SWRInfiniteCacheValue,
-  SWRInfiniteCompareFn
+  SWRInfiniteCompareFn,
+  SWRInfiniteMutatorOptions
 } from './types'
 import { useSyncExternalStore } from 'use-sync-external-store/shim/index.js'
 import { getFirstPageKey } from './serialize'
@@ -55,7 +55,7 @@ export const infinite = (<Data, Error>(useSWRNext: SWRHook) =>
     fn: BareFetcher<Data> | null,
     config: Omit<typeof SWRConfig.defaultValue, 'fetcher'> &
       Omit<SWRInfiniteConfiguration<Data, Error>, 'fetcher'>
-  ): SWRInfiniteResponse<Data, Error> => {
+  ) => {
     const didMountRef = useRef<boolean>(false)
     const {
       cache,
@@ -140,6 +140,8 @@ export const infinite = (<Data, Error>(useSWRNext: SWRHook) =>
       async key => {
         // get the revalidate context
         const forceRevalidateAll = get()._i
+        const shouldRevalidatePage = get()._r
+        set({ _r: UNDEFINED })
 
         // return an array of page data
         const data: Data[] = []
@@ -187,7 +189,12 @@ export const infinite = (<Data, Error>(useSWRNext: SWRHook) =>
             (cacheData &&
               !isUndefined(cacheData[i]) &&
               !config.compare(cacheData[i], pageData))
-          if (fn && shouldFetchPage) {
+          if (
+            fn &&
+            (typeof shouldRevalidatePage === 'function'
+              ? shouldRevalidatePage(pageData, pageArg)
+              : shouldFetchPage)
+          ) {
             const revalidate = async () => {
               const hasPreloadedRequest = pageKey in PRELOAD
               if (!hasPreloadedRequest) {
@@ -238,7 +245,7 @@ export const infinite = (<Data, Error>(useSWRNext: SWRHook) =>
           | Data[]
           | Promise<Data[] | undefined>
           | MutatorCallback<Data[]>,
-        opts?: undefined | boolean | MutatorOptions<Data[], T>
+        opts?: undefined | boolean | SWRInfiniteMutatorOptions<Data[], T>
       ) {
         // When passing as a boolean, it's explicitly used to disable/enable
         // revalidation.
@@ -253,10 +260,10 @@ export const infinite = (<Data, Error>(useSWRNext: SWRHook) =>
         if (shouldRevalidate) {
           if (!isUndefined(data)) {
             // We only revalidate the pages that are changed
-            set({ _i: false })
+            set({ _i: false, _r: options.revalidate })
           } else {
             // Calling `mutate()`, we revalidate all pages
-            set({ _i: true })
+            set({ _i: true, _r: options.revalidate })
           }
         }
 
diff --git a/src/infinite/types.ts b/src/infinite/types.ts
@@ -4,7 +4,9 @@ import type {
   Arguments,
   BareFetcher,
   State,
-  StrictTupleKey
+  StrictTupleKey,
+  MutatorOptions,
+  MutatorCallback
 } from '../_internal'
 
 type FetcherResponse<Data = unknown> = Data | Promise<Data>
@@ -41,12 +43,29 @@ export interface SWRInfiniteConfiguration<
   compare?: SWRInfiniteCompareFn<Data>
 }
 
+interface SWRInfiniteRevalidateFn<Data = any> {
+  (data: Data, key: Arguments): boolean
+}
+
+type InfiniteKeyedMutator<Data> = <MutationData = Data>(
+  data?: Data | Promise<Data | undefined> | MutatorCallback<Data>,
+  opts?: boolean | SWRInfiniteMutatorOptions<Data, MutationData>
+) => Promise<Data | MutationData | undefined>
+
+export interface SWRInfiniteMutatorOptions<Data = any, MutationData = Data>
+  extends Omit<MutatorOptions<Data, MutationData>, 'revalidate'> {
+  revalidate?:
+    | boolean
+    | SWRInfiniteRevalidateFn<Data extends unknown[] ? Data[number] : never>
+}
+
 export interface SWRInfiniteResponse<Data = any, Error = any>
-  extends SWRResponse<Data[], Error> {
+  extends Omit<SWRResponse<Data[], Error>, 'mutate'> {
   size: number
   setSize: (
     size: number | ((_size: number) => number)
   ) => Promise<Data[] | undefined>
+  mutate: InfiniteKeyedMutator<Data[]>
 }
 
 export interface SWRInfiniteHook {
@@ -134,4 +153,5 @@ export interface SWRInfiniteCacheValue<Data = any, Error = any>
   // same key.
   _l?: number
   _k?: Arguments
+  _r?: boolean | SWRInfiniteRevalidateFn
 }
diff --git a/src/mutation/types.ts b/src/mutation/types.ts
@@ -1,4 +1,4 @@
-import type { SWRResponse, Key } from '../core'
+import type { SWRResponse, Key, Arguments } from '../core'
 
 type FetcherResponse<Data> = Data | Promise<Data>
 
@@ -25,7 +25,7 @@ export type SWRMutationConfiguration<
   ExtraArg = any,
   SWRData = any
 > = {
-  revalidate?: boolean
+  revalidate?: boolean | ((data: Data, key: Arguments) => boolean)
   populateCache?:
     | boolean
     | ((result: Data, currentData: SWRData | undefined) => SWRData)
diff --git a/test/use-swr-infinite.test.tsx b/test/use-swr-infinite.test.tsx
@@ -1843,4 +1843,48 @@ describe('useSWRInfinite', () => {
     screen.getByText('data:apple, banana, pineapple,')
     expect(previousPageDataLogs.every(d => d === null)).toBeTruthy()
   })
+
+  it('should support revalidate as a function', async () => {
+    // mock api
+    let pageData = ['apple', 'banana', 'pineapple']
+
+    const key = createKey()
+    function Page() {
+      const { data, mutate: boundMutate } = useSWRInfinite(
+        index => [key, index],
+        ([_, index]) => createResponse(pageData[index]),
+        {
+          initialSize: 3
+        }
+      )
+
+      return (
+        <div
+          onClick={() => {
+            boundMutate(undefined, {
+              // only revalidate 'apple' & 'pineapple' (page=2)
+              revalidate: (d, [_, i]: [string, number]) => {
+                return d === 'apple' || i === 2
+              }
+            })
+          }}
+        >
+          data:{Array.isArray(data) && data.join(',')}
+        </div>
+      )
+    }
+
+    renderWithConfig(<Page />)
+    screen.getByText('data:')
+
+    await screen.findByText('data:apple,banana,pineapple')
+
+    // update response data
+    pageData = pageData.map(data => `[${data}]`)
+
+    // revalidate
+    fireEvent.click(screen.getByText('data:apple,banana,pineapple'))
+
+    await screen.findByText('data:[apple],banana,[pineapple]')
+  })
 })
diff --git a/test/use-swr-local-mutation.test.tsx b/test/use-swr-local-mutation.test.tsx
@@ -1810,4 +1810,119 @@ describe('useSWR - local mutation', () => {
       [key, 'inf', 1]
     ])
   })
+  it('should support revalidate as a function', async () => {
+    let value = 0,
+      mutate
+    const key = createKey()
+    function Page() {
+      mutate = useSWRConfig().mutate
+      const { data } = useSWR(key, () => value++)
+      return <div>data: {data}</div>
+    }
+
+    renderWithConfig(<Page />)
+    screen.getByText('data:')
+
+    // mount
+    await screen.findByText('data: 0')
+
+    act(() => {
+      // value 0 -> 0
+      mutate(key, 100, { revalidate: () => false })
+    })
+    await screen.findByText('data: 100')
+
+    act(() => {
+      // value 0 -> 1
+      mutate(key, 200, { revalidate: () => true })
+    })
+    await screen.findByText('data: 200')
+    await screen.findByText('data: 1')
+  })
+
+  it('the function-style relivadate option receives the key and current data', async () => {
+    let value = 0,
+      mutate
+    const key = createKey()
+    function Page() {
+      mutate = useSWRConfig().mutate
+      const { data } = useSWR(key, () => value++)
+      return <div>data: {data}</div>
+    }
+
+    renderWithConfig(<Page />)
+    screen.getByText('data:')
+
+    // mount
+    await screen.findByText('data: 0')
+
+    act(() => {
+      // value 0 -> 0
+      mutate(key, 100, { revalidate: (d, k) => k === key && d === 200 }) // revalidate = false
+    })
+    await screen.findByText('data: 100')
+
+    act(() => {
+      // value 0 -> 1
+      mutate(key, 200, { revalidate: (d, k) => k === key && d === 200 }) // revalidate = true
+    })
+    await screen.findByText('data: 200')
+    await screen.findByText('data: 1')
+  })
+
+  it('the function-style relivadate option works with mutate filter', async () => {
+    const key1 = createKey()
+    const key2 = createKey()
+    const key3 = createKey()
+
+    let mockData = {
+      [key1]: 'page1',
+      [key2]: 'page2',
+      [key3]: 'page3'
+    }
+    function Page() {
+      const mutate = useSWRConfig().mutate
+      const { data: data1 } = useSWR(key1, () => mockData[key1])
+      const { data: data2 } = useSWR(key2, () => mockData[key2])
+      const { data: data3 } = useSWR(key3, () => mockData[key3])
+
+      return (
+        <>
+          <div>data1: {data1}</div>
+          <div>data2: {data2}</div>
+          <div>data3: {data3}</div>
+          <button
+            onClick={() => {
+              // key1 is filtered
+              mutate(k => k !== key1, 'updated', {
+                // only revalidate key3
+                revalidate: (d, k) => d === 'updated' && k === key3
+              })
+            }}
+          >
+            click
+          </button>
+        </>
+      )
+    }
+
+    renderWithConfig(<Page />)
+
+    // mount
+    await screen.findByText('data1: page1')
+    await screen.findByText('data2: page2')
+    await screen.findByText('data3: page3')
+
+    mockData = {
+      [key1]: '<page1>',
+      [key2]: '<page2>',
+      [key3]: '<page3>'
+    }
+
+    fireEvent.click(screen.getByText('click'))
+
+    await screen.findByText('data1: page1')
+    await screen.findByText('data2: updated')
+    await screen.findByText('data3: <page3>')
+  })
 })
diff --git a/test/use-swr-remote-mutation.test.tsx b/test/use-swr-remote-mutation.test.tsx
@@ -2,7 +2,7 @@ import { act, fireEvent, render, screen } from '@testing-library/react'
 import React, { useState } from 'react'
 import useSWR from 'swr'
 import useSWRMutation from 'swr/mutation'
-import { createKey, sleep, nextTick } from './utils'
+import { createKey, sleep, nextTick, createResponse } from './utils'
 
 const waitForNextTick = () => act(() => sleep(1))
 
@@ -1033,4 +1033,52 @@ describe('useSWR - remote mutation', () => {
     await screen.findByText('data:1,count:1')
     expect(logs).toEqual([0, 1])
   })
+
+  it('should support revalidate as a function', async () => {
+    const key = createKey()
+
+    let value = 0
+
+    function Page() {
+      const { data } = useSWR(key, () => createResponse(++value))
+      const { trigger } = useSWRMutation(key, () => {
+        value += 10
+        return createResponse(value)
+      })
+
+      return (
+        <div>
+          <button
+            onClick={() =>
+              trigger(undefined, {
+                revalidate: (d, k) => k === key && d < 30,
+                populateCache: true
+              })
+            }
+          >
+            trigger
+          </button>
+          <div>data:{data || 'none'}</div>
+        </div>
+      )
+    }
+
+    render(<Page />)
+
+    // mount
+    await screen.findByText('data:1')
+
+    fireEvent.click(screen.getByText('trigger'))
+    await screen.findByText('data:12')
+    fireEvent.click(screen.getByText('trigger'))
+    await screen.findByText('data:23')
+    fireEvent.click(screen.getByText('trigger'))
+    await screen.findByText('data:33')
+
+    // stop revalidation because value > 30
+    fireEvent.click(screen.getByText('trigger'))
+    await screen.findByText('data:43')
+    fireEvent.click(screen.getByText('trigger'))
+    await screen.findByText('data:53')
+  })
 })
