feat(middleware/persist): add skip hydration option #405 (#1647)

* add manual hydration option to persist middleware

* rename variable, update jsdoc, remove from oldImpl

* add tests for persist skipHydration

* add docs

* Update docs/integrations/persisting-store-data.md

Co-authored-by: Blazej Sewera <code@sewera.dev>

* Update docs/integrations/persisting-store-data.md

Co-authored-by: Blazej Sewera <code@sewera.dev>

* Update src/middleware/persist.ts

---------

Co-authored-by: Daishi Kato <dai-shi@users.noreply.github.com>
Co-authored-by: Blazej Sewera <code@sewera.dev>

docs/integrations/persisting-store-data.md

@@ -261,6 +261,50 @@ export const useBoundStore = create(
 )
 ```
 
+### `skipHydration`
+
+> Type: `boolean | undefined`
+
+> Default: `undefined`
+
+By default the store with be hydrated on initialization.
+
+In some applications you may need to control when the first hydration occurs.
+For example, in server-rendered apps.
+
+If you set `skipHydration`, the initial call for hydration isn't called,
+and it is left up to you to manually call `reHydrate()`.
+
+```ts
+export const useBoundStore = create(
+  persist(
+    () => ({
+      count: 0,
+      // ...
+    }),
+    {
+      // ...
+      skipHydration: true,
+    }
+  )
+)
+```
+
+```tsx
+export function StoreConsumer(){
+  const store = useBoundStore();
+
+  // hydrate persisted store after on mount
+  useEffect(() => {
+    store.persist.reHydrate();
+  }, [])
+
+  return (
+    //...
+  )
+}
+```
+
 ## API
 
 > Version: >=3.6.3

src/middleware/persist.ts

@@ -119,6 +119,16 @@ export interface PersistOptions<S, PersistedState = S> {
    * By default, this function does a shallow merge.
    */
   merge?: (persistedState: unknown, currentState: S) => S
+
+  /**
+   * An optional boolean that will prevent the persist middleware from triggering hydration on initialization,
+   * This allows you to call `rehydrate()` at a specific point in your apps rendering life-cycle.
+   *
+   * This is useful in SSR application.
+   *
+   * @default false
+   */
+  skipHydration?: boolean
 }
 
 type PersistListener<S> = (state: S) => void
@@ -491,7 +501,9 @@ const newImpl: PersistImpl = (config, baseOptions) => (set, get, api) => {
     },
   }
 
-  hydrate()
+  if (!options.skipHydration) {
+    hydrate()
+  }
 
   return stateFromStorage || configResult
 }

tests/persistAsync.test.tsx

@@ -561,4 +561,55 @@ describe('persist middleware with async configuration', () => {
     await useBoundStore.persist.rehydrate()
     expect(useBoundStore.persist.hasHydrated()).toBe(true)
   })
+
+  it('can skip initial hydration', async () => {
+    const storage = {
+      getItem: async (name: string) => ({
+        state: { count: 42, name },
+        version: 0,
+      }),
+      setItem: () => {},
+      removeItem: () => {},
+    }
+
+    const onRehydrateStorageSpy = jest.fn()
+    const useBoundStore = create(
+      persist(
+        () => ({
+          count: 0,
+          name: 'empty',
+        }),
+        {
+          name: 'test-storage',
+          storage: storage,
+          onRehydrateStorage: () => onRehydrateStorageSpy,
+          skipHydration: true,
+        }
+      )
+    )
+
+    expect(useBoundStore.getState()).toEqual({
+      count: 0,
+      name: 'empty',
+    })
+
+    // Because `skipHydration` is only in newImpl and the hydration function for newImpl is now a promise
+    // In the default case we would need to await `onFinishHydration` to assert the auto hydration has completed
+    // As we are testing the skip hydration case we await nextTick, to make sure the store is initialised
+    await new Promise((resolve) => process.nextTick(resolve))
+
+    // Asserting store hasn't hydrated from nextTick
+    expect(useBoundStore.persist.hasHydrated()).toBe(false)
+
+    await useBoundStore.persist.rehydrate()
+
+    expect(useBoundStore.getState()).toEqual({
+      count: 42,
+      name: 'test-storage',
+    })
+    expect(onRehydrateStorageSpy).toBeCalledWith(
+      { count: 42, name: 'test-storage' },
+      undefined
+    )
+  })
 })

tests/persistSync.test.tsx

@@ -548,4 +548,55 @@ describe('persist middleware with sync configuration', () => {
     expect(onFinishHydrationSpy1).not.toBeCalledTimes(2)
     expect(onFinishHydrationSpy2).toBeCalledWith({ count: 2 })
   })
+
+  it('can skip initial hydration', async () => {
+    const storage = {
+      getItem: (name: string) => ({
+        state: { count: 42, name },
+        version: 0,
+      }),
+      setItem: () => {},
+      removeItem: () => {},
+    }
+
+    const onRehydrateStorageSpy = jest.fn()
+    const useBoundStore = create(
+      persist(
+        () => ({
+          count: 0,
+          name: 'empty',
+        }),
+        {
+          name: 'test-storage',
+          storage: storage,
+          onRehydrateStorage: () => onRehydrateStorageSpy,
+          skipHydration: true,
+        }
+      )
+    )
+
+    expect(useBoundStore.getState()).toEqual({
+      count: 0,
+      name: 'empty',
+    })
+
+    // Because `skipHydration` is only in newImpl and the hydration function for newImpl is now a promise
+    // In the default case we would need to await `onFinishHydration` to assert the auto hydration has completed
+    // As we are testing the skip hydration case we await nextTick, to make sure the store is initialised
+    await new Promise((resolve) => process.nextTick(resolve))
+
+    // Asserting store hasn't hydrated from nextTick
+    expect(useBoundStore.persist.hasHydrated()).toBe(false)
+
+    await useBoundStore.persist.rehydrate()
+
+    expect(useBoundStore.getState()).toEqual({
+      count: 42,
+      name: 'test-storage',
+    })
+    expect(onRehydrateStorageSpy).toBeCalledWith(
+      { count: 42, name: 'test-storage' },
+      undefined
+    )
+  })
 })